/*
   ===================================================================
   IOT PROJECT: ACCESSIBLE SMART CITY STATION (UPB)
   ===================================================================
 * * This project simulates a multi-function, accessible smart city station
   using an ESP32 in the Wokwi environment.
 * * It monitors real-world data (environment, waste levels, traffic)
   and calculates real-time bus ETAs using an RTC, publishing everything
   via MQTT.
 * * It is fully bidirectional, allowing a user (via a Flutter app) to:
   1. Request a pedestrian crossing (triggering LEDs and a buzzer).
   2. Send a "locate" ping to the station's buzzer for accessibility.
   3. Push emergency alerts to the station's OLED display.
   4. Poll for sensor data on demand.
 * * Hardware: ESP32, DHT22, MQ135, LDR, HC-SR04, DS1307, SSD1306,
   LEDs, Buzzer, Pushbutton.
   Protocol: MQTT (broker.hivemq.com)
*/

// --- PROJECT-WIDE LIBRARIES ---
#include <WiFi.h>
#include <PubSubClient.h>   // For MQTT
#include <ArduinoJson.h>    // For sending and receiving JSON
#include <Wire.h>           // For I2C (OLED and RTC)
#include <DHT.h>              // For the DHT22 sensor
#include <RTClib.h>           // For the DS1307 clock
#include <Adafruit_SSD1306.h> // For the OLED display
#include <Adafruit_GFX.h>   // Graphics for OLED
#include <NewPing.h>        // For the HCSR04.h

// --- 1. NETWORK CONFIGURATION ---
const char* WIFI_SSID = "Wokwi-GUEST";
const char* WIFI_PASS = "";
const char* MQTT_BROKER = "broker.hivemq.com";
const int   MQTT_PORT = 1883;
const char* MQTT_CLIENT_ID = "upb_station_01"; // Must be unique

// --- 2. PINS  ---
#define DHT_PIN 15
#define MQ135_PIN 1        
#define LDR_PIN 2             
#define HC_TRIG_PIN 10
#define HC_ECHO_PIN 11
#define LED_RED_PIN 12
#define LED_AMBER_PIN 13
#define LED_GREEN_PIN 14
#define BUTTON_PIN 4          
#define BUZZER_PIN 16
#define I2C_SDA_PIN 8         
#define I2C_SCL_PIN 9         

// --- 3. MQTT TOPICS ---
const char* TOPIC_ENV_OUT = "upb/smart-city/environment";
const char* TOPIC_TRANSPORT_OUT = "upb/smart-city/transport";
const char* TOPIC_TRAFFIC_OUT = "upb/smart-city/traffic/state";
const char* TOPIC_REQUEST_IN = "upb/smart-city/traffic/request";
const char* TOPIC_LOCATE_IN = "upb/smart-city/traffic/locate";
const char* TOPIC_DISPLAY_IN = "upb/smart-city/display/set_message";
const char* TOPIC_POLL_IN = "upb/smart-city/environment/poll";

// --- 4. GLOBAL OBJECTS ---
WiFiClient   wifiClient;
PubSubClient mqttClient(wifiClient);

// Sensor Objects (Part 2)
DHT dht(DHT_PIN, DHT22);
const float BIN_DEPTH_CM = 400.0;

// Max distance 400cm
NewPing sonar(HC_TRIG_PIN, HC_ECHO_PIN, 400); // 400cm for a bin is quite high, just simulation

// I2C, RTC and OLED Objects (Part 3)
RTC_DS1307 rtc;
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// --- 5. TIMERS & STATE VARIABLES ---
unsigned long lastSensorPoll = 0;
const long sensorInterval = 15000;
unsigned long lastBusUpdate = 0;
const long busInterval = 30000;
unsigned long lastTrafficCheck = 0;
const long trafficInterval = 500;
bool oledAlertMode = false;
unsigned long lastAlertTime = 0;
const long alertDuration = 20000;
int lastDisplayedEta = -1;
enum TrafficState { CARS_GREEN, CARS_AMBER, CARS_RED, PED_WALK };
TrafficState currentTrafficState;
unsigned long lastStateChange = 0;
unsigned long lastBeep = 0;
String lastPublishedState = "";
volatile bool pedestrianRequest = false;
const long MANDATORY_PED_TIME = 3000; // ms of guaranteed pedestrian time inside CARS_RED
bool autoPedestrianGiven = false;     // tracks if that auto-window already used for this red cycle
const long CARS_GREEN_TIME = 10000;
const long CARS_AMBER_TIME = 3000;
const long CARS_RED_TIME = 8000;
const long PED_WALK_TIME = 8000;
const long BEEP_INTERVAL = 500;

// --- Function declarations ---
void publishEnvironmentData();
void calculateAndPublishBusData();
void updateOledEta(int eta);
void showOledAlert(String message);
void manageTrafficLight();
void setState(TrafficState newState);
void publishTrafficState(String stateName);


// ===================================================================
// MQTT CALLBACK FUNCTION (The Command Brain!)
// ===================================================================
void mqttCallback(char* topic, byte* payload, unsigned int length) {
  Serial.print("Message received on [");
  Serial.print(topic);
  Serial.print("]: ");

  String message;
  for (int i = 0; i < length; i++) {
    message += (char)payload[i];
  }
  Serial.println(message);

  if (strcmp(topic, TOPIC_POLL_IN) == 0) {
    Serial.println(">> COMMAND: Sensor poll request received.");
    publishEnvironmentData();
    lastSensorPoll = millis();
  }
  else if (strcmp(topic, TOPIC_REQUEST_IN) == 0) {
    Serial.println(">> COMMAND: Pedestrian cross request received.");
    pedestrianRequest = true;
  }
  else if (strcmp(topic, TOPIC_LOCATE_IN) == 0) {
    Serial.println(">> COMMAND: Locate (Ping) request received.");
    tone(BUZZER_PIN, 1200, 150);
  }
  else if (strcmp(topic, TOPIC_DISPLAY_IN) == 0) {
    Serial.println(">> COMMAND: Message for OLED display received.");
    StaticJsonDocument<128> doc;
    deserializeJson(doc, payload, length);
    const char* alertText = doc["text"];

    if (alertText) {
      showOledAlert(alertText);
    }
  }
}


// ===================================================================
// SETUP FUNCTIONS (WiFi & MQTT)
// ===================================================================
void setup() {
  Serial.begin(115200);
  Serial.println("Starting Smart City Station...");

  // --- 1. Configure Pins ---
  pinMode(LED_RED_PIN, OUTPUT);
  pinMode(LED_AMBER_PIN, OUTPUT);
  pinMode(LED_GREEN_PIN, OUTPUT);
  pinMode(BUZZER_PIN, OUTPUT);
  pinMode(BUTTON_PIN, INPUT_PULLUP);

  // --- 2. Start I2C ---
  Wire.begin(I2C_SDA_PIN, I2C_SCL_PIN);

  // --- 3. Initialize RTC (Part 3) ---
  if (!rtc.begin()) {
    Serial.println("Couldn't find RTC!");
    while (1);
  }

  if (!rtc.isrunning()) {
    Serial.println("RTC is NOT running, setting time to compile time");
    rtc.adjust(DateTime(F(__DATE__), F(__TIME__)));
  }
  Serial.println("RTC initialized.");

  // --- 4. Initialize OLED (Part 3) ---
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(10, 20);
  display.println("Smart City Station");
  display.println("     Booting...");
  display.display();
  delay(1000);

  // --- 5. Connect WiFi ---
  initWiFi();

  // --- 6. Configure MQTT ---
  initMQTT();

  // --- 7. Initialize Sensors (Part 2) ---
  dht.begin();
  Serial.println("Sensors initialized.");

  // --- 8. Initialize Traffic State (Part 4) ---
  Serial.println("Initializing traffic system...");
  setState(CARS_RED);

  Serial.println("Setup complete. Entering main loop.");
}

void initWiFi() {
  Serial.print("Connecting to WiFi ");
  Serial.print(WIFI_SSID);
  WiFi.begin(WIFI_SSID, WIFI_PASS);
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
  }
  Serial.println("\nWiFi connected!");
  Serial.print("IP address: ");
  Serial.println(WiFi.localIP());
}

void initMQTT() {
  mqttClient.setServer(MQTT_BROKER, MQTT_PORT);
  mqttClient.setCallback(mqttCallback);
}

void reconnectMQTT() {
  while (!mqttClient.connected()) {
    Serial.print("Attempting MQTT connection...");
    if (mqttClient.connect(MQTT_CLIENT_ID)) {
      Serial.println("Connected!");
      mqttClient.subscribe(TOPIC_REQUEST_IN);
      mqttClient.subscribe(TOPIC_LOCATE_IN);
      mqttClient.subscribe(TOPIC_DISPLAY_IN);
      mqttClient.subscribe(TOPIC_POLL_IN);
      Serial.println("Subscribed to all command topics.");
    } else {
      Serial.print("failed, rc=");
      Serial.print(mqttClient.state());
      Serial.println(" try again in 5 seconds");
      delay(5000);
    }
  }
}

// ===================================================================
// SENSOR PUBLISHING FUNCTION (Part 2)
// ===================================================================
void publishEnvironmentData() {
  Serial.println("Reading all sensors...");

  float temp = dht.readTemperature();
  float hum = dht.readHumidity();
  int airQuality = analogRead(MQ135_PIN);
  int lightLevel = analogRead(LDR_PIN);

  float distance = sonar.ping_cm();

  // NewPing returns 0 for no ping
  if (distance == 0) {
    Serial.println("ERROR: Sensor out of distance!");
    // Set to a "full" value so we don't get 100%
    distance = BIN_DEPTH_CM;
  }

  if (isnan(temp) || isnan(hum)) {
    Serial.println("Failed to read from DHT sensor!");
    return;
  }

  float binFullPct = 100.0 - (distance / BIN_DEPTH_CM) * 100.0;
  if (binFullPct < 0) binFullPct = 0;
  if (binFullPct > 100) binFullPct = 100;

  StaticJsonDocument<300> doc;
  doc["temp_c"] = temp;
  doc["humidity"] = hum;
  doc["air_raw"] = airQuality;
  doc["light_raw"] = lightLevel;
  doc["bin_pct"] = (int)binFullPct;

  char jsonBuffer[300];
  serializeJson(doc, jsonBuffer);

  mqttClient.publish(TOPIC_ENV_OUT, jsonBuffer);

  Serial.println("Environmental data published.");
}


// ===================================================================
// BUS & OLED FUNCTIONS (Part 3)
// ===================================================================
void calculateAndPublishBusData() {
  DateTime now = rtc.now();
  int currentMinute = now.minute();

  int eta_min = 0;
  if (currentMinute < 15)       eta_min = 15 - currentMinute;
  else if (currentMinute < 30)  eta_min = 30 - currentMinute;
  else if (currentMinute < 45)  eta_min = 45 - currentMinute;
  else                          eta_min = (60 - currentMinute) + 15;

  Serial.printf("RTC Time: %02d:%02d. Next bus (381) in %d min.\n",
                now.hour(), now.minute(), eta_min);

  StaticJsonDocument<100> doc;
  doc["line"] = "381";
  doc["eta_min"] = eta_min;

  char jsonBuffer[100];
  serializeJson(doc, jsonBuffer);
  mqttClient.publish(TOPIC_TRANSPORT_OUT, jsonBuffer);

  if (!oledAlertMode) {
    updateOledEta(eta_min);
  }
}

void updateOledEta(int eta) {
  if (eta == lastDisplayedEta) {
    return;
  }
  lastDisplayedEta = eta;

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("BUS 381 (Piata Romana)");
  display.drawFastHLine(0, 10, display.width(), SSD1306_WHITE);
  display.setTextSize(3);
  display.setCursor(20, 25);
  display.print(eta);
  display.setTextSize(2);
  display.setCursor(70, 32);
  display.print("min");
  display.display();
}

void showOledAlert(String message) {
  oledAlertMode = true;
  lastAlertTime = millis();
  lastDisplayedEta = -1;

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(15, 0);
  display.println("-- STATION ALERT --");
  display.drawFastHLine(0, 10, display.width(), SSD1306_WHITE);
  display.setTextSize(2);
  display.setCursor(0, 25);
  display.println(message);
  display.display();
}

// ===================================================================
// TRAFFIC LIGHT FUNCTIONS (Part 4)
// ===================================================================
void manageTrafficLight() {
  unsigned long now = millis();

  switch (currentTrafficState) {
    case CARS_GREEN:
      if (now - lastStateChange >= CARS_GREEN_TIME) {
        setState(CARS_AMBER);
      }
      break;

    case CARS_AMBER:
      if (now - lastStateChange >= CARS_AMBER_TIME) {
        setState(CARS_RED);
      }
      break;

    case CARS_RED:
      if (digitalRead(BUTTON_PIN) == LOW) {
        pedestrianRequest = true;
      }

      // Decide transition to PED_WALK:
      // - if a pedestrian pressed the button, give them the walk immediately
      // - otherwise, when the red has progressed enough to reach the mandatory window,
      //   grant the automatic pedestrian walk (once per red cycle)
      if (!autoPedestrianGiven) {
        unsigned long elapsed = now - lastStateChange;
        // start auto-window at the end of red: when elapsed >= (CARS_RED_TIME - MANDATORY_PED_TIME)
        if (pedestrianRequest || (elapsed >= (unsigned long)(CARS_RED_TIME - MANDATORY_PED_TIME))) {
          autoPedestrianGiven = true;
          setState(PED_WALK);
          break;
        }
      }

      // If no pedestrian and mandatory window already passed+used, allow usual cycle to green
      if (now - lastStateChange >= CARS_RED_TIME) {
        setState(CARS_GREEN);
      }
      break;

    case PED_WALK:
      if (now - lastBeep >= BEEP_INTERVAL) {
        lastBeep = now;
        tone(BUZZER_PIN, 880, 200);
      }
      if (now - lastStateChange >= PED_WALK_TIME) {
        setState(CARS_RED);
      }
      break;
  }
}


void setState(TrafficState newState) {
  if (newState == currentTrafficState) return;

  currentTrafficState = newState;
  lastStateChange = millis();
  String stateName = "";

  digitalWrite(LED_RED_PIN, LOW);
  digitalWrite(LED_AMBER_PIN, LOW);
  digitalWrite(LED_GREEN_PIN, LOW);
  noTone(BUZZER_PIN);

  switch (newState) {
    case CARS_GREEN:
      digitalWrite(LED_GREEN_PIN, HIGH);
      stateName = "CARS_GREEN";
      break;
    case CARS_AMBER:
      digitalWrite(LED_AMBER_PIN, HIGH);
      stateName = "CARS_AMBER";
      break;
    case CARS_RED:
      digitalWrite(LED_RED_PIN, HIGH);
      stateName = "CARS_RED";
      autoPedestrianGiven = false; // <-- reset on every entry to CARS_RED
      break;
    case PED_WALK:
      digitalWrite(LED_RED_PIN, HIGH);
      stateName = "PED_WALK";
      pedestrianRequest = false;
      lastBeep = millis();
      break;
  }

  Serial.printf("TRAFFIC: State changed to %s\n", stateName.c_str());
  publishTrafficState(stateName);
}

void publishTrafficState(String stateName) {
  if (stateName == "" || stateName == lastPublishedState) {
    return;
  }
  lastPublishedState = stateName;

  StaticJsonDocument<50> doc;
  doc["state"] = stateName;

  char jsonBuffer[50];
  serializeJson(doc, jsonBuffer);
  mqttClient.publish(TOPIC_TRAFFIC_OUT, jsonBuffer);
}


// ===================================================================
// MAIN LOOP (The "Task Scheduler")
// ===================================================================
void loop() {

  if (!mqttClient.connected()) {
    reconnectMQTT();
  }
  mqttClient.loop();

  unsigned long now = millis();

  if (now - lastSensorPoll >= sensorInterval) {
    lastSensorPoll = now;
    publishEnvironmentData();
  }

  if (now - lastBusUpdate >= busInterval) {
    lastBusUpdate = now;
    calculateAndPublishBusData();
  }

  if (now - lastTrafficCheck >= trafficInterval) {
    lastTrafficCheck = now;
    manageTrafficLight();
  }

  if (oledAlertMode && (now - lastAlertTime > alertDuration)) {
    Serial.println("OLED Alert timed out. Reverting to ETA display.");
    oledAlertMode = false;
    lastDisplayedEta = -1;
    calculateAndPublishBusData();
  }
}