#include "stm32c0xx_hal.h"
#include <stdbool.h>

// === Pin Definitions ===
#define CARGREEN_PIN    GPIO_PIN_5   // PA5
#define CARYELLOW_PIN   GPIO_PIN_6   // PA6
#define CARRED_PIN      GPIO_PIN_7   // PA7
#define PEDGREEN_PIN    GPIO_PIN_2   // PA2
#define PEDRED_PIN      GPIO_PIN_3   // PA3
#define PEDBUTTON_PIN   GPIO_PIN_8   // PA8
#define BLUELED_PIN     GPIO_PIN_15  // PA15
#define LED_PORT        GPIOA
#define BUTTON_PORT     GPIOA

//7-Segment pins (common anode logic)
#define SEG_A_PIN  GPIO_PIN_8
#define SEG_B_PIN  GPIO_PIN_9
#define SEG_C_PIN  GPIO_PIN_5
#define SEG_D_PIN  GPIO_PIN_4
#define SEG_E_PIN  GPIO_PIN_6
#define SEG_F_PIN  GPIO_PIN_3
#define SEG_G_PIN  GPIO_PIN_10
#define SEG_PORT   GPIOB

#define BLINKS 8
#define BLINKSPEED 250

enum State { ALERT, STOP, DRIVE, BRAKE, WALK, BLINK };
State state = STOP;

uint32_t start = 0;
uint8_t blinkCounter = 0;
bool pedestrianRequest = false;

// === 7-Segment Display ===
void displayDigit(uint8_t digit) {
  const bool digits[10][7] = {
    {1,1,1,1,1,1,0}, {0,1,1,0,0,0,0}, {1,1,0,1,1,0,1}, {1,1,1,1,0,0,1},
    {0,1,1,0,0,1,1}, {1,0,1,1,0,1,1}, {1,0,1,1,1,1,1}, {1,1,1,0,0,0,0},
    {1,1,1,1,1,1,1}, {1,1,1,1,0,1,1}
  };

  HAL_GPIO_WritePin(SEG_PORT, SEG_A_PIN, digits[digit][0] ? GPIO_PIN_RESET : GPIO_PIN_SET);
  HAL_GPIO_WritePin(SEG_PORT, SEG_B_PIN, digits[digit][1] ? GPIO_PIN_RESET : GPIO_PIN_SET);
  HAL_GPIO_WritePin(SEG_PORT, SEG_C_PIN, digits[digit][2] ? GPIO_PIN_RESET : GPIO_PIN_SET);
  HAL_GPIO_WritePin(SEG_PORT, SEG_D_PIN, digits[digit][3] ? GPIO_PIN_RESET : GPIO_PIN_SET);
  HAL_GPIO_WritePin(SEG_PORT, SEG_E_PIN, digits[digit][4] ? GPIO_PIN_RESET : GPIO_PIN_SET);
  HAL_GPIO_WritePin(SEG_PORT, SEG_F_PIN, digits[digit][5] ? GPIO_PIN_RESET : GPIO_PIN_SET);
  HAL_GPIO_WritePin(SEG_PORT, SEG_G_PIN, digits[digit][6] ? GPIO_PIN_RESET : GPIO_PIN_SET);
}

void clearDisplay() {
  HAL_GPIO_WritePin(SEG_PORT, SEG_A_PIN | SEG_B_PIN | SEG_C_PIN |
                    SEG_D_PIN | SEG_E_PIN | SEG_F_PIN | SEG_G_PIN, GPIO_PIN_SET);
}

// === Light Control ===
void setLights(bool carGreen, bool carYellow, bool carRed, bool pedGreen, bool pedRed) {
  if (carGreen && pedGreen) {
    state = ALERT;
    start = HAL_GetTick();
  } else {
    HAL_GPIO_WritePin(LED_PORT, CARGREEN_PIN, carGreen ? GPIO_PIN_SET : GPIO_PIN_RESET);
    HAL_GPIO_WritePin(LED_PORT, CARYELLOW_PIN, carYellow ? GPIO_PIN_SET : GPIO_PIN_RESET);
    HAL_GPIO_WritePin(LED_PORT, CARRED_PIN, carRed ? GPIO_PIN_SET : GPIO_PIN_RESET);
    HAL_GPIO_WritePin(LED_PORT, PEDGREEN_PIN, pedGreen ? GPIO_PIN_SET : GPIO_PIN_RESET);
    HAL_GPIO_WritePin(LED_PORT, PEDRED_PIN, pedRed ? GPIO_PIN_SET : GPIO_PIN_RESET);
  }
}

void stopState() {
  // Check and remember button press
  if (HAL_GPIO_ReadPin(BUTTON_PORT, PEDBUTTON_PIN) == GPIO_PIN_SET) {
    pedestrianRequest = true;
    HAL_GPIO_WritePin(LED_PORT, BLUELED_PIN, GPIO_PIN_SET); // Turn on blue LED
  }

  // If request remembered, initiate walk
  if (pedestrianRequest) {
    state = WALK;
    start = HAL_GetTick();
    setLights(false, false, true, true, false);
    pedestrianRequest = false; // Clear request
    return;
  }

  if (HAL_GetTick() - start > 2000) {
    state = DRIVE;
    start = HAL_GetTick();
    setLights(true, false, false, false, true);
  }
}

void drive() {
  if (HAL_GPIO_ReadPin(BUTTON_PORT, PEDBUTTON_PIN) == GPIO_PIN_SET) {
    pedestrianRequest = true;
    HAL_GPIO_WritePin(LED_PORT, BLUELED_PIN, GPIO_PIN_SET); // Turn on blue LED
  }

  if (HAL_GetTick() - start > 5000) {
    state = BRAKE;
    start = HAL_GetTick();
    setLights(false, true, false, false, true);
  }
}

void brake() {
  if (HAL_GPIO_ReadPin(BUTTON_PORT, PEDBUTTON_PIN) == GPIO_PIN_SET) {
    pedestrianRequest = true;
    HAL_GPIO_WritePin(LED_PORT, BLUELED_PIN, GPIO_PIN_SET); // Turn on blue LED
  }

  if (HAL_GetTick() - start > 2000) {
    state = STOP;
    start = HAL_GetTick();
    setLights(false, false, true, false, true);
  }
}

void walk() {
  uint32_t elapsed = HAL_GetTick() - start;
  if (elapsed <= 4000) {
    uint8_t secondsLeft = 4 - (elapsed / 1000);
    displayDigit(secondsLeft);
  }

  if (elapsed > 4000) {
    state = BLINK;
    start = HAL_GetTick();
    setLights(false, false, true, false, false);
    blinkCounter = 2 * BLINKS;
  }
}

void blink() {
  if (HAL_GetTick() - start > BLINKSPEED) {
    GPIO_PinState current = HAL_GPIO_ReadPin(LED_PORT, PEDGREEN_PIN);
    HAL_GPIO_WritePin(LED_PORT, PEDGREEN_PIN, current == GPIO_PIN_SET ? GPIO_PIN_RESET : GPIO_PIN_SET);

    blinkCounter--;
    start = HAL_GetTick();

    if (blinkCounter % 2 == 0 && blinkCounter > 0) {
      displayDigit(blinkCounter / 2);
    }

    if (blinkCounter == 0) {
      state = STOP;
      start = HAL_GetTick();
      setLights(false, false, true, false, true);
      clearDisplay();
      HAL_GPIO_WritePin(LED_PORT, BLUELED_PIN, GPIO_PIN_RESET); // Turn off blue LED after blinking ends
    }
  }
}

void alert() {
  if (HAL_GetTick() - start > BLINKSPEED) {
    GPIO_PinState val = HAL_GPIO_ReadPin(LED_PORT, CARYELLOW_PIN);
    HAL_GPIO_WritePin(LED_PORT, CARYELLOW_PIN, val == GPIO_PIN_SET ? GPIO_PIN_RESET : GPIO_PIN_SET);
    start = HAL_GetTick();
  }
}

// === GPIO Init Helper ===
void initPin(GPIO_TypeDef* port, uint16_t pin, uint32_t mode, uint32_t pull) {
  GPIO_InitTypeDef GPIO_InitStruct = {0};
  GPIO_InitStruct.Pin = pin;
  GPIO_InitStruct.Mode = mode;
  GPIO_InitStruct.Pull = pull;
  GPIO_InitStruct.Speed = GPIO_SPEED_FREQ_LOW;
  HAL_GPIO_Init(port, &GPIO_InitStruct);
}

// === Main ===
int main(void) {
  HAL_Init();
  __HAL_RCC_GPIOA_CLK_ENABLE();
  __HAL_RCC_GPIOB_CLK_ENABLE();

  // Output pins (LEDs and 7-segment)
  initPin(LED_PORT, CARGREEN_PIN | CARYELLOW_PIN | CARRED_PIN |
          PEDGREEN_PIN | PEDRED_PIN | BLUELED_PIN, GPIO_MODE_OUTPUT_PP, GPIO_NOPULL);

  initPin(BUTTON_PORT, PEDBUTTON_PIN, GPIO_MODE_INPUT, GPIO_PULLDOWN);

  initPin(SEG_PORT, SEG_A_PIN | SEG_B_PIN | SEG_C_PIN | SEG_D_PIN |
          SEG_E_PIN | SEG_F_PIN | SEG_G_PIN, GPIO_MODE_OUTPUT_PP, GPIO_NOPULL);

  // Initial state
  start = HAL_GetTick();
  setLights(false, false, true, false, true); // Initial STOP
  clearDisplay();

  while (1) {
    switch (state) {
      case STOP:  stopState(); break;
      case DRIVE: drive(); break;
      case BRAKE: brake(); break;
      case WALK:  walk(); break;
      case BLINK: blink(); break;
      default:    alert(); break;
    }
    HAL_Delay(10);
  }
}