/*
 * IoT Edge Device with MQTT Integration
 * ESP32-based environmental monitoring system with alarm capabilities
 * Features: DHT22 sensor, I2C LCD, MQTT communication, NTP sync, low power mode
 */

#include <WiFi.h>
#include <PubSubClient.h>
#include <DHT.h>
#include <LiquidCrystal_I2C.h>
#include <ArduinoJson.h>
#include <time.h>
#include <esp_sleep.h>

// Hardware pin definitions
#define DHT_PIN 4
#define BUTTON_PIN 2
#define LED_PIN 5
#define DHT_TYPE DHT22  // Changed from DHT11 to DHT22

// I2C LCD configuration (16x2 display)
#define LCD_ADDR 0x27
#define LCD_COLS 16
#define LCD_ROWS 2

// Network configuration
const char* ssid = "Wokwi-GUEST";
const char* password = "";

// MQTT configuration
const char* mqtt_server = "broker.hivemq.com";
const int mqtt_port = 1883;
const char* device_id = "esp32_iot_001";
const char* mqtt_user = "";
const char* mqtt_password = "";

// NTP configuration
const char* ntp_server = "pool.ntp.org";
const long gmt_offset_sec = 28800; // GMT+8 for Singapore
const int daylight_offset_sec = 0;

// MQTT Topics
String topic_status = String(device_id) + "/device/status";
String topic_lwt = String(device_id) + "/device/lwt";
String topic_sensor = String(device_id) + "/sensor/environment";
String topic_system = String(device_id) + "/sensor/system";
String topic_alarm = String(device_id) + "/alarm/trigger";
String topic_alarm_status = String(device_id) + "/alarm/status";
String topic_power_mode = String(device_id) + "/control/power_mode";
String topic_config = String(device_id) + "/control/config";

// Timing variables
unsigned long last_sensor_read = 0;
unsigned long last_system_report = 0;
unsigned long alarm_start_time = 0;
unsigned long button_last_press = 0;
unsigned long last_dht_read = 0;  // Added for DHT22 timing

// Configuration variables
int sensor_interval = 5000; // Changed to 5 seconds for better responsiveness during testing
int system_report_interval = 300000; // 5 minutes
int alarm_duration = 30000; // 30 seconds
bool low_power_mode = false;
int sleep_duration = 300; // 5 minutes in low power mode
const unsigned long DHT_READ_INTERVAL = 2000; // DHT22 minimum read interval

// State variables
bool alarm_active = false;
bool button_pressed = false;
bool wifi_connected = false;
bool mqtt_connected = false;
float last_temperature = 0.0;
float last_humidity = 0.0;
bool sensor_initialized = false;  // Added to track sensor initialization

// Hardware objects
DHT dht(DHT_PIN, DHT_TYPE);
LiquidCrystal_I2C lcd(LCD_ADDR, LCD_COLS, LCD_ROWS);
WiFiClient espClient;
PubSubClient mqtt_client(espClient);

void setup() {
  Serial.begin(115200);
  Serial.println("IoT Edge Device Starting...");
  
  // Initialize hardware
  initializeHardware();
  
  // Connect to WiFi
  connectWiFi();
  
  // Initialize NTP
  initializeNTP();
  
  // Connect to MQTT
  connectMQTT();
  
  // Display startup message
  displayStartupMessage();
  
  Serial.println("System initialized successfully");
}

void loop() {
  // Handle WiFi connection
  if (WiFi.status() != WL_CONNECTED) {
    wifi_connected = false;
    reconnectWiFi();
  } else {
    wifi_connected = true;
  }
  
  // Handle MQTT connection
  if (!mqtt_client.connected()) {
    mqtt_connected = false;
    reconnectMQTT();
  } else {
    mqtt_connected = true;
    mqtt_client.loop();
  }
  
  // Read sensors periodically
  if (millis() - last_sensor_read >= sensor_interval) {
    readAndPublishSensors();
    last_sensor_read = millis();
  }
  
  // System health report
  if (millis() - last_system_report >= system_report_interval) {
    publishSystemHealth();
    last_system_report = millis();
  }
  
  // Handle button press for alarm
  handleButtonPress();
  
  // Handle alarm state
  handleAlarmState();
  
  // Update LCD display
  updateDisplay();
  
  // Handle low power mode
  if (low_power_mode && !alarm_active) {
    enterLowPowerMode();
  }
  
  delay(100); // Small delay to prevent watchdog issues
}

void initializeHardware() {
  // Initialize pins
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  pinMode(LED_PIN, OUTPUT);
  digitalWrite(LED_PIN, LOW);
  
  // Initialize DHT sensor with longer delay for DHT22
  Serial.println("Initializing DHT22 sensor...");
  dht.begin();
  delay(2000); // Give DHT22 time to stabilize
  
  // Test sensor reading
  float test_temp = dht.readTemperature();
  float test_hum = dht.readHumidity();
  
  if (!isnan(test_temp) && !isnan(test_hum)) {
    sensor_initialized = true;
    Serial.println("DHT22 sensor initialized successfully");
    Serial.print("Initial readings - Temp: ");
    Serial.print(test_temp);
    Serial.print("°C, Humidity: ");
    Serial.print(test_hum);
    Serial.println("%");
  } else {
    Serial.println("Warning: DHT22 sensor not responding properly");
  }
  
  // Initialize LCD
  lcd.init();
  lcd.backlight();
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("IoT Device");
  lcd.setCursor(0, 1);
  lcd.print("Initializing...");
  
  Serial.println("Hardware initialized");
}

void connectWiFi() {
  WiFi.begin(ssid, password);
  Serial.print("Connecting to WiFi");
  
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 20) {
    delay(500);
    Serial.print(".");
    digitalWrite(LED_PIN, !digitalRead(LED_PIN)); // Blink LED during connection
    attempts++;
  }
  
  if (WiFi.status() == WL_CONNECTED) {
    wifi_connected = true;
    digitalWrite(LED_PIN, HIGH); // Solid LED when connected
    Serial.println();
    Serial.print("WiFi connected! IP: ");
    Serial.println(WiFi.localIP());
  } else {
    Serial.println();
    Serial.println("WiFi connection failed!");
    digitalWrite(LED_PIN, LOW);
  }
}

void reconnectWiFi() {
  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("Reconnecting to WiFi...");
    WiFi.disconnect();
    delay(1000);
    WiFi.begin(ssid, password);
  }
}

void initializeNTP() {
  if (wifi_connected) {
    configTime(gmt_offset_sec, daylight_offset_sec, ntp_server);
    Serial.println("NTP initialized");
  }
}

void connectMQTT() {
  mqtt_client.setServer(mqtt_server, mqtt_port);
  mqtt_client.setCallback(mqttCallback);
  reconnectMQTT();
}

void reconnectMQTT() {
  while (!mqtt_client.connected() && wifi_connected) {
    Serial.print("Attempting MQTT connection...");

    String client_id = String(device_id) + "_client";
    String lwt_payload = createLWTMessage();

    // Connect with Last Will and Testament
    if (mqtt_client.connect(
          client_id.c_str(),
          mqtt_user,
          mqtt_password,
          topic_lwt.c_str(),   // LWT topic
          1,                   // QoS
          true,                // Retain
          lwt_payload.c_str()  // LWT message
        )) {
      Serial.println("connected");
      mqtt_connected = true;

      // Subscribe to control topics
      mqtt_client.subscribe(topic_power_mode.c_str());
      mqtt_client.subscribe(topic_config.c_str());

      // Publish device online status
      publishDeviceStatus("online");

    } else {
      Serial.print("failed, rc=");
      Serial.print(mqtt_client.state());
      Serial.println(" try again in 5 seconds");
      delay(5000);
    }
  }
}

void mqttCallback(char* topic, byte* payload, unsigned int length) {
  String message = "";
  for (int i = 0; i < length; i++) {
    message += (char)payload[i];
  }
  
  Serial.print("Message arrived [");
  Serial.print(topic);
  Serial.print("] ");
  Serial.println(message);
  
  // Parse JSON message
  DynamicJsonDocument doc(1024);
  deserializeJson(doc, message);
  
  // Handle power mode control
  if (String(topic) == topic_power_mode) {
    String mode = doc["mode"];
    if (mode == "low_power") {
      low_power_mode = true;
      sleep_duration = doc["sleep_duration"] | 300;
    } else if (mode == "normal") {
      low_power_mode = false;
    }
  }
  
  // Handle configuration updates
  if (String(topic) == topic_config) {
    sensor_interval = (doc["sensor_interval"] | 60) * 1000;
    alarm_duration = (doc["alarm_duration"] | 30) * 1000;
  }
}

void readAndPublishSensors() {
  // Ensure minimum time between DHT22 readings
  if (millis() - last_dht_read < DHT_READ_INTERVAL) {
    return;
  }
  
  Serial.println("Reading DHT22 sensor...");
  
  // Read DHT22 sensor with retries
  float humidity = NAN;
  float temperature = NAN;
  int retries = 3;
  
  for (int i = 0; i < retries; i++) {
    humidity = dht.readHumidity();
    temperature = dht.readTemperature();
    
    // Check if readings are valid
    if (!isnan(humidity) && !isnan(temperature)) {
      break; // Successful reading
    }
    
    Serial.print("DHT22 read attempt ");
    Serial.print(i + 1);
    Serial.println(" failed, retrying...");
    delay(500); // Wait before retry
  }
  
  last_dht_read = millis();
  
  if (isnan(humidity) || isnan(temperature)) {
    Serial.println("Failed to read from DHT22 sensor after retries!");
    
    // Check sensor connection
    Serial.println("Sensor connection issue detected. Check wiring:");
    Serial.println("- VCC to 3.3V or 5V");
    Serial.println("- GND to Ground");
    Serial.println("- DATA to pin 4 with 10kΩ pull-up resistor");
    
    return;
  }
  
  // Update global variables with valid readings
  last_temperature = temperature;
  last_humidity = humidity;
  
  Serial.print("Temperature: ");
  Serial.print(temperature);
  Serial.print("°C, Humidity: ");
  Serial.print(humidity);
  Serial.println("%");
  
  // Create sensor data message
  DynamicJsonDocument doc(512);
  doc["temperature"] = temperature;
  doc["humidity"] = humidity;
  doc["timestamp"] = getCurrentTimestamp();
  doc["unit_temp"] = "°C";
  doc["unit_hum"] = "%";
  doc["sensor_type"] = "DHT22";
  
  String payload;
  serializeJson(doc, payload);
  
  // Publish sensor data
  if (mqtt_connected) {
    mqtt_client.publish(topic_sensor.c_str(), payload.c_str());
    Serial.println("Sensor data published: " + payload);
  }
}

void publishSystemHealth() {
  DynamicJsonDocument doc(512);
  doc["uptime"] = millis() / 1000;
  doc["free_heap"] = ESP.getFreeHeap();
  doc["wifi_rssi"] = WiFi.RSSI();
  doc["timestamp"] = getCurrentTimestamp();
  doc["sensor_initialized"] = sensor_initialized;
  
  String payload;
  serializeJson(doc, payload);
  
  if (mqtt_connected) {
    mqtt_client.publish(topic_system.c_str(), payload.c_str());
    Serial.println("System health published: " + payload);
  }
}

void handleButtonPress() {
  bool current_button = !digitalRead(BUTTON_PIN); // Inverted due to pull-up
  
  // Debounce button
  if (current_button && !button_pressed && (millis() - button_last_press > 200)) {
    button_pressed = true;
    button_last_press = millis();
    triggerAlarm();
  } else if (!current_button) {
    button_pressed = false;
  }
}

void triggerAlarm() {
  alarm_active = true;
  alarm_start_time = millis();
  
  // Create alarm message
  DynamicJsonDocument doc(512);
  doc["alarm_type"] = "manual";
  doc["severity"] = "high";
  doc["message"] = "Manual alarm triggered";
  doc["timestamp"] = getCurrentTimestamp();
  
  String payload;
  serializeJson(doc, payload);
  
  // Publish alarm
  if (mqtt_connected) {
    mqtt_client.publish(topic_alarm.c_str(), payload.c_str(), true);
    Serial.println("Alarm triggered: " + payload);
  }
  
  // Update alarm status
  publishAlarmStatus(true);
}

void handleAlarmState() {
  if (alarm_active) {
    // Blink LED during alarm
    digitalWrite(LED_PIN, (millis() / 250) % 2);
    
    // Check if alarm duration has elapsed
    if (millis() - alarm_start_time >= alarm_duration) {
      alarm_active = false;
      digitalWrite(LED_PIN, wifi_connected ? HIGH : LOW);
      publishAlarmStatus(false);
    }
  }
}

void publishAlarmStatus(bool active) {
  DynamicJsonDocument doc(512);
  doc["active"] = active;
  doc["type"] = "manual";
  doc["duration"] = alarm_duration / 1000;
  doc["timestamp"] = getCurrentTimestamp();
  
  String payload;
  serializeJson(doc, payload);
  
  if (mqtt_connected) {
    mqtt_client.publish(topic_alarm_status.c_str(), payload.c_str(), true);
  }
}

void publishDeviceStatus(String status) {
  DynamicJsonDocument doc(512);
  doc["status"] = status;
  doc["timestamp"] = getCurrentTimestamp();
  doc["ip"] = WiFi.localIP().toString();
  
  String payload;
  serializeJson(doc, payload);
  
  mqtt_client.publish(topic_status.c_str(), payload.c_str(), true);
  Serial.println("Device status published: " + payload);
}

String createLWTMessage() {
  DynamicJsonDocument doc(512);
  doc["status"] = "offline";
  doc["timestamp"] = getCurrentTimestamp();
  doc["reason"] = "unexpected_disconnect";
  
  String payload;
  serializeJson(doc, payload);
  return payload;
}

void updateDisplay() {
  static unsigned long last_display_update = 0;
  
  // Update display every 2 seconds
  if (millis() - last_display_update >= 2000) {
    lcd.clear();
    
    if (alarm_active) {
      lcd.setCursor(0, 0);
      lcd.print("*** ALARM ***");
      lcd.setCursor(0, 1);
      lcd.print("Press to trigger");
    } else {
      // Display temperature and humidity
      lcd.setCursor(0, 0);
      if (last_temperature != 0.0 || last_humidity != 0.0) {
        lcd.print("T:");
        lcd.print(last_temperature, 1);
        lcd.print("C H:");
        lcd.print(last_humidity, 0);
        lcd.print("%");
      } else {
        lcd.print("Sensor reading...");
      }
      
      // Display connection status
      lcd.setCursor(0, 1);
      if (wifi_connected && mqtt_connected) {
        lcd.print("Online ");
      } else if (wifi_connected) {
        lcd.print("WiFi OK ");
      } else {
        lcd.print("Offline ");
      }
      
      // Display current time
      lcd.print(getCurrentTime());
    }
    
    last_display_update = millis();
  }
}

void displayStartupMessage() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("System Ready");
  lcd.setCursor(0, 1);
  lcd.print("IoT Device v1.0");
  delay(2000);
}

String getCurrentTimestamp() {
  struct tm timeinfo;
  if (!getLocalTime(&timeinfo)) {
    return "1970-01-01T00:00:00Z";
  }
  
  char timestamp[32];
  strftime(timestamp, sizeof(timestamp), "%Y-%m-%dT%H:%M:%SZ", &timeinfo);
  return String(timestamp);
}

String getCurrentTime() {
  struct tm timeinfo;
  if (!getLocalTime(&timeinfo)) {
    return "00:00";
  }
  
  char time_str[8];
  strftime(time_str, sizeof(time_str), "%H:%M", &timeinfo);
  return String(time_str);
}

void enterLowPowerMode() {
  Serial.println("Entering low power mode...");
  
  // Publish offline status
  if (mqtt_connected) {
    publishDeviceStatus("sleeping");
    delay(100); // Allow message to be sent
  }
  
  // Turn off LCD backlight
  lcd.noBacklight();
  lcd.clear();
  
  // Configure wake up timer
  esp_sleep_enable_timer_wakeup(sleep_duration * 1000000); // Convert to microseconds
  
  // Enter deep sleep
  esp_deep_sleep_start();
}