#include <RTClib.h>      // Define RTC_DS1307 and DateTime
#include "my_lcd.h"
#include "my_wifi.h"
#include "PsychicMqttClient.h"
#include <DHT.h>         // DHT sensor library

// --- GPIO Pin ---
#define DHTPIN 32      
#define DHTTYPE DHT22  
#define RED_LED_PIN 26
#define ORANGE_LED_PIN 19
#define PUSH_BUTTON_PIN 34
#define BLUE_LED_PIN 27 // Corrected pin name for the warning LED

// --- MQTT Configuration ---
const char MQTTserver[] = "mqtts://34373d9086c442d3bc260aca9b1928bb.s1.eu.hivemq.cloud:8883";
const char MQTT_CLIENT_ID[] = "MengMilestone2"; 
const char MQTT_USERNAME[] = "IOT_Milestone";
const char MQTT_PASSWORD[] = "Admin@12345"; 
const char* LWT_TOPIC = "Icecreampolur5/Status/Device";
const char* LWT_MESSAGE_OFFLINE = "Offline"; // Last Will Testment
const char* LWT_MESSAGE_ONLINE = "Online"; // Last Will Testment
const int LWT_QOS = 1;             // Quality of Service for the LWT message
const bool LWT_RETAIN = true;        // Retain flag for the LWT message (important for status updates)

// --- MQTT Status Topic ---
const char* MQTT_TEMP_TOPIC = "Icecreampolur5/Status/Temp";
const char* MQTT_RH_TOPIC = "Icecreampolur5/Status/RH";
const char* MQTT_LIGHT_TOPIC = "Icecreampolur5/Status/Light";
const char* MQTT_COMMAND_TOPIC = "Icecreampolur5/Command/Light";
const char* MQTT_WARNING_TOPIC = "Icecreampolur5/Warning/Env";

// --- MQTT Alarm Topic ---
const char* MQTT_ALARM_FIRE_TOPIC = "Icecreampolur5/Alarm/Fire"; // Topic for Fire Alarm
const char* MQTT_STATE_ON = "Alarm ON";
const char* MQTT_STATE_OFF = "Alarm OFF";
const char* MQTT_STATELIGHT_ON = "Light ON";
const char* MQTT_STATELIGHT_OFF = "Light OFF";
const char* MQTT_WARNING_HIGH_TEMP = "High Temperature Warning";
const char* MQTT_WARNING_HIGH_RH = "High Humidity Warning";
const char* MQTT_WARNING_OK = "Environmental Conditions OK";

// --- MQTT Command Topic for Alarm Control ---
const char* MQTT_COMMAND_ON = "ON";  // Payload to turn Air-Cond ON
const char* MQTT_COMMAND_OFF = "OFF"; // Payload to turn Air-Cond OFF
const char* PAYLOAD_ON = "ON";
const char* PAYLOAD_OFF = "OFF";

// --- Global Object ---
MyLCD       lcd;
RTC_DS1307 rtc;
MyWiFi      wifi;
PsychicMqttClient mqttClient;
DHT dht(DHTPIN, DHTTYPE); // Initialize DHT object with pin and type

// --- Global Alarm State Variables ---
bool fireAlarmActive = false; // Initialize to false so the LED starts off
bool LightActive = false;  
bool previousLightState = false; 
bool topicsSubscribed = false; 
bool highTempWarning = false;
bool highRhWarning = false;

// --- Millis() variables for non-blocking operations ---
unsigned long previousLcdDisplayMillis = 0;
const long lcdDisplayInterval = 2000; // LCD display updates every 2 seconds

unsigned long previousMqttPublishMillis = 0;
const long mqttPublishInterval = 5000; // MQTT data publishes every 5 seconds

// --- Button Debounce Variables ---
unsigned long lastDebounceTime = 0; // The last time the input pin was toggled
const long debounceDelay = 100;      // The debounce time; increase if not stable

// --- Variables to hold the current and previous reading from the button pin ---
int buttonReading;      // The current raw reading from the input pin
int buttonState = LOW;      // The debounced current state of the button
int lastButtonState = LOW;  // The debounced previous state of the button

// --- Keep track of LCD state for cycling through different displays ---
enum LcdState {
  STATE_WELCOME_MESSAGE,
  STATE_DATETIME_ALARM,
  STATE_SENSOR_WARNING
};
LcdState currentLcdState = STATE_WELCOME_MESSAGE; // Start with the welcome message
// --------------------------

// --- Variables to store sensor readings ---
float h = 0.0; // Humidity value
float t = 0.0; // Temperature value

// --- Callback function for MQTT Alarm Command topic ---
void mqttLightCommand(
  const char *topic,
  const char *payload,
  int retain,
  int qos,
  bool dup
) {
  Serial.printf("Received Light Command on topic: %s, payload: %s\r\n", topic, payload);
  if (strcmp(payload, MQTT_COMMAND_ON) == 0) {
    LightActive = true;
    Serial.println("Light commanded ON via MQTT.");
  } else if (strcmp(payload, MQTT_COMMAND_OFF) == 0) {
    LightActive = false;
    Serial.println("Light commanded OFF via MQTT.");
  } else {
    Serial.println("Unknown Light Command payload. Use 'ON' or 'OFF'.");
  }
}


void subscribe_mqtt_topics() { // Renamed for clarity
  Serial.println("Attempting to subscribe to MQTT topics...");
  mqttClient.onTopic(MQTT_COMMAND_TOPIC, 1, mqttLightCommand);
  Serial.println("Subscription calls made for topics.");
}


void setup() {
  Serial.begin(115200); // Initialize serial communication
  lcd.begin();          // Initialize the LCD

  // --- Set up GPIO pins ---
  pinMode(RED_LED_PIN, OUTPUT);      // Red LED pin
  digitalWrite(RED_LED_PIN, LOW);    // Ensure LED is off initially
  pinMode(PUSH_BUTTON_PIN, INPUT_PULLDOWN); // Push Button pin with internal pulldown
  pinMode(ORANGE_LED_PIN, OUTPUT);      // Light LED pin
  digitalWrite(ORANGE_LED_PIN, LOW);    // Ensure LED is off initially
  pinMode(BLUE_LED_PIN, OUTPUT);        // Blue LED pin
  digitalWrite(BLUE_LED_PIN, LOW);      // Ensure Blue LED is off initially
  
  // --- Initialize RTC (Real-Time Clock) ---
  if (!rtc.begin()) {  // Check if RTC module is found
    Serial.println(F("Error: RTC not found!"));
    lcd.print("RTC not found!");
    while (1);    // Halt execution if RTC fails, as time is critical
  }
  wifi.begin(&rtc);     // Initialize WiFi, passing RTC for time sync

  // --- Initialize DHT sensor ---
  dht.begin();
  Serial.println("DHT sensor initialized.");

 // --- MQTT Client Configuration ---
  mqttClient.setServer(MQTTserver); 
  mqttClient.setClientId(MQTT_CLIENT_ID);
  mqttClient.setCredentials(MQTT_USERNAME, MQTT_PASSWORD);
  Serial.printf("Set MQTT Client ID: %s, Username: %s\r\n", MQTT_CLIENT_ID, MQTT_USERNAME);

  mqttClient.setKeepAlive(30);
  Serial.println("Set MQTT Keep Alive to 30 seconds.");

  mqttClient.attachArduinoCACertBundle();
  Serial.println("Attached Arduino CA Certificate Bundle for TLS/SSL.");

  // --- Set Last Will and Testament BEFORE connecting ---
  mqttClient.setWill(LWT_TOPIC, 1, true, LWT_MESSAGE_OFFLINE);
  Serial.printf("Set Last Will (Topic: %s, Message: %s)\r\n", LWT_TOPIC, LWT_MESSAGE_OFFLINE);

  // --- Attach MQTT Event Handlers ---
  mqttClient.onConnect([&](bool sessionPresent) { // Added capture list [&]
    Serial.printf("MQTT Client Connected! Session Present: %d\r\n", sessionPresent);
    if (!sessionPresent || !topicsSubscribed) { // Re-subscribe if new session or flag is false
    subscribe_mqtt_topics();
    topicsSubscribed = true;
    }
    // --- Publish "Online" status after successful connection ---
    mqttClient.publish(LWT_TOPIC, 1, true, LWT_MESSAGE_ONLINE); // Publish online message
    Serial.printf("Published Online status: %s\r\n", LWT_MESSAGE_ONLINE);
  });

  mqttClient.onDisconnect([&](bool sessionPresent) { // Added capture list [&]
    Serial.printf("MQTT Client Disconnected! Session Present: %d\r\n", sessionPresent);
    topicsSubscribed = false; // Reset flag so we resubscribe on reconnect
  });
  mqttClient.onError([](esp_mqtt_error_codes_t error) {
    Serial.printf("MQTT Error occurred: %d\r\n", error);
  });
  mqttClient.onSubscribe([](int msgId) {
    Serial.printf("MQTT Subscribed (Msg ID: %d)\r\n", msgId);
  });
  mqttClient.onPublish([](int msgId) {
    Serial.printf("MQTT Published (Msg ID: %d)\r\n", msgId);
  });

  // --- Publish initial Light state ---
  mqttClient.publish(MQTT_LIGHT_TOPIC, 1, true, LightActive ? MQTT_STATELIGHT_ON : MQTT_STATELIGHT_OFF);
  Serial.printf("Published Initial Light State: %s\r\n", LightActive ? MQTT_STATELIGHT_ON : MQTT_STATELIGHT_OFF);
  previousLightState = LightActive; // Initialize previous state
}

void loop() {
  // --- WiFi & MQTT Connection Management ---
  if (wifi.isTimeSynched()) {
    if (!mqttClient.connected()) {
      Serial.println("MQTT Client not connected. Attempting to connect...");
      mqttClient.connect();
    }
  } else {
    Serial.println("WiFi not connected or time not synchronized.");
    topicsSubscribed = false; // Ensure resubscription if WiFi drops and reconnects
    previousMqttPublishMillis = millis() - mqttPublishInterval;
  }

  // --- Non-blocking Button Read and Debounce Logic ---
  unsigned long currentMillis = millis(); // Get current time for timing checks

  buttonReading = digitalRead(PUSH_BUTTON_PIN); // Read the raw button state

  // --- If the switch changed, due to pressing: ---
  if (buttonReading != lastButtonState) {
    // --- reset the debouncing timer ---
    lastDebounceTime = currentMillis;
  }

  if ((currentMillis - lastDebounceTime) > debounceDelay) {
    // Whatever the reading is at, it's been there for longer than the debounce delay,
    // so take it as the actual current debounced state:
    if (buttonReading != buttonState) {
      buttonState = buttonReading;

      // Only act if the button is pressed (HIGH, due to INPUT_PULLDOWN to VCC)
      if (buttonState == HIGH) {
        Serial.println("Button Pressed!");
        // Toggle the alarm state when button is pressed
        fireAlarmActive = !fireAlarmActive; // Inverts the current alarm state
        Serial.printf("Alarm state toggled to: %s\r\n", fireAlarmActive ? "ON" : "OFF");

        // Publish MQTT message about button press/release
        if (wifi.isTimeSynched() && mqttClient.connected()) {
          if (fireAlarmActive) { // If alarm just turned ON by button
            mqttClient.publish(MQTT_ALARM_FIRE_TOPIC, 1, false, MQTT_STATE_ON);
            Serial.println("MQTT Alarm Sent: Button Pressed (Alarm ON)");
          } else { // If alarm just turned OFF by button
            mqttClient.publish(MQTT_ALARM_FIRE_TOPIC, 1, false, MQTT_STATE_OFF);
            Serial.println("MQTT Alarm Sent: Button Released (Alarm OFF)");
          }
        }
      }
    }
  }
  // Save the reading. Next time through the loop, it'll be the lastButtonState:
  lastButtonState = buttonReading; // This should always be the raw reading for debounce


  // --- Centralized LED Control ---
  digitalWrite(ORANGE_LED_PIN, LightActive ? HIGH : LOW);
  digitalWrite(RED_LED_PIN, fireAlarmActive ? HIGH : LOW);
  
  // Turn on the Purple LED if a warning condition is met
  bool warningActive = highTempWarning || highRhWarning;
  digitalWrite(BLUE_LED_PIN, warningActive ? HIGH : LOW);

  // --- Check and publish Light state only if it changes ---
  if (LightActive != previousLightState) {
    if (wifi.isTimeSynched() && mqttClient.connected()) {
      mqttClient.publish(MQTT_LIGHT_TOPIC, 1, true, LightActive ? MQTT_STATELIGHT_ON : MQTT_STATELIGHT_OFF);
      Serial.printf("Light State Changed and Published: %s\r\n", LightActive ? MQTT_STATELIGHT_ON : MQTT_STATELIGHT_OFF);
    }
    previousLightState = LightActive; // Update previous state after publishing
  }

  // --- Non-blocking LCD Display Logic ---
  if (currentMillis - previousLcdDisplayMillis >= lcdDisplayInterval) {
    previousLcdDisplayMillis = currentMillis; // Reset the timer for the next LCD update

    lcd.clear(); // Clear the LCD before displaying new content

    // Cycle through LCD display states
    switch (currentLcdState) {
      case STATE_WELCOME_MESSAGE:
        lcd.display("ESP32-DevKitC-v4", "Class DEE1", "Hui Meng");
        currentLcdState = STATE_DATETIME_ALARM; // Transition to next state
        break;

      case STATE_DATETIME_ALARM: {
        DateTime now = rtc.now(); // Get current date and time from RTC
        char dtbuf[21] = "YYYY-MM-DD hh:mm:ss"; // Buffer for formatted datetime string
        now.toString(dtbuf); // Convert DateTime object to string
        lcd.datetime(dtbuf); // Display datetime on LCD

        if (wifi.isTimeSynched()) { // Only display sensor data if time is synched (implies WiFi up)
          char temp_str[10]; // Buffer for temperature string
          char rh_str[10];   // Buffer for humidity string
          dtostrf(t, 4, 1, temp_str); // Format temperature float to string (e.g., "30.5")
          dtostrf(h, 4, 1, rh_str);   // Format humidity float to string (e.g., "65.2")

          lcd.alarm(1, "Temp xxxx:", temp_str, 0); // Temp Room1, with a '0' as placeholder for 4th arg
          lcd.alarm(2, "RH   xxxx:", rh_str, "%"); // RH Room1, with "%" as unit

          // --- ORANGE LED STATUS: Ensure 4 arguments are provided
          const char* lightStatus = LightActive ? "ON" : "OFF";
          lcd.alarm(3, "Light    :", lightStatus, "");
          
          // --- ADD THIS LINE TO DISPLAY FIRE ALARM STATUS ---
          const char* fireAlarmStatus = fireAlarmActive ? "ON" : "OFF";
          lcd.alarm(4, "FireAlarm:", fireAlarmStatus, "");
        }
        currentLcdState = highTempWarning || highRhWarning ? STATE_SENSOR_WARNING : STATE_WELCOME_MESSAGE; // Transition based on warning
        break;
      }

      case STATE_SENSOR_WARNING: {
        lcd.clear();
        lcd.print("WARNING!");
        if (highTempWarning) {
          lcd.setCursor(0, 1);
          lcd.print("HIGH TEMP > 40C");
        }
        if (highRhWarning) {
          lcd.setCursor(0, 2);
          lcd.print("HIGH RH > 50%");
        }
        currentLcdState = STATE_DATETIME_ALARM; // Transition back to normal display
        break;
      }
    }
  }

  // --- Non-blocking MQTT Publishing & DHT Reading Logic ---
  // Only attempt to read DHT and publish if WiFi is connected and MQTT is connected
  if (wifi.isTimeSynched() && mqttClient.connected()) {
    if (currentMillis - previousMqttPublishMillis >= mqttPublishInterval) {
      previousMqttPublishMillis = currentMillis; // Reset the timer for the next publish

      // Read temperature and humidity from DHT sensor
      h = dht.readHumidity();
      t = dht.readTemperature(); // Read temperature in Celsius

      // Check if DHT reads failed (returns NaN)
      if (isnan(h) || isnan(t)) {
        Serial.println(F("Failed to read from DHT sensor! Retrying next interval."));
        return; // Exit this block early, don't publish bad data or trigger alarms with bad data
      }

      Serial.printf("Publishing MQTT data: Temp=%.1f°C, Humidity=%.1f%%\r\n", t, h);

      // --- Environmental Warning Logic and MQTT Publishing ---
      bool currentHighTemp = t > 40.0;
      bool currentHighRh = h > 50.0;

      if (currentHighTemp != highTempWarning || currentHighRh != highRhWarning) {
        if (currentHighTemp && !currentHighRh) {
          mqttClient.publish(MQTT_WARNING_TOPIC, 1, false, MQTT_WARNING_HIGH_TEMP);
          Serial.println("Published High Temperature Warning.");
        } else if (!currentHighTemp && currentHighRh) {
          mqttClient.publish(MQTT_WARNING_TOPIC, 1, false, MQTT_WARNING_HIGH_RH);
          Serial.println("Published High Humidity Warning.");
        } else if (currentHighTemp && currentHighRh) {
          // You could combine this into a single message
          mqttClient.publish(MQTT_WARNING_TOPIC, 1, false, "High Temp & RH Warning");
          Serial.println("Published High Temp & RH Warning.");
        } else {
          mqttClient.publish(MQTT_WARNING_TOPIC, 1, false, MQTT_WARNING_OK);
          Serial.println("Published Environmental Conditions OK.");
        }
      }

      // Update the global state variables
      highTempWarning = currentHighTemp;
      highRhWarning = currentHighRh;

      // Convert float readings to string for MQTT payload
      char temp_payload[10];
      char rh_payload[10];
      dtostrf(t, 4, 1, temp_payload); // Format temperature for MQTT
      dtostrf(h, 4, 1, rh_payload);   // Format humidity for MQTT

      mqttClient.publish(MQTT_TEMP_TOPIC, 2, 1, temp_payload); // Publish actual temperature with QoS 2, Retain 1
      mqttClient.publish(MQTT_RH_TOPIC, 1, 1, rh_payload);     // Publish actual humidity with QoS 1, Retain 1
    }
  }
}