#include <WiFi.h>
#include <PubSubClient.h>
#include <IRremote.h>
#include <DHT.h>
#include <Wire.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_GFX.h>

// ===== OLED Configuration =====
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_SDA 33
#define OLED_SCL 32
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire);

// ===== IoT Configuration =====
const char* ssid = "Wokwi-GUEST";
const char* password = "";
const char* mqttServer = "demo.thingsboard.io";
const char* token = "gLpxYXede8M3Ps6cd2Qy";
const int mqttPort = 1883;

// Pin Definitions
#define IR_RECEIVE_PIN 35
#define DHT_PIN 25
#define DHT_TYPE DHT22
#define MOTION_PIN 34
#define MQ2_ANALOG_PIN 33 // Shared with OLED SDA
#define MQ2_DIGITAL_PIN 32 // Shared with OLED SCL
const int relayPins[8] = {23, 22, 21, 4, 2, 15, 16, 17};

// Timing Intervals (ms)
const unsigned long dhtInterval = 2000;
const unsigned long motionDebounce = 500;
const unsigned long statusUpdateInterval = 5000;
const unsigned long displayRefreshInterval = 1000;

// State Variables
bool relayStates[8] = {false};
bool motionDetected = false;
float currentTemp = 0;
float currentHum = 0;
bool gasDetected = false;
unsigned long lastDisplayUpdate = 0;
String serialBuffer = ""; // Buffer for serial messages

// Global Objects
WiFiClient espClient;
PubSubClient client(espClient);
IRrecv irReceiver(IR_RECEIVE_PIN);
DHT dht(DHT_PIN, DHT_TYPE);

void setup() {
  Serial.begin(115200);
  
  // Initialize OLED
  Wire.begin(OLED_SDA, OLED_SCL);
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED allocation failed");
    while(1);
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(0,0);
  display.println("Initializing...");
  display.display();
  
  // Initialize components
  initWiFi();
  initSensors();
  initRelays();
  initIRReceiver();
  
  // MQTT setup
  client.setBufferSize(2048);
  client.setServer(mqttServer, mqttPort);
  client.setCallback(mqttCallback);
  
  // Initial display
  updateDisplay(true);
}

void loop() {
  // Maintain MQTT connection
  if (!client.connected()) reconnectMQTT();
  client.loop();

  // Handle sensors
  handleDHT();
  handleMotion();
  handleMQ2();
  handleIRRemote();
  
  // Update display periodically
  if (millis() - lastDisplayUpdate >= displayRefreshInterval) {
    updateDisplay(false);
  }
  
  // Handle serial input for OLED display
  handleSerialInput();
}

void handleSerialInput() {
  while (Serial.available()) {
    char c = Serial.read();
    if (c == '\n') {
      // Display serial message on OLED
      display.clearDisplay();
      display.setCursor(0,0);
      display.println("Serial Message:");
      display.println(serialBuffer);
      display.display();
      serialBuffer = "";
      lastDisplayUpdate = millis();
    } else {
      serialBuffer += c;
    }
  }
}

// ===== Display Functions =====
void updateDisplay(bool forceUpdate) {
  static float lastShownTemp = -100;
  static float lastShownHum = -100;
  static bool lastShownMotion = false;
  static bool lastShownGas = false;
  static bool lastShownRelays[8] = {false};
  
  // Check if any value changed
  bool needsUpdate = forceUpdate ||
                    (currentTemp != lastShownTemp) ||
                    (currentHum != lastShownHum) ||
                    (motionDetected != lastShownMotion) ||
                    (gasDetected != lastShownGas);
  
  // Check relay states
  for (int i = 0; i < 8 && !needsUpdate; i++) {
    if (relayStates[i] != lastShownRelays[i]) {
      needsUpdate = true;
    }
  }
  
  if (!needsUpdate) return;
  
  // Update last shown values
  lastShownTemp = currentTemp;
  lastShownHum = currentHum;
  lastShownMotion = motionDetected;
  lastShownGas = gasDetected;
  for (int i = 0; i < 8; i++) {
    lastShownRelays[i] = relayStates[i];
  }
  
  // Prepare serial output
  String serialOutput = "\n=== System Status ===\n";
  serialOutput += "Temperature: " + String(currentTemp, 1) + "°C\n";
  serialOutput += "Humidity: " + String(currentHum, 1) + "%\n";
  serialOutput += "Motion: " + String(motionDetected ? "Detected" : "Clear") + "\n";
  serialOutput += "Gas: " + String(gasDetected ? "Detected" : "Clear") + "\n";
  serialOutput += "Relay States:\n";
  for (int i = 0; i < 8; i++) {
    serialOutput += "  Relay " + String(i+1) + ": " + (relayStates[i] ? "ON" : "OFF") + "\n";
  }
  serialOutput += "=====================";
  
  // Send to Serial Monitor
  Serial.println(serialOutput);
  
  // Update OLED Display
  display.clearDisplay();
  display.setCursor(0,0);
  
  // Line 1: Temp/Hum
  display.printf("T:%.1fC H:%.1f%%", currentTemp, currentHum);
  
  // Line 2: Motion/Gas
  display.setCursor(0, 16);
  display.print("M:");
  display.print(motionDetected ? "Y " : "N ");
  display.print("G:");
  display.print(gasDetected ? "Y" : "N");
  
  // Relay Status (2 rows of 4)
  display.setCursor(0, 32);
  for (int i = 0; i < 4; i++) {
    display.printf("%d:%s ", i+1, relayStates[i] ? "ON" : "OFF");
  }
  
  display.setCursor(0, 48);
  for (int i = 4; i < 8; i++) {
    display.printf("%d:%s ", i+1, relayStates[i] ? "ON" : "OFF");
  }
  
  display.display();
  lastDisplayUpdate = millis();
}

// ===== Sensor Handlers =====
void handleDHT() {
  static unsigned long lastRead = 0;
  if (millis() - lastRead >= dhtInterval) {
    lastRead = millis();
    
    float temp = dht.readTemperature();
    float hum = dht.readHumidity();
    
    if (!isnan(temp) && !isnan(hum)) {
      if (temp != currentTemp || hum != currentHum) {
        currentTemp = temp;
        currentHum = hum;
        sendTelemetry("temperature", String(temp, 1));
        sendTelemetry("humidity", String(hum, 1));
        updateDisplay(true);
      }
    }
  }
}

void handleMotion() {
  static unsigned long lastDebounce = 0;
  bool currentState = digitalRead(MOTION_PIN);
  
  if (millis() - lastDebounce > motionDebounce) {
    if (currentState != motionDetected) {
      motionDetected = currentState;
      setRelay(5, motionDetected);
      sendTelemetry("motion", motionDetected ? "true" : "false");
      updateDisplay(true);
    }
    lastDebounce = millis();
  }
}

void handleMQ2() {
  static unsigned long lastRead = 0;
  if (millis() - lastRead >= 1000) {
    lastRead = millis();
    bool currentGas = digitalRead(MQ2_DIGITAL_PIN);
    if (currentGas != gasDetected) {
      gasDetected = currentGas;
      sendTelemetry("gas_detected", gasDetected ? "true" : "false");
      updateDisplay(true);
    }
  }
}

void handleIRRemote() {
  if (irReceiver.decode()) {
    unsigned long command = irReceiver.decodedIRData.command;
    String irMessage = "IR Command: 0x" + String(command, HEX);
    Serial.println(irMessage);
    
    // Show IR command on OLED for 2 seconds
    display.clearDisplay();
    display.setCursor(0,0);
    display.println(irMessage);
    display.display();
    delay(2000);
    updateDisplay(true);
    
    switch (command) {
      case 0x30: toggleRelay(0); break;
      case 0x18: toggleRelay(1); break;
      case 0x7A: toggleRelay(2); break;
      case 0x10: toggleRelay(3); break;
      case 0x38: toggleRelay(4); break;
      case 0x5A: toggleRelay(5); break;
      case 0x42: toggleRelay(6); break;
      case 0x4A: toggleRelay(7); break;
    }
    irReceiver.resume();
  }
}

// ===== Relay Control =====
void setRelay(int index, bool state) {
  if (index < 0 || index >= 8) return;
  
  relayStates[index] = state;
  digitalWrite(relayPins[index], state ? HIGH : LOW);
  sendRelayStatus(index);
  updateDisplay(true);
}

void toggleRelay(int index) {
  setRelay(index, !relayStates[index]);
}

// ===== MQTT Functions =====
void reconnectMQTT() {
  static unsigned long lastAttempt = 0;
  if (millis() - lastAttempt < 5000) return;
  lastAttempt = millis();
  
  Serial.print("Connecting to MQTT...");
  if (client.connect("ESP32_Client", token, NULL)) {
    Serial.println("connected");
    client.subscribe("v1/devices/me/rpc/request/+");
  } else {
    Serial.print("failed, rc=");
    Serial.println(client.state());
  }
}

void mqttCallback(char* topic, byte* payload, unsigned int length) {
  payload[length] = '\0';
  String message = String((char*)payload);
  String mqttMessage = "MQTT RX: " + message;
  Serial.println(mqttMessage);
  
  // Show MQTT message on OLED for 2 seconds
  display.clearDisplay();
  display.setCursor(0,0);
  display.println(mqttMessage);
  display.display();
  delay(2000);
  updateDisplay(true);

  for (int i = 0; i < 8; i++) {
    String target = "\"method\":\"setRelay" + String(i+1) + "\"";
    if (message.indexOf(target) != -1) {
      bool state = message.indexOf("\"params\":true") != -1;
      setRelay(i, state);
      break;
    }
  }
}

void sendTelemetry(const String& key, const String& value) {
  if (!client.connected()) return;
  
  String payload = "{\"" + key + "\":" + value + "}";
  String txMessage = "MQTT TX: " + payload;
  
  if (client.publish("v1/devices/me/telemetry", payload.c_str())) {
    Serial.println(txMessage);
    
    // Show MQTT TX on OLED for 1 second
    display.clearDisplay();
    display.setCursor(0,0);
    display.println(txMessage);
    display.display();
    delay(1000);
    updateDisplay(true);
  } else {
    Serial.println("MQTT TX failed!");
  }
}

void sendRelayStatus(int index) {
  String key = "relay" + String(index + 1);
  sendTelemetry(key, relayStates[index] ? "true" : "false");
}

// ===== Initialization Functions =====
void initWiFi() {
  Serial.print("Connecting to WiFi...");
  WiFi.begin(ssid, password);
  
  unsigned long startTime = millis();
  while (WiFi.status() != WL_CONNECTED && millis() - startTime < 10000) {
    delay(500);
    Serial.print(".");
  }
  
  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("\nWiFi connected");
    Serial.print("IP: ");
    Serial.println(WiFi.localIP());
  } else {
    Serial.println("\nWiFi failed!");
  }
}

void initSensors() {
  dht.begin();
  pinMode(MOTION_PIN, INPUT_PULLUP);
  pinMode(MQ2_DIGITAL_PIN, INPUT);
  Serial.println("Sensors initialized");
}

void initRelays() {
  for (int i = 0; i < 8; i++) {
    pinMode(relayPins[i], OUTPUT);
    digitalWrite(relayPins[i], LOW);
  }
  Serial.println("Relays initialized");
}

void initIRReceiver() {
  irReceiver.enableIRIn();
  Serial.println("IR Receiver ready");
}