#include <WiFi.h>
#include <PubSubClient.h>
#include "SimpleJson.h"

// -- WiFi ------------------------------------------------------------
const char* WIFI_SSID = "Wokwi-GUEST";
const char* WIFI_PASS = "";

// -- MQTT ------------------------------------------------------------
const char* MQTT_BROKER = "broker.hivemq.com";
const int   MQTT_PORT   = 1883;
const char* TOPIC_DATA  = "iem/task3/arushi/data";
const char* TOPIC_CMD   = "iem/task3/arushi/cmd";

// -- Robustness ------------------------------------------------------
const char*          SHARED_TOKEN     = "iem2026";
const char*          EXPECTED_SOURCE  = "nano";
const unsigned long  WATCHDOG_TIMEOUT = 10000;
const int            INTERVAL_MIN     = 50;
const int            INTERVAL_MAX     = 2000;

// -- Hardware --------------------------------------------------------
const int LED_PIN    = 28;
const int BUTTON_PIN = 2;
const int POT_PIN    = 26;

// -- State -----------------------------------------------------------
bool          blinkEnabled     = false;
bool          ledState         = false;
int           overrideInterval = -1;
unsigned long lastValidCmd     = 0;
bool          inSafeState      = false;
int           expectedSeqNr    = -1;
unsigned long seqOutgoing      = 0;

// -- Button ----------------------------------------------------------
bool          lastButtonState  = HIGH;
unsigned long lastPressTime    = 0;

// -- Timing ----------------------------------------------------------
unsigned long lastBlinkTime    = 0;
unsigned long lastPublish      = 0;
unsigned long lastHeartbeat    = 0;

// -- Statistics ------------------------------------------------------
unsigned long msgAccepted      = 0;
unsigned long msgRejectedJson  = 0;
unsigned long msgRejectedAuth  = 0;
unsigned long msgSeqGaps       = 0;

WiFiClient   wifiClient;
PubSubClient mqtt(wifiClient);
SimpleJson   jsonOut, jsonIn;

// ── WiFi setup ──────────────────────────────────────────────────────
void setupWiFi() {
  Serial1.print("Connecting to WiFi");
  WiFi.begin(WIFI_SSID, WIFI_PASS);
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 30) {
    delay(500);
    Serial1.print(".");
    attempts++;
  }
  if (WiFi.status() == WL_CONNECTED) {
    Serial1.println("\nWiFi connected!");
  } else {
    Serial1.println("\nWiFi FAILED");
  }
}

// ── Safe state ──────────────────────────────────────────────────────
// Activates when no valid command received for WATCHDOG_TIMEOUT ms
void enterSafeState() {
  if (!inSafeState) {
    inSafeState      = true;
    overrideInterval = -1;
    Serial1.println("WATCHDOG: Safe state – poti control restored");
  }
}

// ── Leave safe state ─────────────────────────────────────────────────
void leaveSafeState() {
  if (inSafeState) {
    inSafeState = false;
    Serial1.println("WATCHDOG: Remote control active");
  }
}

// ── Message validation ───────────────────────────────────────────────
// Rejects messages with wrong token or wrong source
bool validateMessage(const SimpleJson& msg) {
  if (strcmp(msg.getString("token"), SHARED_TOKEN) != 0) {
    Serial1.println("REJECT: bad token");
    msgRejectedAuth++;
    return false;
  }
  if (strcmp(msg.getString("source"), EXPECTED_SOURCE) != 0) {
    Serial1.println("REJECT: bad source");
    msgRejectedAuth++;
    return false;
  }
  return true;
}

// ── Sequence number check ────────────────────────────────────────────
// Detects gaps, replays, and restarts
bool checkSequence(const SimpleJson& msg) {
  int seq = msg.getInt("seq", -1);
  if (seq < 0) return true;
  if (seq == 0) {
    expectedSeqNr = 1;
    Serial1.println("SEQ: sender restarted, resyncing");
    return true;
  }
  if (expectedSeqNr < 0) {
    expectedSeqNr = seq + 1;
    return true;
  }
  if (seq < expectedSeqNr) {
    Serial1.println("SEQ: replay discarded");
    return false;
  }
  if (seq > expectedSeqNr) {
    Serial1.println("SEQ: gap detected");
    msgSeqGaps++;
  }
  expectedSeqNr = seq + 1;
  return true;
}

// ── Process command fields ───────────────────────────────────────────
void processCommand(const SimpleJson& cmd) {
  // blinkEnabled
  if (cmd.hasKey("blinkEnabled")) {
    blinkEnabled = cmd.getBool("blinkEnabled");
    if (!blinkEnabled) digitalWrite(LED_PIN, LOW);
    Serial1.print("CMD: blinkEnabled=");
    Serial1.println(blinkEnabled);
  }
  // ledOn – permanently on/off
  if (cmd.hasKey("ledOn")) {
    bool on = cmd.getBool("ledOn");
    blinkEnabled = false;
    ledState     = on;
    digitalWrite(LED_PIN, on ? HIGH : LOW);
    Serial1.print("CMD: ledOn="); Serial1.println(on);
  }
  // interval override with range check
  if (cmd.hasKey("interval")) {
    int val = cmd.getInt("interval");
    if (val >= INTERVAL_MIN && val <= INTERVAL_MAX) {
      overrideInterval = val;
      Serial1.print("CMD: interval="); Serial1.println(val);
    } else {
      Serial1.println("CMD: interval out of range");
    }
  }
  // useLocalPot – reset interval override
  if (cmd.hasKey("useLocalPot") && cmd.getBool("useLocalPot")) {
    overrideInterval = -1;
    Serial1.println("CMD: poti control restored");
  }
}

// ── MQTT callback ────────────────────────────────────────────────────
void mqttCallback(char* topic, byte* payload, unsigned int length) {
  char buf[256];
  unsigned int len = min(length, (unsigned int)255);
  memcpy(buf, payload, len);
  buf[len] = 0;
  Serial1.print("RX: "); Serial1.println(buf);

  jsonIn.parse(buf);
  if (jsonIn.count == 0) {
    Serial1.println("REJECT: invalid JSON");
    msgRejectedJson++;
    return;
  }
  if (!validateMessage(jsonIn)) return;
  if (!checkSequence(jsonIn)) return;

  lastValidCmd = millis();
  leaveSafeState();
  msgAccepted++;
  processCommand(jsonIn);
}

// ── MQTT reconnect ───────────────────────────────────────────────────
void mqttReconnect() {
  while (!mqtt.connected()) {
    Serial1.print("Connecting to MQTT...");
    String clientId = "pico-" + String(random(0xffff), HEX);
    if (mqtt.connect(clientId.c_str())) {
      Serial1.println("connected");
      mqtt.subscribe(TOPIC_CMD);
      Serial1.print("Subscribed to: "); Serial1.println(TOPIC_CMD);
    } else {
      Serial1.print("failed rc=");
      Serial1.print(mqtt.state());
      Serial1.println(" retrying in 3s");
      delay(3000);
    }
  }
}

// ── Publish sensor data ──────────────────────────────────────────────
// Sends Pico state to ESP32 every 2000ms
void publishSensorData(int potValue, unsigned long blinkInterval) {
  jsonOut.clear();
  jsonOut.setString("token",        SHARED_TOKEN);
  jsonOut.setString("source",       "pico");
  jsonOut.setInt   ("seq",          (int)seqOutgoing++);
  jsonOut.setInt   ("potValue",     potValue);
  jsonOut.setBool  ("blinkEnabled", blinkEnabled);
  jsonOut.setInt   ("interval",     (int)blinkInterval);
  jsonOut.setBool  ("ledState",     ledState);
  jsonOut.setBool  ("safeState",    inSafeState);
  jsonOut.setInt   ("uptime",       (int)(millis()/1000));

  char buf[256];
  jsonOut.toCharArray(buf, sizeof(buf));
  mqtt.publish(TOPIC_DATA, buf);
}

// ── setup ────────────────────────────────────────────────────────────
void setup() {
  pinMode(LED_PIN,    OUTPUT);
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  Serial1.begin(115200);
  delay(1000);
  Serial1.println("=== Pico W MQTT Task 3 ===");
  setupWiFi();
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(mqttCallback);
  lastValidCmd = millis();
  Serial1.println("Pico W ready");
}

// ── loop ─────────────────────────────────────────────────────────────
void loop() {
  unsigned long now = millis();

  // MQTT connection + processing
  if (!mqtt.connected()) mqttReconnect();
  mqtt.loop();

  // Watchdog – enter safe state if no command for 10s
  if (!inSafeState && (now - lastValidCmd > WATCHDOG_TIMEOUT))
    enterSafeState();

  // Heartbeat – onboard LED every 300ms
  if (now - lastHeartbeat >= 300) {
    lastHeartbeat = now;
    digitalWrite(LED_BUILTIN, !digitalRead(LED_BUILTIN));
  }

  // Button – edge detection with time guard
  bool currentButton = digitalRead(BUTTON_PIN);
  if (lastButtonState == HIGH && currentButton == LOW
      && (now - lastPressTime) > 300) {
    lastPressTime = now;
    blinkEnabled  = !blinkEnabled;
    if (!blinkEnabled) digitalWrite(LED_PIN, LOW);
    Serial1.println(blinkEnabled ? "BTN: blink ON" : "BTN: blink OFF");
  }
  lastButtonState = currentButton;

  // Potentiometer + interval determination
  int  potValue      = analogRead(POT_PIN);
  long blinkInterval = (overrideInterval > 0)
                         ? overrideInterval
                         : map(potValue, 0, 4095, 50, 1000);

  // Non-blocking blink
  if (blinkEnabled && (now - lastBlinkTime >= (unsigned long)blinkInterval)) {
    lastBlinkTime = now;
    ledState      = !ledState;
    digitalWrite(LED_PIN, ledState ? HIGH : LOW);
  }

  // Publish sensor data every 2000ms
  if (now - lastPublish >= 2000) {
    lastPublish = now;
    publishSensorData(potValue, blinkInterval);
  }
}