#include <WiFi.h>
#include <Firebase_ESP_Client.h>
#include <ESP32Servo.h>
#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <addons/TokenHelper.h>
#include <addons/RTDBHelper.h>

#define WIFI_SSID "Wokwi-GUEST"
#define WIFI_PASSWORD ""

#define API_KEY "AIzaSyB654P2Pdrx7EUkD1RmLFIZq5jFo2RAki4"
#define DATABASE_URL "https://smart-traffic-system-4ac4b-default-rtdb.firebaseio.com/"
#define USER_EMAIL "[email protected]"
#define USER_PASSWORD "1234567890"

FirebaseData fbdo;
FirebaseAuth auth;
FirebaseConfig config;

#define RED_PIN 12
#define YELLOW_PIN 18
#define GREEN_PIN 5
#define PEDESTRIAN_BUTTON 4

#define SERVO_PEDESTRIAN 14
#define SERVO_EMERGENCY 27

LiquidCrystal_I2C lcd(0x27, 16, 2);

Servo pedestrianServo;
Servo emergencyServo;

#define NORMAL_GREEN 20000UL
#define NORMAL_RED   20000UL
#define LONG_GREEN   30000UL
#define SHORT_RED    10000UL
#define YELLOW_TIME  3000UL

#define PED_WAIT_TIME 10000UL
#define PED_CROSS_TIME 20000UL
#define BUTTON_DEBOUNCE 200UL

enum TrafficState { RED, YELLOW, GREEN };
TrafficState trafficState = GREEN;
TrafficState nextAfterYellow = RED;

unsigned long stateStart = 0;
unsigned long currentGreen = NORMAL_GREEN;
unsigned long currentRed   = NORMAL_RED;

bool pedRequest = false;
bool pedActive = false;
unsigned long pedRequestTime = 0;
unsigned long pedStart = 0;
unsigned long lastButtonTime = 0;

bool emergencyFB = false;
int carCount = 0;

unsigned long fbPrevMillis = 0;

void connectWiFi() {
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
  while (WiFi.status() != WL_CONNECTED) delay(200);
}

void connectFirebase() {
  config.api_key = API_KEY;
  auth.user.email = USER_EMAIL;
  auth.user.password = USER_PASSWORD;
  config.database_url = DATABASE_URL;
  Firebase.begin(&config, &auth);
  Firebase.reconnectWiFi(true);
}

void setTraffic(TrafficState s, TrafficState next) {
  trafficState = s;
  nextAfterYellow = next;
  stateStart = millis();
  digitalWrite(RED_PIN, s == RED);
  digitalWrite(YELLOW_PIN, s == YELLOW);
  digitalWrite(GREEN_PIN, s == GREEN);
}

void trafficAI() {
  if (carCount > 2 && carCount < 7) {
    currentGreen = NORMAL_GREEN;
    currentRed   = NORMAL_RED;
  } else {
    currentGreen = LONG_GREEN;
    currentRed   = SHORT_RED;
  }
}

void trafficFSM(unsigned long now) {
  unsigned long elapsed = now - stateStart;

  if (trafficState == GREEN) {
    if (elapsed >= currentGreen)
      setTraffic(YELLOW, RED);
  }
  else if (trafficState == YELLOW) {
    if (elapsed >= YELLOW_TIME)
      setTraffic(nextAfterYellow, nextAfterYellow == RED ? GREEN : RED);
  }
  else {
    if (!pedActive && elapsed >= currentRed)
      setTraffic(YELLOW, GREEN);
  }
}

void pedestrianLogic(unsigned long now) {
  if (digitalRead(PEDESTRIAN_BUTTON) == LOW &&
      now - lastButtonTime > BUTTON_DEBOUNCE) {
    pedRequest = true;
    pedRequestTime = now;
    lastButtonTime = now;
  }

  if (pedRequest && !pedActive) {
    if (now - pedRequestTime < PED_WAIT_TIME) return;

    setTraffic(RED, GREEN);
    pedActive = true;
    pedRequest = false;
    pedStart = now;
    pedestrianServo.write(0);
  }

  if (pedActive && now - pedStart >= PED_CROSS_TIME) {
    pedActive = false;
    pedestrianServo.write(90);
    stateStart = now;
  }
}

void emergencyLogic() {
  if (emergencyFB) emergencyServo.write(0);
  else emergencyServo.write(90);
}

void firebaseTask(unsigned long now) {
  if (!Firebase.ready()) return;
  if (now - fbPrevMillis < 1000) return;
  fbPrevMillis = now;

  if (Firebase.RTDB.getBool(&fbdo, "/ai/detections/latest/is_emergency"))
    emergencyFB = fbdo.boolData();

  if (Firebase.RTDB.getInt(&fbdo, "/ai/detections/latest/car_count"))
    carCount = fbdo.intData();
}

void sendStatusToFirebase(unsigned long now) {
  static unsigned long lastSend = 0;
  if (!Firebase.ready()) return;
  if (now - lastSend < 1000) return;
  lastSend = now;

  Firebase.RTDB.setBool(&fbdo,
    "/esp32/status/pedestrian_active",
    pedActive
  );

  int pedRemain = 0;
  if (pedActive)
    pedRemain = (PED_CROSS_TIME - (now - pedStart)) / 1000;

  Firebase.RTDB.setInt(&fbdo,
    "/esp32/status/pedestrian_remaining",
    pedRemain
  );

  Firebase.RTDB.setInt(&fbdo,
    "/esp32/status/pollution",
    50
  );

  String color =
    trafficState == GREEN ? "GREEN" :
    trafficState == YELLOW ? "YELLOW" : "RED";

  Firebase.RTDB.setString(&fbdo,
    "/esp32/status/traffic_color",
    color
  );

  unsigned long remain =
    trafficState == GREEN ? (currentGreen - (now - stateStart)) :
    trafficState == YELLOW ? (YELLOW_TIME - (now - stateStart)) :
    (currentRed - (now - stateStart));

  Firebase.RTDB.setInt(&fbdo,
    "/esp32/status/traffic_remaining",
    remain / 1000
  );
}

void updateLCD(unsigned long now) {
  static unsigned long last = 0;
  if (now - last < 300) return;
  last = now;

  lcd.clear();

  if (pedActive) {
    lcd.setCursor(0,0);
    lcd.print("PEDESTRIAN");
    lcd.setCursor(0,1);
    lcd.print("CROSS ");
    lcd.print((PED_CROSS_TIME - (now - pedStart)) / 1000);
    lcd.print(" SEC");
    return;
  }

  if (pedRequest) {
    lcd.setCursor(0,0);
    lcd.print("RED IN");
    lcd.setCursor(0,1);
    lcd.print((PED_WAIT_TIME - (now - pedRequestTime)) / 1000);
    lcd.print(" SEC");
    return;
  }

  lcd.setCursor(0,0);
  lcd.print("TRAFFIC ");
  lcd.print(trafficState == GREEN ? "GREEN" :
            trafficState == YELLOW ? "YELLOW" : "RED");

  unsigned long remain =
    trafficState == GREEN ? (currentGreen - (now - stateStart)) :
    trafficState == YELLOW ? (YELLOW_TIME - (now - stateStart)) :
    (currentRed - (now - stateStart));

  lcd.setCursor(0,1);
  lcd.print("REMAIN ");
  lcd.print(remain / 1000);
  lcd.print(" SEC");
}

void setup() {
  pinMode(RED_PIN, OUTPUT);
  pinMode(YELLOW_PIN, OUTPUT);
  pinMode(GREEN_PIN, OUTPUT);
  pinMode(PEDESTRIAN_BUTTON, INPUT_PULLUP);

  pedestrianServo.attach(SERVO_PEDESTRIAN);
  emergencyServo.attach(SERVO_EMERGENCY);
  pedestrianServo.write(90);
  emergencyServo.write(90);

  Wire.begin();
  lcd.init();
  lcd.backlight();

  connectWiFi();
  connectFirebase();

  setTraffic(GREEN, RED);
}

void loop() {
  unsigned long now = millis();

  firebaseTask(now);
  trafficAI();
  emergencyLogic();
  pedestrianLogic(now);
  trafficFSM(now);
  updateLCD(now);
  sendStatusToFirebase(now);
}