// ====================================================================
// Task 3 – Pico W  |  MQTT + JSON
// Single file – no nano.ino needed, no compilation order issues.
//
// How it works:
//   Serial Monitor → type command → published to broker as source="nano"
//   → travels through fbe-mqtt.hs-weingarten.de → received back by Pico
//   → validated → LED controlled
// ====================================================================

#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 = "fbe-mqtt.hs-weingarten.de";
const int   MQTT_PORT   = 1883;
const char* TOPIC_CMD   = "iem/task3/pico/cmd";   // subscribe + publish
const char* TOPIC_DATA  = "iem/task3/pico/data";  // publish sensor data

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

// ── Pins ─────────────────────────────────────────────────────────────
const int LED_PIN       = 28;
const int HEARTBEAT_PIN = LED_BUILTIN;
const int BUTTON_PIN    = 2;
const int POT_PIN       = 26;

// ── LED / blink state ────────────────────────────────────────────────
bool          blinkEnabled     = true;
bool          ledState         = false;
int           overrideInterval = -1;
unsigned long lastBlink        = 0;

// ── Watchdog ─────────────────────────────────────────────────────────
bool          remoteMode   = false;
bool          inSafeState  = false;
unsigned long lastValidCmd = 0;

// ── Timing ───────────────────────────────────────────────────────────
unsigned long lastHeartbeat  = 0;
bool          heartbeatState = false;
unsigned long lastPublish    = 0;

// ── Sequence numbers ─────────────────────────────────────────────────
int           rxExpectedSeq = -1;   // incoming commands
unsigned long txDataSeq     = 0;    // outgoing sensor data
unsigned long txCmdSeq      = 0;    // outgoing commands (from Serial)

// ── Statistics ───────────────────────────────────────────────────────
unsigned long msgAccepted = 0;
unsigned long msgRejJson  = 0;
unsigned long msgRejAuth  = 0;
unsigned long msgSeqGaps  = 0;

// ── Objects ──────────────────────────────────────────────────────────
WiFiClient   wifiClient;
PubSubClient mqtt(wifiClient);
SimpleJson   jsonIn, jsonOut;
String       inputBuffer = "";


// ====================================================================
// WiFi
// ====================================================================
void setupWiFi() {
  Serial.println("[WiFi] Connecting...");
  WiFi.begin(WIFI_SSID, WIFI_PASS);
  int tries = 0;
  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
    if (++tries > 40) {
      Serial.println("\n[WiFi] FAILED – is Wokwi IoT Gateway running?");
      return;
    }
  }
  Serial.print("\n[WiFi] Connected – IP: ");
  Serial.println(WiFi.localIP());
}


// ====================================================================
// Watchdog
// ====================================================================
void enterSafeState() {
  if (inSafeState) return;
  inSafeState      = true;
  overrideInterval = -1;
  blinkEnabled     = true;
  Serial.println("[WDG] 10s timeout → SAFE STATE (local poti)");
}

void leaveSafeState() {
  if (!inSafeState) return;
  inSafeState = false;
  Serial.println("[WDG] Valid message → leaving safe state");
}


// ====================================================================
// Validation  (getString() returns arduino::String → need .c_str())
// ====================================================================
bool validateMessage(const SimpleJson& msg) {
  if (!msg.hasKey("token") ||
      strcmp(msg.getString("token").c_str(), SHARED_TOKEN) != 0) {
    Serial.println("[AUTH] Rejected: bad/missing token");
    msgRejAuth++;
    return false;
  }
  if (!msg.hasKey("source") ||
      strcmp(msg.getString("source").c_str(), EXPECTED_SOURCE) != 0) {
    Serial.print("[AUTH] Rejected: unexpected source = ");
    Serial.println(msg.getString("source"));
    msgRejAuth++;
    return false;
  }
  return true;
}


// ====================================================================
// Sequence check
// ====================================================================
bool checkSequence(const SimpleJson& msg) {
  if (!msg.hasKey("seq")) return true;
  int seq = msg.getInt("seq");

  if (rxExpectedSeq == -1 || seq == 0) {
    rxExpectedSeq = seq + 1;
    return true;
  }
  if (seq < rxExpectedSeq) return false;   // replay – discard
  if (seq > rxExpectedSeq) {
    Serial.print("[SEQ] Gap: expected "); Serial.print(rxExpectedSeq);
    Serial.print(" got "); Serial.println(seq);
    msgSeqGaps++;
  }
  rxExpectedSeq = seq + 1;
  return true;
}


// ====================================================================
// Process incoming command fields
// ====================================================================
void processCommand(const SimpleJson& cmd) {
  if (cmd.hasKey("blinkEnabled")) {
    blinkEnabled = cmd.getBool("blinkEnabled");
    Serial.print("[CMD] blinkEnabled=");
    Serial.println(blinkEnabled ? "true" : "false");
  }
  if (cmd.hasKey("ledOn")) {
    if (cmd.getBool("ledOn")) {
      blinkEnabled = false; ledState = true;
      digitalWrite(LED_PIN, HIGH);
      Serial.println("[CMD] LED forced ON");
    } else {
      ledState = false;
      digitalWrite(LED_PIN, LOW);
      Serial.println("[CMD] LED forced OFF");
    }
  }
  if (cmd.hasKey("interval")) {
    int iv = cmd.getInt("interval");
    if (iv < INTERVAL_MIN || iv > INTERVAL_MAX) {
      Serial.print("[CMD] Interval rejected (out of range): ");
      Serial.println(iv);
    } else {
      overrideInterval = iv;
      Serial.print("[CMD] interval="); Serial.print(iv); Serial.println("ms");
    }
  }
  if (cmd.hasKey("useLocalPot") && cmd.getBool("useLocalPot")) {
    overrideInterval = -1;
    Serial.println("[CMD] useLocalPot → poti restored");
  }
}


// ====================================================================
// MQTT callback – fires when a message arrives on TOPIC_CMD
// ====================================================================
void mqttCallback(char* topic, byte* payload, unsigned int length) {
  char msg[256];
  if (length >= sizeof(msg)) length = sizeof(msg) - 1;
  memcpy(msg, payload, length);
  msg[length] = '\0';                       // mandatory null-terminate

  Serial.print("[MQTT IN] "); Serial.println(msg);

  if (!jsonIn.parse(msg))       { msgRejJson++; return; }
  if (!validateMessage(jsonIn)) return;
  if (!checkSequence(jsonIn))   return;

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


// ====================================================================
// MQTT reconnect + subscribe
// ====================================================================
void mqttReconnect() {
  String id = "iem-pico-";
  id += String(random(0xFFFF), HEX);
  Serial.print("[MQTT] Connecting as "); Serial.print(id); Serial.print("... ");
  if (mqtt.connect(id.c_str())) {
    Serial.println("OK");
    mqtt.subscribe(TOPIC_CMD);
  } else {
    Serial.print("FAIL rc="); Serial.println(mqtt.state());
    delay(2000);
  }
}


// ====================================================================
// Publish sensor data → TOPIC_DATA  (Pico → broker)
// ====================================================================
void publishSensorData(int potRaw, unsigned long activeInterval) {
  jsonOut.clear();
  jsonOut.setString("token",        SHARED_TOKEN);
  jsonOut.setString("source",       "pico");
  jsonOut.setInt   ("seq",          (int)txDataSeq++);
  jsonOut.setInt   ("potValue",     potRaw);
  jsonOut.setBool  ("blinkEnabled", blinkEnabled);
  jsonOut.setInt   ("interval",     (int)activeInterval);
  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);
  Serial.print("[DATA OUT] "); Serial.println(buf);
}


// ====================================================================
// Publish command → TOPIC_CMD  (simulates Nano sending to broker)
// source="nano" so the Pico's own validateMessage() accepts it back
// ====================================================================
void publishCommand() {
  jsonOut.setString("token",  SHARED_TOKEN);
  jsonOut.setString("source", "nano");
  jsonOut.setInt   ("seq",    (int)txCmdSeq++);

  char buf[256];
  jsonOut.toCharArray(buf, sizeof(buf));
  mqtt.publish(TOPIC_CMD, buf);
  Serial.print("[CMD OUT] "); Serial.println(buf);
}


// ====================================================================
// Serial Monitor command parser  (simulates Nano side)
// ====================================================================
void processSerialInput() {
  while (Serial.available()) {
    char c = (char)Serial.read();

    if (c == '\n' || c == '\r') {
      if (inputBuffer.length() == 0) return;

      String input = inputBuffer;
      inputBuffer = "";
      input.trim();
      input.toUpperCase();

      if (input == "ON") {
        jsonOut.clear();
        jsonOut.setBool("blinkEnabled", false);
        jsonOut.setBool("ledOn",        true);
        publishCommand();
        Serial.println("ACK:ON");

      } else if (input == "OFF") {
        jsonOut.clear();
        jsonOut.setBool("blinkEnabled", false);
        jsonOut.setBool("ledOn",        false);
        publishCommand();
        Serial.println("ACK:OFF");

      } else if (input == "BLINK") {
        jsonOut.clear();
        jsonOut.setBool("blinkEnabled", true);
        publishCommand();
        Serial.println("ACK:BLINK");

      } else if (input == "NOBLINK") {
        jsonOut.clear();
        jsonOut.setBool("blinkEnabled", false);
        publishCommand();
        Serial.println("ACK:NOBLINK");

      } else if (input.startsWith("INTERVAL")) {
        int spIdx = input.indexOf(' ');
        if (spIdx == -1) {
          Serial.println("Usage: INTERVAL <ms>  e.g. INTERVAL 200");
        } else {
          int ms = input.substring(spIdx + 1).toInt();
          if (ms < INTERVAL_MIN || ms > INTERVAL_MAX) {
            Serial.print("Error: must be ");
            Serial.print(INTERVAL_MIN); Serial.print("–");
            Serial.print(INTERVAL_MAX); Serial.println(" ms");
          } else {
            jsonOut.clear();
            jsonOut.setInt("interval", ms);
            publishCommand();
            Serial.print("ACK:INTERVAL:"); Serial.println(ms);
          }
        }

      } else if (input == "POT") {
        jsonOut.clear();
        jsonOut.setBool("useLocalPot", true);
        publishCommand();
        Serial.println("ACK:POT");

      } else if (input == "STATUS") {
        int pot = analogRead(POT_PIN);
        int iv  = (overrideInterval > 0)
                    ? overrideInterval
                    : (int)map(pot, 0, 1023, INTERVAL_MIN, 1000);
        Serial.println("── Pico W status ──────────────────");
        Serial.print("   LED      : "); Serial.println(ledState     ? "ON"  : "OFF");
        Serial.print("   Blink    : "); Serial.println(blinkEnabled ? "ON"  : "OFF");
        Serial.print("   Interval : "); Serial.print(iv); Serial.println(" ms");
        Serial.print("   Pot raw  : "); Serial.println(pot);
        Serial.print("   SafeState: "); Serial.println(inSafeState  ? "YES" : "NO");
        Serial.print("   Uptime   : "); Serial.print(millis()/1000); Serial.println(" s");
        Serial.println("───────────────────────────────────");

      } else if (input == "STATS") {
        Serial.println("── MQTT statistics ────────────────");
        Serial.print("   Accepted  : "); Serial.println(msgAccepted);
        Serial.print("   Rej JSON  : "); Serial.println(msgRejJson);
        Serial.print("   Rej Auth  : "); Serial.println(msgRejAuth);
        Serial.print("   Seq gaps  : "); Serial.println(msgSeqGaps);
        Serial.println("───────────────────────────────────");

      } else if (input == "HELP") {
        Serial.println("── Commands ───────────────────────");
        Serial.println("   ON               LED on");
        Serial.println("   OFF              LED off");
        Serial.println("   BLINK            Start blinking");
        Serial.println("   NOBLINK          Stop blinking");
        Serial.println("   INTERVAL <ms>    Set interval (50–2000)");
        Serial.println("   POT              Use local potentiometer");
        Serial.println("   STATUS           Show current state");
        Serial.println("   STATS            Show message statistics");
        Serial.println("   HELP             Show this list");
        Serial.println("───────────────────────────────────");

      } else {
        Serial.print("Unknown: "); Serial.print(input);
        Serial.println("  (type HELP)");
      }

    } else {
      inputBuffer += c;
    }
  }
}


// ====================================================================
// Setup
// ====================================================================
void setup() {
  pinMode(LED_PIN,       OUTPUT);
  pinMode(HEARTBEAT_PIN, OUTPUT);
  pinMode(BUTTON_PIN,    INPUT_PULLUP);

  Serial.begin(115200);
  delay(300);
  Serial.println("\n=== Task 3 – Pico W MQTT ===");
  Serial.println("Type HELP for commands.\n");

  setupWiFi();
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(mqttCallback);
}


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

  // Heartbeat ~300 ms
  if (now - lastHeartbeat >= 300) {
    lastHeartbeat  = now;
    heartbeatState = !heartbeatState;
    digitalWrite(HEARTBEAT_PIN, heartbeatState);
  }

  // MQTT – must call mqtt.loop() every iteration
  if (!mqtt.connected()) mqttReconnect();
  mqtt.loop();

  // Watchdog
  if (remoteMode && (now - lastValidCmd >= WATCHDOG_MS)) enterSafeState();

  // Button – debounced edge detection
  static bool          lastBtn   = HIGH;
  static bool          stableBtn = HIGH;
  static unsigned long debounceT = 0;
  bool curBtn = digitalRead(BUTTON_PIN);
  if (curBtn != lastBtn) debounceT = now;
  if ((now - debounceT) > 50 && curBtn != stableBtn) {
    stableBtn = curBtn;
    if (stableBtn == LOW) {
      blinkEnabled = !blinkEnabled;
      Serial.print("[BTN] Blink: "); Serial.println(blinkEnabled ? "ON" : "OFF");
    }
  }
  lastBtn = curBtn;

  // Potentiometer
  int potRaw = analogRead(POT_PIN);
  unsigned long potInterval    = map(potRaw, 0, 1023, INTERVAL_MIN, 1000);
  unsigned long activeInterval = (overrideInterval > 0)
                                   ? (unsigned long)overrideInterval
                                   : potInterval;

  // Blink (non-blocking)
  if (blinkEnabled && (now - lastBlink >= activeInterval)) {
    lastBlink = now;
    ledState  = !ledState;
    digitalWrite(LED_PIN, ledState);
  }

  // Publish sensor data every 500 ms
  if (now - lastPublish >= 500) {
    lastPublish = now;
    publishSensorData(potRaw, activeInterval);
  }

  // Serial input
  processSerialInput();
}