#include <WiFi.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <time.h>
#include <RTClib.h>
#include <EEPROM.h>
#include "DHT.h"

// ------------ WIFI ------------
#define WIFI_SSID "GalaxyA34"
#define WIFI_PASSWORD "galaxya34"

// ------------ EEPROM ------------
#define EEPROM_SIZE 4096
#define EEPROM_BLOCK_SIZE 800

// ------------ CHANNELS ------------
#define NUM_CHANNELS 5   // 4 relay/mcb/dimmer + 1 sensor

// ------------ MQTT ------------
const char* mqtt_server = "broker.hivemq.com";
const int mqtt_port = 1883;

WiFiClient espClient;
PubSubClient client(espClient);

// ------------ DEVICE ID ------------
const char* deviceID = "25082024233005";

// ------------ RELAYS ------------
int relayPins[4] = {13, 12, 14, 27};
int manualPins[4] = {26, 25, 33, 32};

bool relayState[4] = {false, false, false, false};

int prevManualState[4] = {HIGH, HIGH, HIGH, HIGH};

// ------------ DHT ------------
#define DHTPIN 4
#define DHTTYPE DHT11
DHT dht(DHTPIN, DHTTYPE);

// ------------ RTC ------------
RTC_DS3231 rtc;

// ------------ MQTT TOPICS ------------
char topic_relay[4][32];
char topic_sensor[32];

// ------------ CHANNEL STRUCT ------------
struct Channel {
  String chanelID;
  String type;
  String currentState;
  int intensity;
  float loadCurrent;
  float temperature;
  float humidity;
};

Channel channels[NUM_CHANNELS];

// ----------------------------------------------------
//                     TIME (EPOCH)
// ----------------------------------------------------
unsigned long long getEpochTime() {
  time_t now;
  time(&now);
  // return milliseconds since epoch
  return (unsigned long long)now * 1000ULL;
}

// ----------------------------------------------------
//                 JSON FORMATTER
// ----------------------------------------------------
String createJSON(const Channel &ch, const char *source) {
  // larger buffer to be safe
  DynamicJsonDocument doc(512);

  doc["chanelID"] = ch.chanelID;
  doc["type"] = ch.type;
  doc["timestamp"] = getEpochTime();
  doc["source"] = source;

  if (ch.type == "relay" || ch.type == "dimmer" || ch.type == "mcb")
    doc["currentState"] = ch.currentState;

  if (ch.type == "dimmer")
    doc["intensity"] = ch.intensity;

  if (ch.type == "mcb")
    doc["loadCurrent"] = ch.loadCurrent;

  if (ch.type == "temp_sensor") {
    doc["temperature"] = ch.temperature;
    doc["humidity"] = ch.humidity;
  }

  String output;
  serializeJson(doc, output);
  return output;
}

// ----------------------------------------------------
//               EEPROM SAVE PER CHANNEL
// ----------------------------------------------------
void saveChanelDataToEEPROM(int chIndex, const String &data) {
  int baseAddr = chIndex * EEPROM_BLOCK_SIZE;
  int len = data.length();

  if (len >= EEPROM_BLOCK_SIZE - 1) {
    // avoid overflow
    len = EEPROM_BLOCK_SIZE - 2;
  }

  for (int i = 0; i < len; i++)
    EEPROM.write(baseAddr + i, data[i]);

  EEPROM.write(baseAddr + len, '\0');
  EEPROM.commit();
}

// ----------------------------------------------------
//               EEPROM READ PER CHANNEL
// ----------------------------------------------------
String readChannelFromEEPROM(int chIndex) {
  int baseAddr = chIndex * EEPROM_BLOCK_SIZE;
  String data = "";

  for (int i = 0; i < EEPROM_BLOCK_SIZE; i++) {
    char c = EEPROM.read(baseAddr + i);
    if (c == '\0') break;
    data += c;
  }
  return data;
}

// ----------------------------------------------------
//           RESTORE RELAYS ON BOOT
// ----------------------------------------------------
void loadRelayStatesFromEEPROM() {
  Serial.print("Stored Data: ");
  for (int i = 0; i < 4; i++) {
    String data = readChannelFromEEPROM(i);
    Serial.print("Index: ");
    Serial.print(i+1);
    Serial.print(" \n ");
    Serial.println(data);
    

    if (data.length() == 0) continue;

    DynamicJsonDocument doc(512);
    DeserializationError err = deserializeJson(doc, data);
    if (err) {
      Serial.print("EEPROM JSON parse error channel ");
      Serial.print(i);
      Serial.print(": ");
      Serial.println(err.c_str());
      continue;
    }

    const char* state = doc["currentState"] | "off";
    relayState[i] = (String(state) == "on");
    digitalWrite(relayPins[i], relayState[i] ? HIGH : LOW);

    channels[i].currentState = String(state);

    Serial.printf("Restored relay %d -> %s\n", i + 1, relayState[i] ? "ON" : "OFF");
  
  }
}

// ----------------------------------------------------
//                  WIFI CONNECT
// ----------------------------------------------------
void connectWiFi() {
  WiFi.mode(WIFI_STA);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
  Serial.print("Connecting to WiFi");

  unsigned long start = millis();
  const unsigned long timeout = 20000; // 20s timeout

  while (WiFi.status() != WL_CONNECTED) {
    delay(300);
    Serial.print(".");
    if (millis() - start > timeout) {
      Serial.println();
      Serial.println("WiFi connect timeout, retrying...");
      start = millis();
      WiFi.disconnect();
      WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
    }
  }
  Serial.println();
  Serial.print("Connected, IP: ");
  Serial.println(WiFi.localIP());
}

// ----------------------------------------------------
//              PUBLISH RELAY STATE (FIXED)
// ----------------------------------------------------
void publishRelayState(int index) {
  String jsonData = createJSON(channels[index], "device");

  String ackTopic = String(topic_relay[index]) + "/ack";
  client.publish(ackTopic.c_str(), jsonData.c_str(), true); // retained ack

  saveChanelDataToEEPROM(index, jsonData);
}

// ----------------------------------------------------
//                  MQTT CONNECT (FIXED)
// ----------------------------------------------------
void connectMQTT() {
  // if already connected, nothing to do
  if (client.connected()) return;

  // Prepare will topic (LWT)
  char willTopic[64];
  snprintf(willTopic, sizeof(willTopic), "%s/lwt", deviceID);
  const char* willMessage = "offline";

  Serial.print("Connecting to MQTT... ");

  // No authentication example — if you do have username/password, replace NULLs accordingly.
  // connect(clientID, username, password, willTopic, willQos, willRetain, willMessage)
  // using willQos = 1, willRetain = true
  if (client.connect(deviceID, NULL, NULL, willTopic, 1, true, willMessage)) {

    Serial.println("connected");


    client.subscribe(deviceID);
    // subscribe to relay channels
    for (int i = 0; i < 4; i++) {
      client.subscribe(topic_relay[i]);
      Serial.print("Subscribed: ");
      Serial.println(topic_relay[i]);

      // publish last known channel payload (if any) to ack topic (retained)
      String ackTopic = String(topic_relay[i]) + "/ack";
      String payload = readChannelFromEEPROM(i);
      if (payload.length() > 0) {
        client.publish(ackTopic.c_str(), payload.c_str(), true);
      }
    }

    // publish device online on a device ack or status topic (retained)
    String devAck = String(deviceID) + "/ack";
    client.publish(devAck.c_str(), "online", true);

  } else {
    Serial.print("failed, rc=");
    Serial.println(client.state());
    // a short delay to avoid tight loop - loop() will call connect again
    delay(1000);
  }
}

// ----------------------------------------------------
//                  MQTT CALLBACK
// ----------------------------------------------------
void callback(char* topic, byte* payload, unsigned int length) {

  char buffer[length + 1];
  memcpy(buffer, payload, length);
  buffer[length] = '\0';

  Serial.print("\nMessage on ");
  Serial.print(topic);
  Serial.print(" : ");
  Serial.println(buffer);

  // device commands (topic == deviceID)
  if (strcmp(topic, deviceID) == 0) {
    if (strcmp(buffer, "reboot") == 0 || strcmp(buffer, "restart") == 0) {
      String ackTopic = String(deviceID) + "/ack";
      client.publish(ackTopic.c_str(), "restarting", false);
      delay(200);
      ESP.restart();
      return;
    }
  }

  // parse JSON safely
  DynamicJsonDocument doc(512);
  DeserializationError err = deserializeJson(doc, buffer);
  if (err) {
    Serial.print("JSON parse error: ");
    Serial.println(err.c_str());
    return;
  }

  const char* state = doc["currentState"];
  const char* type = doc["type"];
  if (!state) return;

  if (strcmp(state, "on") != 0 && strcmp(state, "off") != 0) return;

  // match channel
  for (int i = 0; i < 4; i++) {
    if (strcmp(topic, topic_relay[i]) == 0) {

      bool target = (strcmp(state, "on") == 0);

      relayState[i] = target;
      channels[i].currentState = target ? "on" : "off";

      //if (channels[i].type != "dimmer")
      if(strcmp(type, "mcb") == 0)
        digitalWrite(relayPins[i], target ? HIGH : LOW);

      if(strcmp(type, "relay") == 0)
        digitalWrite(relayPins[i], target ? HIGH : LOW);
      
      if(strcmp(type, "dimmer") == 0)
        digitalWrite(relayPins[i], target ? HIGH : LOW);

      // save and publish ack
      saveChanelDataToEEPROM(i, String(buffer));
      publishRelayState(i);

      Serial.printf("MQTT Relay %d -> %s\n", i + 1, target ? "ON" : "OFF");
      break;
    }
  }
}

// ----------------------------------------------------
//              MANUAL SWITCH HANDLER
// ----------------------------------------------------
/*
void checkManualButtons() {
  static unsigned long lastDebounce[4] = {0};
  static int lastRaw[4] = {HIGH, HIGH, HIGH, HIGH};

  const unsigned long DEBOUNCE_MS = 80;

  for (int i = 0; i < 4; i++) {

    if (channels[i].type == "dimmer") continue;

    int raw = digitalRead(manualPins[i]);
    bool switchOn = (raw == LOW);

    if (raw != lastRaw[i]) {
      lastDebounce[i] = millis();
      lastRaw[i] = raw;
      continue;
    }

    if (millis() - lastDebounce[i] >= DEBOUNCE_MS) {
      if (switchOn != relayState[i]) {

        relayState[i] = switchOn;
        channels[i].currentState = switchOn ? "on" : "off";

        digitalWrite(relayPins[i], switchOn ? HIGH : LOW);
        publishRelayState(i);

        Serial.printf("Manual Relay %d -> %s\n", i + 1, switchOn ? "ON" : "OFF");
      }
    }
  }
}
*/

void checkManualButtons()
{
  for (int i = 0; i < NUM_CHANNELS; i++) 
  {
    int currState = digitalRead(manualPins[i]);
    //Serial.printf("Manual pin %d changed to: %d\n",i + 1,currState);
    if (currState != prevManualState[i]) 
    {
      prevManualState[i] = currState;
      Serial.printf("Manual pin %d changed to: %d\n",i + 1,currState);
      digitalWrite(relayPins[i],currState == 0 ? HIGH : LOW);
    }
  }
}


// ----------------------------------------------------
//                    RTC SYNC
// ----------------------------------------------------
void syncTimeFromRTC() {
  if (!rtc.begin()) {
    Serial.println("RTC not found!");
    return;
  }

  // If RTC lost power, optionally set to compile time (commented)
  // if (rtc.lostPower()) { rtc.adjust(DateTime(F(__DATE__), F(__TIME__))); }

  DateTime now = rtc.now();

  struct tm t;
  t.tm_year = now.year() - 1900;
  t.tm_mon  = now.month() - 1;
  t.tm_mday = now.day();
  t.tm_hour = now.hour();
  t.tm_min  = now.minute();
  t.tm_sec  = now.second();
  t.tm_isdst = -1;

  time_t rtcTime = mktime(&t);
  struct timeval tv = {rtcTime, 0};
  settimeofday(&tv, NULL);

  Serial.println("RTC synced");
}

// ----------------------------------------------------
//                  SENSOR PUBLISH
// ----------------------------------------------------
void publishSensor() {
  float t = dht.readTemperature();
  float h = dht.readHumidity();
  if (isnan(t) || isnan(h)) {
    Serial.println("DHT read failed");
    return;
  }

  channels[4].temperature = t;
  channels[4].humidity = h;

  String jsonData = createJSON(channels[4], "device");
  client.publish(topic_sensor, jsonData.c_str(), true); // retained
  saveChanelDataToEEPROM(4, jsonData);

  Serial.print("Sensor published: ");
  Serial.println(jsonData);
}

// ----------------------------------------------------
//                    SETUP
// ----------------------------------------------------
void setup() {
  Serial.begin(115200);
  delay(100);

  EEPROM.begin(EEPROM_SIZE);

  channels[0] = {"", "relay", "off", 0, 0, 0, 0};
  channels[1] = {"", "relay", "off", 0, 0, 0, 0};
  channels[2] = {"", "dimmer", "off", 0, 0, 0, 0};
  channels[3] = {"", "mcb", "off", 0, 0, 0, 0};
  channels[4] = {"", "temp_sensor", "", 0, 0, 0, 0};


  for (int i = 0; i < 4; i++) {
    pinMode(relayPins[i], OUTPUT);
    pinMode(manualPins[i], INPUT_PULLUP);
    digitalWrite(relayPins[i], LOW);
  }

  // restore previous states (before connecting MQTT)
  loadRelayStatesFromEEPROM();

  for (int i = 0; i < NUM_CHANNELS; i++)
    channels[i].chanelID = String(deviceID) + "C" + String(i + 1);

  for (int i = 0; i < NUM_CHANNELS; i++)
    snprintf(topic_relay[i], sizeof(topic_relay[i]), "%sC%d", deviceID, i + 1);
  snprintf(topic_sensor, sizeof(topic_sensor), "%sC5", deviceID);

  dht.begin();
  Wire.begin();

  // attempt RTC sync (if available)
  syncTimeFromRTC();

  // start WiFi and MQTT
  connectWiFi();
  client.setServer(mqtt_server, mqtt_port);
  client.setCallback(callback);
}

// ----------------------------------------------------
//                      LOOP
// ----------------------------------------------------
void loop() {
  if (WiFi.status() != WL_CONNECTED) {
    connectWiFi();
  }

  if (!client.connected()) {
    connectMQTT();
  }
  client.loop();

  static unsigned long lastSensor = 0;
  if (millis() - lastSensor > 10000) {
    publishSensor();
    lastSensor = millis();
  }

  checkManualButtons();

  // small delay to let network tasks run
  delay(10);
}