#include <WiFi.h>
#include <PubSubClient.h>
#include <Wire.h>
#include <Adafruit_MCP23X17.h>
#include <esp_task_wdt.h>
#include <Preferences.h>
#include <ArduinoJson.h>
#include <ESPAsyncWebServer.h>
#include <AsyncTCP.h>
#include <ArduinoOTA.h>

// =====================================================
// CONFIG
// =====================================================
const char* ssid         = "YOUR_WIFI";
const char* password     = "YOUR_PASS";

const char* mqtt_server  = "192.168.1.10";
const uint16_t mqtt_port = 1883;

const char* hostname     = "esp32-12ch-relay";

#define FW_VERSION              "v13.4-PRO"

#define TOTAL_RELAYS            12
#define TOTAL_INPUTS            12

#define WDT_TIMEOUT             15

#define I2C_SDA                 21
#define I2C_SCL                 22
#define I2C_FREQ                100000UL
#define I2C_TIMEOUT_MS          50

#define MCP_RELAY_ADDR          0x20
#define MCP_INPUT_ADDR          0x21

#define RELAY_ON                LOW
#define RELAY_OFF               HIGH

#define INPUT_SCAN_INTERVAL     10
#define INPUT_DEBOUNCE          80
#define MIN_RELAY_INTERVAL      150UL

#define NVS_SAVE_DELAY          30000UL
#define MIN_NVS_WRITE_INTERVAL  120000UL
#define TELEMETRY_INTERVAL      60000UL
#define MQTT_FAILSAFE_TIMEOUT   300000UL   // 5 menit tanpa activity
#define MCP_HEALTH_INTERVAL     30000UL
#define MCP_FAILURE_THRESHOLD   5

#define SAFE_MODE_PIN           0

// Security
const char* webUser = "admin";
const char* webPass = "admin123";
const char* otaPassword = "CHANGE_ME_NOW_strong_password";

// MQTT Topics
#define MQTT_CMD_BASE   "home/cmd/relay/"
#define MQTT_ACK_BASE   "home/ack/relay/"
#define MQTT_STAT_BASE  "home/stat/relay/"
#define MQTT_INPUT_BASE "home/input/"

// MCP Registers
#define MCP_IOCON       0x0A
#define MCP_OLATA       0x14

// =====================================================
// OBJECTS
// =====================================================
WiFiClient espClient;
PubSubClient client(espClient);

Adafruit_MCP23X17 mcpRelay;
Adafruit_MCP23X17 mcpInput;

AsyncWebServer server(80);
Preferences prefs;

SemaphoreHandle_t relayMutex = NULL;
SemaphoreHandle_t i2cMutex   = NULL;

// =====================================================
// DATA
// =====================================================
bool relayState[TOTAL_RELAYS];
bool relayDirty[TOTAL_RELAYS];
const uint8_t relayPin[TOTAL_RELAYS] = {0,1,2,3,4,5,6,7,8,9,10,11};

bool rawInput[TOTAL_INPUTS];
bool stableInput[TOTAL_INPUTS];
unsigned long debounceTimer[TOTAL_INPUTS];
unsigned long lastRelayChange[TOTAL_RELAYS];

uint16_t interlockMask[TOTAL_RELAYS];

enum RestorePolicy { RESTORE, FORCE_OFF, FORCE_ON };
enum FailSafePolicy { HOLD_LAST, FORCE_OFF_ON_FAIL, FORCE_ON_ON_FAIL };

RestorePolicy restorePolicy[TOTAL_RELAYS];
FailSafePolicy failSafePolicy[TOTAL_RELAYS];

// =====================================================
// GLOBAL
// =====================================================
char mqttClientId[40];

bool dirtyPending = false;
bool safeMode = false;

unsigned long lastWiFiCheck = 0;
unsigned long lastMQTTAttempt = 0;
unsigned long lastTelemetry = 0;
unsigned long lastInputScan = 0;
unsigned long nextSaveTime = 0;
unsigned long lastNvsWrite = 0;
unsigned long lastMQTTActivity = 0;
unsigned long lastMCPCheck = 0;

uint8_t mcpFailureCount = 0;

uint32_t mqttBackoff = 3000;

struct Telemetry {
  uint32_t reconnectCount = 0;
  uint32_t publishFail = 0;
  uint32_t minFreeHeap = 0;
} telemetry;

// =====================================================
// HELPER
// =====================================================
void generateClientID() {
  uint8_t mac[6];
  WiFi.macAddress(mac);
  snprintf(mqttClientId, sizeof(mqttClientId), "ESP32_%02X%02X%02X", mac[3], mac[4], mac[5]);
}

bool recoverI2C() {
  Serial.println("I2C Recovery...");

  bool locked = false;
  if (i2cMutex) locked = (xSemaphoreTake(i2cMutex, pdMS_TO_TICKS(100)) == pdTRUE);

  pinMode(I2C_SDA, INPUT_PULLUP);
  pinMode(I2C_SCL, OUTPUT_OPEN_DRAIN);

  for (uint8_t i = 0; i < 9; i++) {
    digitalWrite(I2C_SCL, LOW); delayMicroseconds(5);
    digitalWrite(I2C_SCL, HIGH); delayMicroseconds(5);
  }

  // STOP condition
  digitalWrite(I2C_SDA, LOW); delayMicroseconds(5);
  digitalWrite(I2C_SCL, HIGH); delayMicroseconds(5);
  digitalWrite(I2C_SDA, HIGH); delayMicroseconds(5);

  Wire.begin(I2C_SDA, I2C_SCL);
  Wire.setClock(I2C_FREQ);
  Wire.setTimeout(I2C_TIMEOUT_MS);

  if (locked) xSemaphoreGive(i2cMutex);
  return true;
}

bool checkMCPHealth() {
  if (!i2cMutex || xSemaphoreTake(i2cMutex, pdMS_TO_TICKS(50)) != pdTRUE) return false;

  Wire.beginTransmission(MCP_RELAY_ADDR);
  Wire.write(MCP_IOCON);
  if (Wire.endTransmission(false) != 0) {
    xSemaphoreGive(i2cMutex);
    return false;
  }

  Wire.requestFrom(MCP_RELAY_ADDR, 1);
  bool ok = Wire.available() && (Wire.read() != 0xFF);

  xSemaphoreGive(i2cMutex);
  return ok;
}

void reconfigureMCP() {
  if (i2cMutex) xSemaphoreTake(i2cMutex, pdMS_TO_TICKS(200));

  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    mcpRelay.pinMode(relayPin[i], OUTPUT);
    mcpRelay.digitalWrite(relayPin[i], relayState[i] ? RELAY_ON : RELAY_OFF);
  }
  for (uint8_t i = 0; i < TOTAL_INPUTS; i++) {
    mcpInput.pinMode(i, INPUT_PULLUP);
  }

  if (i2cMutex) xSemaphoreGive(i2cMutex);
}

void loadRelayState() {
  prefs.begin("relay", true);
  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    char key[8];
    snprintf(key, sizeof(key), "r%d", i);
    relayState[i] = prefs.getBool(key, false);

    snprintf(key, sizeof(key), "p%d", i);
    restorePolicy[i] = (RestorePolicy)prefs.getUChar(key, RESTORE);

    snprintf(key, sizeof(key), "f%d", i);
    failSafePolicy[i] = (FailSafePolicy)prefs.getUChar(key, HOLD_LAST);

    relayDirty[i] = false;
    lastRelayChange[i] = 0;
  }
  prefs.end();
}

void saveDirtyStates() {
  if (!dirtyPending || (millis() - lastNvsWrite < MIN_NVS_WRITE_INTERVAL)) return;

  prefs.begin("relay", false);
  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    if (relayDirty[i]) {
      char key[8];
      snprintf(key, sizeof(key), "r%d", i);
      prefs.putBool(key, relayState[i]);
      relayDirty[i] = false;
    }
  }
  prefs.end();

  dirtyPending = false;
  lastNvsWrite = millis();
  Serial.println("NVS SAVED");
}

// =====================================================
// RELAY
// =====================================================
bool interlockAllowed(uint8_t ch) {
  uint16_t mask = interlockMask[ch];
  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    if ((mask & (1 << i)) && relayState[i]) return false;
  }
  return true;
}

bool writeRelayHardware(uint8_t ch, bool state) {
  if (ch >= TOTAL_RELAYS) return false;

  for (uint8_t retry = 0; retry < 3; retry++) {
    if (i2cMutex && xSemaphoreTake(i2cMutex, pdMS_TO_TICKS(100)) != pdTRUE) {
      delay(5); continue;
    }

    mcpRelay.digitalWrite(relayPin[ch], state ? RELAY_ON : RELAY_OFF);

    if (i2cMutex) xSemaphoreGive(i2cMutex);

    delay(5);

    if (checkMCPHealth()) {
      relayState[ch] = state;
      return true;
    }

    recoverI2C();
    delay(10);
  }
  return false;
}

bool applyRelay(uint8_t ch, bool state) {
  if (xSemaphoreTakeRecursive(relayMutex, pdMS_TO_TICKS(150)) != pdTRUE) return false;

  bool success = false;
  do {
    if (ch >= TOTAL_RELAYS) break;
    if (relayState[ch] == state) { success = true; break; }
    if (state && !interlockAllowed(ch)) break;
    if (millis() - lastRelayChange[ch] < MIN_RELAY_INTERVAL) break;

    if (writeRelayHardware(ch, state)) {
      relayDirty[ch] = true;
      dirtyPending = true;
      nextSaveTime = millis() + NVS_SAVE_DELAY;
      lastRelayChange[ch] = millis();
      success = true;
    }
  } while(0);

  xSemaphoreGiveRecursive(relayMutex);
  return success;
}

bool setRelay(uint8_t ch, bool state) {
  bool success = applyRelay(ch, state);
  publishRelay(ch);

  char topic[64];
  snprintf(topic, sizeof(topic), "%s%d", MQTT_ACK_BASE, ch + 1);
  const char* ack = success ? "OK" : "BLOCKED";
  client.publish(topic, ack, false);

  return success;
}

void allRelay(bool state) {
  if (xSemaphoreTakeRecursive(relayMutex, pdMS_TO_TICKS(300)) != pdTRUE) return;

  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    applyRelay(i, state);
  }
  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    publishRelay(i);
  }
  xSemaphoreGiveRecursive(relayMutex);
}

// =====================================================
// FAILSAFE
// =====================================================
void applyFailsafe() {
  Serial.println("MQTT Failsafe Activated");
  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    if (failSafePolicy[i] == FORCE_OFF_ON_FAIL) setRelay(i, false);
    else if (failSafePolicy[i] == FORCE_ON_ON_FAIL) setRelay(i, true);
  }
}

// =====================================================
// TELEMETRY
// =====================================================
void publishTelemetry() {
  if (!client.connected()) return;

  StaticJsonDocument<256> doc;
  doc["fw"] = FW_VERSION;
  doc["uptime"] = millis() / 1000;
  doc["heap"] = ESP.getFreeHeap();
  doc["min_heap"] = ESP.getMinFreeHeap();
  doc["stack"] = uxTaskGetStackHighWaterMark(NULL);
  doc["rssi"] = WiFi.RSSI();
  doc["reconnect"] = telemetry.reconnectCount;
  doc["pub_fail"] = telemetry.publishFail;
  doc["reset"] = esp_reset_reason();

  char payload[256];
  serializeJson(doc, payload);
  client.publish("home/stat/system", payload, false);
}

// =====================================================
// MQTT CALLBACK
// =====================================================
void callback(char* topic, byte* payload, unsigned int length) {
  char msg[16] = {0};
  size_t len = min(length, (unsigned int)(sizeof(msg) - 1));
  memcpy(msg, payload, len);
  msg[len] = '\0';

  lastMQTTActivity = millis();

  for (size_t i = 0; i < len; i++) {
    if (msg[i] >= 'a' && msg[i] <= 'z') msg[i] -= 32;
  }

  if (strcmp(topic, "home/cmd/relay/all") == 0) {
    if (strcmp(msg, "ON") == 0) allRelay(true);
    else if (strcmp(msg, "OFF") == 0) allRelay(false);
    return;
  }

  if (strncmp(topic, MQTT_CMD_BASE, strlen(MQTT_CMD_BASE)) == 0) {
    int ch = atoi(topic + strlen(MQTT_CMD_BASE)) - 1;
    if (ch < 0 || ch >= TOTAL_RELAYS) return;

    if (strcmp(msg, "ON") == 0) setRelay(ch, true);
    else if (strcmp(msg, "OFF") == 0) setRelay(ch, false);
  }
}

// =====================================================
// INPUT SCAN
// =====================================================
void scanInputs() {
  unsigned long now = millis();
  if (now - lastInputScan < INPUT_SCAN_INTERVAL) return;
  lastInputScan = now;

  for (uint8_t i = 0; i < TOTAL_INPUTS; i++) {
    bool reading = false;
    if (i2cMutex && xSemaphoreTake(i2cMutex, pdMS_TO_TICKS(20)) == pdTRUE) {
      reading = !mcpInput.digitalRead(i);
      xSemaphoreGive(i2cMutex);
    }

    if (reading != rawInput[i]) {
      rawInput[i] = reading;
      debounceTimer[i] = now;
    }

    if ((now - debounceTimer[i]) > INPUT_DEBOUNCE) {
      if (stableInput[i] != reading) {
        stableInput[i] = reading;
        if (stableInput[i]) setRelay(i, !relayState[i]);
      }
    }
  }
}

// =====================================================
// MCP HEALTH
// =====================================================
void checkMCPHealthPeriodic() {
  static unsigned long last = 0;
  if (millis() - last < MCP_HEALTH_INTERVAL) return;
  last = millis();

  if (!checkMCPHealth()) {
    mcpFailureCount++;
    Serial.printf("MCP Health FAILED (%d/%d)\n", mcpFailureCount, MCP_FAILURE_THRESHOLD);

    recoverI2C();
    mcpRelay.begin_I2C(MCP_RELAY_ADDR);
    mcpInput.begin_I2C(MCP_INPUT_ADDR);
    reconfigureMCP();

    if (mcpFailureCount >= MCP_FAILURE_THRESHOLD) {
      Serial.println("MCP Failure Threshold Exceeded - Restarting");
      ESP.restart();
    }
  } else {
    mcpFailureCount = 0;
  }
}

// =====================================================
// CONNECTIVITY
// =====================================================
void handleConnectivity() {
  unsigned long now = millis();

  if (now - lastWiFiCheck > 5000) {
    lastWiFiCheck = now;
    if (WiFi.status() != WL_CONNECTED) WiFi.reconnect();
  }

  if (WiFi.status() == WL_CONNECTED) {
    if (!client.connected() && (now - lastMQTTAttempt > mqttBackoff)) {
      lastMQTTAttempt = now;
      if (client.connect(mqttClientId, "home/stat/LWT", 0, true, "OFFLINE")) {
        telemetry.reconnectCount++;
        client.publish("home/stat/LWT", "ONLINE", true);

        for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
          char t[64]; snprintf(t, sizeof(t), "%s%d", MQTT_CMD_BASE, i + 1);
          client.subscribe(t);
        }
        client.subscribe("home/cmd/relay/all");

        for (uint8_t i = 0; i < TOTAL_RELAYS; i++) publishRelay(i);
        publishTelemetry();
        mqttBackoff = 3000;
      } else {
        mqttBackoff = min(mqttBackoff * 2, 60000UL);
      }
    }
  }
}

// =====================================================
// WEB + OTA
// =====================================================
void setupWeb() {
  server.on("/", HTTP_GET, [](AsyncWebServerRequest *request) {
    if (!request->authenticate(webUser, webPass)) return request->requestAuthentication();
    AsyncResponseStream *response = request->beginResponseStream("text/html");
    response->print("<html><body><h1>ESP32 Industrial Relay v13.3-PRO</h1>");
    for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
      response->printf("Relay %d : %s<br>", i+1, relayState[i] ? "ON" : "OFF");
    }
    response->print("</body></html>");
    request->send(response);
  });
  server.begin();
}

void setupOTA() {
  ArduinoOTA.setHostname(hostname);
  ArduinoOTA.setPassword(otaPassword);
  ArduinoOTA.begin();
}

// =====================================================
// SETUP
// =====================================================
void setup() {
  Serial.begin(115200);
  delay(1000);
  Serial.println("\n=== ESP32 12CH INDUSTRIAL RELAY CONTROLLER v13.3-PRO ===");

  esp_task_wdt_init(WDT_TIMEOUT, true);
  esp_task_wdt_add(NULL);

  relayMutex = xSemaphoreCreateRecursiveMutex();
  i2cMutex   = xSemaphoreCreateMutex();

  pinMode(SAFE_MODE_PIN, INPUT_PULLUP);
  safeMode = (digitalRead(SAFE_MODE_PIN) == LOW);
  if (safeMode) Serial.println("SAFE MODE ACTIVE");

  Wire.begin(I2C_SDA, I2C_SCL);
  Wire.setClock(I2C_FREQ);
  Wire.setTimeout(I2C_TIMEOUT_MS);

  if (!mcpRelay.begin_I2C(MCP_RELAY_ADDR)) {
    recoverI2C();
    if (!mcpRelay.begin_I2C(MCP_RELAY_ADDR)) {
      Serial.println("MCP Relay Fatal!");
      while(1) esp_task_wdt_reset();
    }
  }
  if (!mcpInput.begin_I2C(MCP_INPUT_ADDR)) {
    Serial.println("MCP Input Fatal!");
    while(1) esp_task_wdt_reset();
  }

  for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
    mcpRelay.pinMode(relayPin[i], OUTPUT);
    mcpRelay.digitalWrite(relayPin[i], RELAY_OFF);
  }
  for (uint8_t i = 0; i < TOTAL_INPUTS; i++) {
    mcpInput.pinMode(i, INPUT_PULLUP);
    bool val = !mcpInput.digitalRead(i);
    rawInput[i] = val;
    stableInput[i] = val;
    debounceTimer[i] = millis();
  }

  loadRelayState();

  if (!safeMode) {
    for (uint8_t i = 0; i < TOTAL_RELAYS; i++) {
      bool state = relayState[i];
      if (restorePolicy[i] == FORCE_OFF) state = false;
      if (restorePolicy[i] == FORCE_ON) state = true;
      writeRelayHardware(i, state);
    }
  } else {
    for (uint8_t i = 0; i < TOTAL_RELAYS; i++) writeRelayHardware(i, false);
  }

  generateClientID();
  WiFi.setHostname(hostname);
  WiFi.mode(WIFI_STA);
  WiFi.begin(ssid, password);

  client.setServer(mqtt_server, mqtt_port);
  client.setCallback(callback);
  client.setKeepAlive(60);
  client.setBufferSize(1024);

  interlockMask[0] |= (1 << 1);
  interlockMask[1] |= (1 << 0);

  setupWeb();
  if (!safeMode) setupOTA();

  Serial.println("SYSTEM READY");
}

// =====================================================
// LOOP
// =====================================================
void loop() {
  esp_task_wdt_reset();

  handleConnectivity();
  if (client.connected()) client.loop();

  if (!safeMode) {
    ArduinoOTA.handle();
    scanInputs();
    checkMCPHealthPeriodic();
  }

  if (dirtyPending && millis() >= nextSaveTime) saveDirtyStates();

  if (millis() - lastTelemetry > TELEMETRY_INTERVAL) {
    lastTelemetry = millis();
    publishTelemetry();
  }

  if (!safeMode && lastMQTTActivity != 0 && 
      (millis() - lastMQTTActivity > MQTT_FAILSAFE_TIMEOUT)) {
    applyFailsafe();
    lastMQTTActivity = 0;
  }

  vTaskDelay(1);
}