// ====================================================================
// Task 3: MQTT + JSON  –  Pico W  (TAB 1 – the "device")
// Board  : Raspberry Pi Pico W
// Wiring : LED on GP28 | Button on GP2 | Potentiometer on GP26
//
// This board:
//   • Connects to WiFi + MQTT broker
//   • Subscribes to TOPIC_CMD  → executes LED commands
//   • Publishes  to TOPIC_DATA → sends poti/LED state every 500 ms
// ====================================================================

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

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

// ── MQTT ─────────────────────────────────────────────────────────────
// If test.mosquitto.org is unreliable, try: broker.hivemq.com
const char* MQTT_BROKER = "test.mosquitto.org";
const int   MQTT_PORT   = 1883;

// ⚠ Change "student01" to your own name so you don't clash with others!
const char* TOPIC_DATA = "iem/task3/student01/pico/data";
const char* TOPIC_CMD  = "iem/task3/student01/pico/cmd";

// ── Robustness ───────────────────────────────────────────────────────
const char*          SHARED_TOKEN    = "iem2026";
const char*          EXPECTED_SOURCE = "nano";        // Pico only accepts msgs from "nano"
const unsigned long  WATCHDOG_TIMEOUT = 10000UL;      // 10 s
const int            INTERVAL_MIN    = 50;
const int            INTERVAL_MAX    = 2000;

// ── Hardware pins ─────────────────────────────────────────────────────
const int LED_PIN    = 28;   // GP28 – external LED
const int BUTTON_PIN =  2;   // GP2  – push button
const int POT_PIN    = 26;   // GP26 – potentiometer (ADC0)

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

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

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

// ── Button edge detection ─────────────────────────────────────────────
bool lastButtonState = HIGH;

// ── Objects ───────────────────────────────────────────────────────────
WiFiClient   wifiClient;
PubSubClient mqtt(wifiClient);
SimpleJson   jsonOut, jsonIn;

// =====================================================================
// WiFi
// =====================================================================
void setupWiFi() {
    Serial.print("[WiFi] Connecting");
    WiFi.begin(WIFI_SSID, WIFI_PASS);
    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 40) {
        delay(500);
        Serial.print(".");
        attempts++;
    }
    if (WiFi.status() == WL_CONNECTED)
        Serial.println("\n[WiFi] Connected! IP: " + WiFi.localIP().toString());
    else
        Serial.println("\n[WiFi] FAILED – check SSID/password");
}

// =====================================================================
// Safe state
// =====================================================================
void enterSafeState() {
    if (!inSafeState) {
        inSafeState      = true;
        overrideInterval = -1;          // fall back to local poti
        Serial.println("[WATCHDOG] Timeout! → SAFE STATE (local poti control)");
    }
}

void leaveSafeState() {
    if (inSafeState) {
        inSafeState = false;
        Serial.println("[WATCHDOG] Valid command received → leaving safe state");
    }
}

// =====================================================================
// Robustness: validate token + source
// =====================================================================
bool validateMessage(const SimpleJson& msg) {
    // Token check
    if (!msg.hasKey("token")) { msgRejectedAuth++; return false; }
    if (strcmp(msg.getString("token"), SHARED_TOKEN) != 0) {
        Serial.println("[AUTH] Invalid token → rejected");
        msgRejectedAuth++;
        return false;
    }
    // Source check  (Pico only accepts messages from "nano")
    if (!msg.hasKey("source")) { msgRejectedAuth++; return false; }
    if (strcmp(msg.getString("source"), EXPECTED_SOURCE) != 0) {
        Serial.println("[AUTH] Invalid source → rejected");
        msgRejectedAuth++;
        return false;
    }
    return true;
}

// =====================================================================
// Robustness: sequence number check
// =====================================================================
bool checkSequence(const SimpleJson& msg) {
    int seq = msg.getInt("seq");

    // seq == 0 means sender restarted → resync
    if (seq == 0) { expectedSeqNr = 1; return true; }

    // First message ever
    if (expectedSeqNr == -1) { expectedSeqNr = seq + 1; return true; }

    // Replay / duplicate → discard
    if (seq < expectedSeqNr) {
        Serial.print("[SEQ] Replay detected (got "); Serial.print(seq);
        Serial.print(", expected "); Serial.print(expectedSeqNr);
        Serial.println(") → discarded");
        return false;
    }

    // Gap → warn but accept
    if (seq > expectedSeqNr) {
        msgSeqGaps++;
        Serial.print("[SEQ] Gap! expected "); Serial.print(expectedSeqNr);
        Serial.print(", got "); Serial.println(seq);
    }

    expectedSeqNr = seq + 1;
    return true;
}

// =====================================================================
// Process incoming command fields
// =====================================================================
void processCommand(const SimpleJson& cmd) {
    // blinkEnabled
    if (cmd.hasKey("blinkEnabled")) {
        blinkEnabled = cmd.getBool("blinkEnabled");
        Serial.print("[CMD] blinkEnabled = ");
        Serial.println(blinkEnabled ? "true" : "false");
    }
    // ledOn – set LED permanently on/off
    if (cmd.hasKey("ledOn")) {
        ledState = cmd.getBool("ledOn");
        if (!blinkEnabled) digitalWrite(LED_PIN, ledState ? HIGH : LOW);
        Serial.print("[CMD] ledOn = ");
        Serial.println(ledState ? "true" : "false");
    }
    // interval override with range check
    if (cmd.hasKey("interval")) {
        int iv = cmd.getInt("interval");
        if (iv >= INTERVAL_MIN && iv <= INTERVAL_MAX) {
            overrideInterval = iv;
            Serial.print("[CMD] interval = "); Serial.println(iv);
        } else {
            Serial.println("[CMD] interval out of range 50–2000 → rejected");
        }
    }
    // useLocalPot – reset override, go back to poti
    if (cmd.hasKey("useLocalPot") && cmd.getBool("useLocalPot")) {
        overrideInterval = -1;
        Serial.println("[CMD] useLocalPot → back to poti control");
    }
}

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

    Serial.print("[MQTT RX] "); Serial.println(message);

    // 2. Parse JSON
    if (!jsonIn.parse(message)) {
        msgRejectedJson++;
        Serial.println("[JSON] Parse failed → rejected");
        return;
    }

    // 3. Validate token + source
    if (!validateMessage(jsonIn)) return;

    // 4. Check sequence number
    if (!checkSequence(jsonIn)) return;

    // 5. Reset watchdog + stats
    lastValidCmd = millis();
    msgAccepted++;
    leaveSafeState();

    // 6. Execute command
    processCommand(jsonIn);
}

// =====================================================================
// MQTT reconnect (also used on first connect)
// =====================================================================
void mqttReconnect() {
    while (!mqtt.connected()) {
        Serial.print("[MQTT] Connecting...");
        // Use a random client ID so two Pico W boards don't clash
        String clientId = "pico-w-" + String(random(0xFFFF), HEX);
        if (mqtt.connect(clientId.c_str())) {
            Serial.println(" connected!");
            mqtt.subscribe(TOPIC_CMD);
            Serial.print("[MQTT] Subscribed to: "); Serial.println(TOPIC_CMD);
        } else {
            Serial.print(" failed, rc="); Serial.print(mqtt.state());
            Serial.println(" – retrying in 3 s");
            delay(3000);
        }
    }
}

// =====================================================================
// Publish sensor data to TOPIC_DATA
// =====================================================================
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);
    Serial.print("[MQTT TX] "); Serial.println(buf);
}

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

    Serial.begin(115200);
    delay(1000);
    Serial.println("=== Task 3 – Pico W (device) ===");

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

    lastValidCmd = millis();   // start watchdog clock
}

// =====================================================================
// loop()
// =====================================================================
void loop() {
    // ── MQTT keep-alive ─────────────────────────────────────────────
    if (!mqtt.connected()) mqttReconnect();
    mqtt.loop();   // ← MUST be called every iteration!

    unsigned long now = millis();

    // ── Heartbeat: LED_BUILTIN blinks every 300 ms ──────────────────
    if (now - lastHeartbeat >= 300) {
        lastHeartbeat  = now;
        heartbeatState = !heartbeatState;
        digitalWrite(LED_BUILTIN, heartbeatState ? HIGH : LOW);
    }

    // ── Watchdog: enter safe state if no valid cmd for 10 s ─────────
    if (!inSafeState && (now - lastValidCmd > WATCHDOG_TIMEOUT)) {
        enterSafeState();
    }

    // ── Button: toggle blink on/off (edge detection) ─────────────────
    bool buttonState = digitalRead(BUTTON_PIN);
    if (buttonState == LOW && lastButtonState == HIGH) {
        blinkEnabled = !blinkEnabled;
        Serial.print("[BTN] Blink = "); Serial.println(blinkEnabled ? "ON" : "OFF");
    }
    lastButtonState = buttonState;

    // ── Potentiometer: map to blink interval ──────────────────────────
    int potValue = analogRead(POT_PIN);
    unsigned long blinkInterval;
    if (overrideInterval != -1)
        blinkInterval = overrideInterval;                   // remote override
    else
        blinkInterval = map(potValue, 0, 1023, 50, 1000);  // local poti

    // ── Blink external LED ────────────────────────────────────────────
    if (blinkEnabled) {
        if (now - lastBlink >= blinkInterval) {
            lastBlink = now;
            ledState  = !ledState;
            digitalWrite(LED_PIN, ledState ? HIGH : LOW);
        }
    }

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