#include <Arduino.h>
#include "DHT.h"
#include <WiFi.h>
#include <PubSubClient.h>
#include <esp_task_wdt.h>
#include <ArduinoJson.h>

// ══════════════════════════════════════════════════════════════
// RELAY CONFIGURATION — READ BEFORE FLASHING
// Most single channel relay modules    → Active HIGH (default)
// Multi channel relay boards (2/4/8ch) → May need Active LOW
// If humidifier turns ON at boot unexpectedly → swap the values
// ══════════════════════════════════════════════════════════════
#define RELAY_ON  HIGH   // ← CHANGE TO LOW  if your relay is active LOW
#define RELAY_OFF LOW    // ← CHANGE TO HIGH if your relay is active LOW

// ── WiFi Credentials ───────────────────────────────────────────
#define WIFI_SSID        "Wokwi-GUEST"
#define WIFI_PASSWORD    ""

// ── MQTT Broker ────────────────────────────────────────────────
#define MQTT_BROKER      "test.mosquitto.org"
#define MQTT_PORT        1883

// ── MQTT Topics — Structured ───────────────────────────────────
#define TOPIC_TELEMETRY        "greenhouse/room1/telemetry"            // Full JSON payload
#define TOPIC_FAULT            "greenhouse/room1/fault"                // Fault and recovery events
#define TOPIC_SYSTEM           "greenhouse/room1/system"               // Heap, RSSI, uptime
#define TOPIC_STATUS           "greenhouse/room1/status"               // LWT: online/offline
#define TOPIC_COMMAND          "greenhouse/room1/command"              // Incoming commands

// ── MQTT Topics — Flat (for panel apps) ────────────────────────
#define TOPIC_TEMP             "greenhouse/room1/telemetry/temp"       // e.g. "27.30"
#define TOPIC_HUMIDITY         "greenhouse/room1/telemetry/humidity"   // e.g. "65.20"
#define TOPIC_HUMIDIFIER_STATE "greenhouse/room1/telemetry/humidifier" // "true"/"false" retained

// ── Pin Definitions ────────────────────────────────────────────
#define DHTPIN           4      // GPIO pin for DHT22 data line
#define DHTTYPE          DHT22
#define HUMIDIFIER_PIN   21     // GPIO pin for relay module

// ── Timing ─────────────────────────────────────────────────────
#define BAUD_RATE        115200
#define READ_INTERVAL    2000   // Sensor read interval (ms)
#define SENSOR_TIMEOUT   30000  // Max time without valid reading before failsafe (ms)
#define SYSTEM_INTERVAL  60000  // System metrics publish interval (ms)
#define HEARTBEAT_MAX    30000  // Force publish telemetry every 30s even if no change
#define WIFI_TIMEOUT     15000  // Max wait for WiFi at startup (ms)
#define MQTT_TIMEOUT     15000  // Max wait for MQTT at startup (ms)

// ── Humidity Thresholds (runtime updatable via MQTT) ───────────
#define HUMIDITY_ON_THRESHOLD_DEFAULT  70.0  // Humidifier triggers below this
#define HUMIDITY_OFF_THRESHOLD_DEFAULT 90.0  // Humidifier stops above this
#define HUMIDITY_MIN_VALID             40.0  // Minimum accepted threshold value
#define HUMIDITY_MAX_VALID             99.0  // Maximum accepted threshold value
#define TRIGGER_DELAY                  10000 // Humidity must stay low for 10s to trigger

// ── Publish change threshold ───────────────────────────────────
#define PUBLISH_DELTA    0.5    // Minimum change to trigger publish

// ── Exponential Backoff ────────────────────────────────────────
#define BACKOFF_MIN      2000   // Start at 2s
#define BACKOFF_MAX      60000  // Cap at 60s

// ── MQTT Payload Safety ────────────────────────────────────────
// Max accepted command payload size in bytes
// Commands exceeding this are rejected to prevent stack overflow
// Max valid command: {"override":"FORCED_ON","threshold_on":99,"threshold_off":99} = ~60 bytes
// 128 bytes gives comfortable headroom
#define MQTT_MAX_PAYLOAD 128

// ══════════════════════════════════════════════════════════════
// CONTROL MODE
// AUTO       → sensor and thresholds decide
// FORCED_ON  → manual override, humidifier ON regardless
// FORCED_OFF → manual override, humidifier OFF regardless
// ══════════════════════════════════════════════════════════════
enum ControlMode {
  AUTO,
  FORCED_ON,
  FORCED_OFF
};

DHT dht(DHTPIN, DHTTYPE);
WiFiClient espClient;
PubSubClient mqtt(espClient);

// ── Runtime State ──────────────────────────────────────────────
ControlMode   mode              = AUTO;   // Current control mode
bool          humidifierOn      = false;  // Actual relay output state
bool          sensorFaultActive = false;  // True if sensor failing >30s
bool          lowHumTimerActive = false;  // True if 10s trigger countdown running

// ── Thresholds (modifiable via MQTT) ──────────────────────────
float humidityOnThreshold  = HUMIDITY_ON_THRESHOLD_DEFAULT;
float humidityOffThreshold = HUMIDITY_OFF_THRESHOLD_DEFAULT;

// ── Timestamps ─────────────────────────────────────────────────
unsigned long lastReadTime         = 0;
unsigned long lastValidReadTime    = 0;
unsigned long lastSystemPublish    = 0;
unsigned long lastTelemetryPublish = 0;
unsigned long lowHumStartTime      = 0;
unsigned long lastReconnectAttempt = 0;
unsigned long backoffDelay         = BACKOFF_MIN;
unsigned long bootTime             = 0;

// ── Last Published Values (delta tracking) ────────────────────
float lastPublishedTemp     = -999;
float lastPublishedHumidity = -999;
bool  lastPublishedHumState = false;

// ── Forward Declarations ───────────────────────────────────────
void setHumidifier(bool state, const char* reason);
void publishTelemetry(float tempC, float humidity);
void publishSystem();
void publishFault(const char* message);
void handleCommand(const char* payload);

// ── Generate unique MQTT client ID — static buffer, no String ──
void getMQTTClientId(char* buffer, size_t size) {
  uint64_t mac = ESP.getEfuseMac();
  snprintf(buffer, size, "ESP32_Humidifier_%04X%08X",
    (uint16_t)(mac >> 32),
    (uint32_t)mac
  );
}

// ── Set relay output and track state ──────────────────────────
void setHumidifier(bool state, const char* reason) {
  if (humidifierOn == state) return; // No change — skip
  humidifierOn = state;
  digitalWrite(HUMIDIFIER_PIN, state ? RELAY_ON : RELAY_OFF);

  // Publish retained humidifier state immediately on change
  // Retained so panel apps always know current state on reconnect
  if (mqtt.connected()) {
    mqtt.publish(TOPIC_HUMIDIFIER_STATE,
      humidifierOn ? "true" : "false", true);
  }

  Serial.printf(">> Humidifier %s — %s\n", state ? "ON" : "OFF", reason);
}

// ── Publish telemetry JSON + flat topics for panel apps ────────
void publishTelemetry(float tempC, float humidity) {

  // Full JSON payload — for Node-RED, Home Assistant etc
  // Schema: temperature(float), humidity(float), humidifier(bool), mode(string)
  char jsonPayload[128];
  snprintf(jsonPayload, sizeof(jsonPayload),
    "{\"temperature\":%.2f,\"humidity\":%.2f,\"humidifier\":%s,\"mode\":\"%s\"}",
    tempC,
    humidity,
    humidifierOn ? "true" : "false",
    mode == AUTO ? "AUTO" : mode == FORCED_ON ? "FORCED_ON" : "FORCED_OFF"
  );
  mqtt.publish(TOPIC_TELEMETRY, jsonPayload);

  // Flat individual topics — for MQTT panel apps
  char buf[16];

  snprintf(buf, sizeof(buf), "%.2f", tempC);
  mqtt.publish(TOPIC_TEMP, buf);                  // e.g. "27.30"

  snprintf(buf, sizeof(buf), "%.2f", humidity);
  mqtt.publish(TOPIC_HUMIDITY, buf);              // e.g. "65.20"

  // Humidifier state published retained — panel apps see correct
  // state immediately on reconnect without waiting for next publish
  mqtt.publish(TOPIC_HUMIDIFIER_STATE,
    humidifierOn ? "true" : "false", true);

  // Track last published values for delta comparison
  lastPublishedTemp     = tempC;
  lastPublishedHumidity = humidity;
  lastPublishedHumState = humidifierOn;
  lastTelemetryPublish  = millis();

  Serial.printf(">> Telemetry published: %s\n", jsonPayload);
}

// ── Publish system diagnostics JSON ───────────────────────────
void publishSystem() {
  char payload[128];
  snprintf(payload, sizeof(payload),
    "{\"rssi\":%d,\"heap\":%lu,\"uptime\":%lu}",
    WiFi.RSSI(),
    (unsigned long)ESP.getFreeHeap(),
    (millis() - bootTime) / 1000
  );
  mqtt.publish(TOPIC_SYSTEM, payload);
  Serial.printf(">> System published: %s\n", payload);
  lastSystemPublish = millis();
}

// ── Publish fault event to dedicated fault topic ───────────────
// Fault topic is separate from telemetry to keep telemetry
// schema stable — consumers parsing telemetry JSON won't see
// unexpected fault payloads breaking their parsers
void publishFault(const char* message) {
  char payload[128];
  snprintf(payload, sizeof(payload),
    "{\"fault\":\"%s\",\"uptime\":%lu}",
    message,
    (millis() - bootTime) / 1000
  );
  mqtt.publish(TOPIC_FAULT, payload);
  Serial.printf(">> FAULT published: %s\n", payload);
}

// ── MQTT command handler ───────────────────────────────────────
void handleCommand(const char* payload) {
  // StaticJsonDocument sized for max expected command payload
  // Max valid command fits well within 128 bytes
  // Oversized payloads already rejected in onMQTTMessage()
  StaticJsonDocument<128> doc;
  DeserializationError err = deserializeJson(doc, payload);

  if (err) {
    Serial.printf(">> Command parse error: %s\n", err.c_str());
    return;
  }

  // Override command
  if (doc.containsKey("override")) {
    const char* val = doc["override"];
    if (strcmp(val, "ON") == 0) {
      mode = FORCED_ON;
      setHumidifier(true, "manual override ON");
    } else if (strcmp(val, "OFF") == 0) {
      mode = FORCED_OFF;
      setHumidifier(false, "manual override OFF");
    } else if (strcmp(val, "AUTO") == 0) {
      mode = AUTO;
      Serial.println(">> Mode set to AUTO");
    } else {
      Serial.println(">> Unknown override value — ignored");
    }
  }

  // Threshold update — validate before applying
  float newOn  = doc.containsKey("threshold_on")  ? (float)doc["threshold_on"]  : humidityOnThreshold;
  float newOff = doc.containsKey("threshold_off") ? (float)doc["threshold_off"] : humidityOffThreshold;

  if (doc.containsKey("threshold_on") || doc.containsKey("threshold_off")) {
    if (newOn  < HUMIDITY_MIN_VALID || newOn  > HUMIDITY_MAX_VALID ||
        newOff < HUMIDITY_MIN_VALID || newOff > HUMIDITY_MAX_VALID ||
        newOn  >= newOff) {
      Serial.printf(">> Threshold rejected — invalid (on:%.1f off:%.1f)\n", newOn, newOff);
      Serial.printf("   Rules: %.0f-%.0f%%, threshold_on < threshold_off\n",
        HUMIDITY_MIN_VALID, HUMIDITY_MAX_VALID);
    } else {
      humidityOnThreshold  = newOn;
      humidityOffThreshold = newOff;
      Serial.printf(">> Thresholds updated — ON:%.1f%% OFF:%.1f%%\n",
        humidityOnThreshold, humidityOffThreshold);
    }
  }
}

// ── MQTT message callback ──────────────────────────────────────
void onMQTTMessage(char* topic, byte* payload, unsigned int length) {

  // Reject oversized payloads before touching the stack
  // Prevents stack overflow from malformed or malicious broker messages
  if (length > MQTT_MAX_PAYLOAD) {
    Serial.printf(">> Command rejected — payload too large (%u bytes, max %d)\n",
      length, MQTT_MAX_PAYLOAD);
    return;
  }

  // Fixed size buffer — safe regardless of broker payload
  char message[MQTT_MAX_PAYLOAD + 1];
  unsigned int copyLen = min(length, (unsigned int)MQTT_MAX_PAYLOAD);
  memcpy(message, payload, copyLen);
  message[copyLen] = '\0';

  Serial.printf(">> Command received [%s]: %s\n", topic, message);
  handleCommand(message);
}

// ── Non-blocking reconnect with exponential backoff ───────────
void maintainConnections() {
  if (millis() - lastReconnectAttempt < backoffDelay) return;
  lastReconnectAttempt = millis();

  // WiFi check
  if (WiFi.status() != WL_CONNECTED) {
    Serial.printf(">> WiFi not connected — reconnecting (backoff: %lus)\n",
      backoffDelay / 1000);
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
    backoffDelay = min(backoffDelay * 2, (unsigned long)BACKOFF_MAX);
    return;
  }

  // MQTT check
  if (!mqtt.connected()) {
    Serial.printf(">> MQTT not connected — reconnecting (backoff: %lus)\n",
      backoffDelay / 1000);
    char clientId[32];
    getMQTTClientId(clientId, sizeof(clientId));
    if (mqtt.connect(clientId, TOPIC_STATUS, 1, true, "offline")) {
      mqtt.publish(TOPIC_STATUS, "online", true); // Retained online status
      mqtt.subscribe(TOPIC_COMMAND);
      Serial.println(">> MQTT reconnected — subscribed to command topic");
      backoffDelay = BACKOFF_MIN; // Reset backoff on success
    } else {
      Serial.printf(">> MQTT failed RC=%d — will retry\n", mqtt.state());
      backoffDelay = min(backoffDelay * 2, (unsigned long)BACKOFF_MAX);
    }
    return;
  }

  // Both connected — reset backoff
  backoffDelay = BACKOFF_MIN;
}

// ── Humidifier auto control logic ─────────────────────────────
void controlHumidifier(float humidity) {
  // Only runs in AUTO mode — overrides bypass this entirely
  if (mode != AUTO) return;

  if (!humidifierOn) {
    if (humidity < humidityOnThreshold) {
      if (!lowHumTimerActive) {
        lowHumStartTime   = millis();
        lowHumTimerActive = true;
        Serial.printf(">> Humidity %.1f%% below %.1f%% — starting 10s timer\n",
          humidity, humidityOnThreshold);
      } else {
        unsigned long elapsed = millis() - lowHumStartTime;
        if (elapsed >= TRIGGER_DELAY) {
          setHumidifier(true, "humidity below threshold for 10s");
          lowHumTimerActive = false;
        } else {
          Serial.printf(">> Timer running: %lus / 10s\n", elapsed / 1000);
        }
      }
    } else {
      if (lowHumTimerActive) {
        lowHumTimerActive = false;
        Serial.println(">> Humidity recovered — 10s timer reset");
      }
    }
  } else {
    if (humidity >= humidityOffThreshold) {
      setHumidifier(false, "humidity reached off threshold");
    }
  }
}

// ── Sensor failsafe check ──────────────────────────────────────
void checkSensorFailsafe() {
  if (millis() - lastValidReadTime <= SENSOR_TIMEOUT) return;

  // Sensor has been failing for >30s
  if (!sensorFaultActive) {
    sensorFaultActive = true;

    if (mode == AUTO) {
      // AUTO mode — failsafe wins, force OFF
      setHumidifier(false, "sensor fault failsafe");
      Serial.println(">> SENSOR FAULT — humidifier forced OFF (AUTO mode)");
      if (mqtt.connected()) publishFault("sensor_timeout_auto_off");
    } else if (mode == FORCED_ON) {
      // Manual override active — respect human decision but warn loudly
      Serial.println(">> SENSOR FAULT — humidifier stays ON (FORCED_ON active)");
      Serial.println("   WARNING: Operating blind — consider switching to AUTO");
      if (mqtt.connected()) publishFault("sensor_timeout_override_active");
    } else {
      // FORCED_OFF — already off, just log
      Serial.println(">> SENSOR FAULT — humidifier already OFF (FORCED_OFF mode)");
      if (mqtt.connected()) publishFault("sensor_timeout_already_off");
    }
  }
}

// ── Sensor recovery handler ────────────────────────────────────
void handleSensorRecovery() {
  if (!sensorFaultActive) return;
  sensorFaultActive = false;
  lowHumTimerActive = false; // Reset timer — start fresh after blind window
  Serial.println(">> Sensor recovered — resuming normal operation");
  if (mqtt.connected()) publishFault("sensor_recovered");
}

void setup() {
  Serial.begin(BAUD_RATE);
  delay(1000);
  bootTime = millis();

  // ── Watchdog ────────────────────────────────────────────────
  // Arduino core already initializes WDT — just register our task
  esp_task_wdt_add(NULL);
  Serial.println("Watchdog registered");

  // ── Humidifier pin ──────────────────────────────────────────
  // Always OFF on boot — sensor decides state fresh
  pinMode(HUMIDIFIER_PIN, OUTPUT);
  digitalWrite(HUMIDIFIER_PIN, RELAY_OFF);
  Serial.println("Humidifier OFF on boot");

  dht.begin();
  lastValidReadTime = millis(); // Give sensor grace period on boot

  // ── WiFi with 15s timeout ───────────────────────────────────
  Serial.printf("Connecting to WiFi: %s\n", WIFI_SSID);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
  unsigned long wifiStart = millis();
  while (WiFi.status() != WL_CONNECTED) {
    esp_task_wdt_reset();
    if (millis() - wifiStart > WIFI_TIMEOUT) {
      Serial.println("\nWiFi timeout — continuing. Humidifier runs regardless.");
      break;
    }
    delay(500);
    Serial.print(".");
  }
  if (WiFi.status() == WL_CONNECTED) {
    Serial.printf("\nWiFi connected — IP: %s\n", WiFi.localIP().toString().c_str());
  }

  // ── MQTT with 15s timeout ───────────────────────────────────
  if (WiFi.status() == WL_CONNECTED) {
    mqtt.setServer(MQTT_BROKER, MQTT_PORT);
    mqtt.setCallback(onMQTTMessage);

    char clientId[32];
    getMQTTClientId(clientId, sizeof(clientId));
    Serial.printf("Connecting to MQTT as: %s\n", clientId);

    unsigned long mqttStart = millis();
    while (!mqtt.connected()) {
      esp_task_wdt_reset();
      if (millis() - mqttStart > MQTT_TIMEOUT) {
        Serial.println("MQTT timeout — continuing. Humidifier runs regardless.");
        break;
      }
      if (mqtt.connect(clientId, TOPIC_STATUS, 1, true, "offline")) {
        mqtt.publish(TOPIC_STATUS, "online", true);
        mqtt.subscribe(TOPIC_COMMAND);
        Serial.println("MQTT connected — subscribed to command topic");
        backoffDelay = BACKOFF_MIN;
      }
      delay(1000);
      Serial.print(".");
    }
  }

  Serial.println("--------------------------------------------------");
  Serial.println("Version 2.1 Ready — ESP32 Humidifier Controller");
  Serial.printf("Thresholds — ON:%.1f%%  OFF:%.1f%%\n",
    humidityOnThreshold, humidityOffThreshold);
  Serial.println("Mode: AUTO");
  Serial.println("--------------------------------------------------");
}

void loop() {
  esp_task_wdt_reset(); // Confirm loop is alive

  // Maintain connections — non-blocking, exponential backoff
  maintainConnections();
  if (mqtt.connected()) mqtt.loop();

  // Always check sensor failsafe — runs regardless of network state
  checkSensorFailsafe();

  // ── Sensor read ─────────────────────────────────────────────
  if (millis() - lastReadTime < READ_INTERVAL) return;
  lastReadTime = millis();

  float humidity = dht.readHumidity();
  float tempC    = dht.readTemperature();

  if (isnan(humidity) || isnan(tempC)) {
    Serial.println("WARNING: DHT22 read failed");
    return;
  }

  // Valid reading received
  lastValidReadTime = millis();
  handleSensorRecovery();

  // Print to serial
  Serial.printf("Temp: %.2f°C  |  Humidity: %.2f%%  |  Humidifier: %s  |  Mode: %s\n",
    tempC, humidity,
    humidifierOn ? "ON" : "OFF",
    mode == AUTO ? "AUTO" : mode == FORCED_ON ? "FORCED_ON" : "FORCED_OFF"
  );

  // Run humidifier logic — AUTO mode only
  controlHumidifier(humidity);

  // ── Smart publish — delta or heartbeat ──────────────────────
  if (mqtt.connected()) {
    bool tempChanged      = abs(tempC    - lastPublishedTemp)     >= PUBLISH_DELTA;
    bool humChanged       = abs(humidity - lastPublishedHumidity) >= PUBLISH_DELTA;
    bool stateChanged     = humidifierOn != lastPublishedHumState;
    bool heartbeatElapsed = (millis() - lastTelemetryPublish)     >= HEARTBEAT_MAX;

    if (tempChanged || humChanged || stateChanged || heartbeatElapsed) {
      publishTelemetry(tempC, humidity);
    }

    // System metrics every 60s
    if (millis() - lastSystemPublish >= SYSTEM_INTERVAL) {
      publishSystem();
    }
  }
}