#include <Arduino.h>
#include <ArduinoJson.h>
#include <ArduinoOTA.h>
#include <ESPmDNS.h>
#include <esp_ota_ops.h>
#include <esp_task_wdt.h>
#include <esp_http_client.h>
#include <esp_https_ota.h>

#include "AutomationManager.h"
#include "StepperManager.h"
#include "TouchManager.h"
#include "async_network.h"
#include "ble_provisioning.h"
#include "config.h"
#include "hardware.h"
#include "mqtt_manager.h"
#include "sinric_manager.h"
#include "storage.h"
#include "diagnostics_manager.h"

// ============================================================
// Serial Tee — Captures Serial output for BLE rescue streaming
// Wraps the hardware serial so all Serial.print() calls also
// feed into the BLE ring buffer for live serial monitoring.
// ============================================================
class SerialTee : public Print {
public:
  SerialTee() : _enabled(false) {}

  void enable() { _enabled = true; }
  void disable() { _enabled = false; }

  size_t write(uint8_t c) override {
    Serial.write(c); // Always write to real serial
    if (_enabled) {
      char ch = (char)c;
      bleProvisioning.captureSerial(&ch, 1);
    }
    return 1;
  }

  size_t write(const uint8_t *buffer, size_t size) override {
    Serial.write(buffer, size); // Always write to real serial
    if (_enabled) {
      bleProvisioning.captureSerial((const char *)buffer, size);
    }
    return size;
  }

private:
  bool _enabled;
};

static SerialTee serialTee;

// ============================================================
// Rescue Mode State
// ============================================================
static bool rescueModeActive = false;
static unsigned long rescueAutoTriggerStart = 0;
static const unsigned long RESCUE_AUTO_TRIGGER_MS = 60000; // 60 seconds

// Global objects
AsyncNetwork network;

// MQTT connection watchdog: reboot if MQTT stays disconnected for this long
// ONLY triggers if the broker was previously reachable (prevents infinite
// reboot loops when running without a broker or internet).
#define MQTT_WATCHDOG_TIMEOUT_MS 1800000UL  // 30 minutes
static unsigned long mqttDisconnectedSince = 0;

void setup() {
  // CRITICAL: Disable the ESP-IDF Task Watchdog Timer (TWDT) immediately.
  // On ESP-IDF 5.x (Arduino-ESP32 3.x), the TWDT is auto-enabled and libraries
  // (SinricPro WebSockets, Arduino core) call esp_task_wdt_reset() from tasks
  // not registered with the TWDT, causing non-stop serial spam:
  //   "E (xxx) task_wdt: esp_task_wdt_reset(707): task not found"
  // This firmware has its own 30-minute MQTT watchdog for crash recovery.
  // Ignore return values — these may fail if TWDT was never initialized,
  // which is fine (it just means there's nothing to disable).
  (void)esp_task_wdt_delete(NULL);  // Remove current task if auto-subscribed
  (void)esp_task_wdt_deinit();      // Tear down TWDT entirely

  Serial.begin(SERIAL_BAUD);
  delay(1000);
  Serial.printf("\n\n--- StepperMote %s ---\n", FIRMWARE_VERSION);

  // Initialize Storage
  storage.begin();

  // ========================================
  // BOOT CRASH COUNTER + SAFE MODE
  // ========================================
  // Detect firmware crash loops and enter safe mode after 3 consecutive crashes.
  // Matches Zaylo's proven pattern: only PANIC and WDT resets are counted.
  // Brownouts indicate power supply issues, not firmware bugs.
  static bool safeMode = false;
  esp_reset_reason_t resetReason = esp_reset_reason();
  Serial.printf("[Boot] Reset reason: %d\n", (int)resetReason);

  if (resetReason == ESP_RST_PANIC || resetReason == ESP_RST_TASK_WDT ||
      resetReason == ESP_RST_INT_WDT) {
    Serial.println("[Boot] WARNING: Abnormal reset detected (firmware crash)");
    config.boot_crash_count++;

    if (config.boot_crash_count >= 3) {
      Serial.println("[Boot] CRITICAL: Boot loop detected — entering SAFE MODE!");
      safeMode = true;
      config.boot_crash_count = 0; // Reset before save so next boot starts clean
    }
    storage.save(); // Single write: either incremented count or reset-to-0
  } else if (resetReason == ESP_RST_BROWNOUT) {
    Serial.println("[Boot] WARNING: Brownout reset (power supply issue — not counted as crash)");
  } else {
    // Normal boot — reset counter if it was non-zero
    if (config.boot_crash_count != 0) {
      config.boot_crash_count = 0;
      storage.save();
    }
  }

  // TRUE SAFE MODE: Block all dangerous initialization.
  // Only serial and auto-reboot are available. Touch sensor is disabled.
  // Auto-recovers after 10 minutes to prevent permanent bricking.
  if (safeMode) {
    Serial.println("[SafeMode] *** SAFE MODE ACTIVE ***");
    Serial.println("[SafeMode] All subsystems disabled to prevent crash");
    Serial.println("[SafeMode] Type 'reboot' over serial to exit");
    Serial.println("[SafeMode] Auto-reboot in 10 minutes");

    unsigned long safeModeStart = millis();
    const unsigned long SAFE_MODE_AUTO_REBOOT_MS = 10UL * 60UL * 1000UL;

    while (true) {
      delay(100);

      // AUTO-RECOVERY: Reboot after 10 minutes
      if (millis() - safeModeStart > SAFE_MODE_AUTO_REBOOT_MS) {
        Serial.println("[SafeMode] Auto-recovery timeout — rebooting normally");
        delay(500);
        ESP.restart();
      }

      // SERIAL ESCAPE: Type "reboot" or "exit" to exit safe mode
      if (Serial.available()) {
        String cmd = Serial.readStringUntil('\n');
        cmd.trim();
        if (cmd == "reboot" || cmd == "exit") {
          Serial.println("[SafeMode] Serial command received — rebooting");
          delay(500);
          ESP.restart();
        }
      }

      // Periodic status
      static unsigned long lastStatusLog = 0;
      if (millis() - lastStatusLog > 10000) {
        unsigned long remaining = (SAFE_MODE_AUTO_REBOOT_MS - (millis() - safeModeStart)) / 1000;
        Serial.printf("[SafeMode] Reboot in %lum%lus\n", remaining / 60, remaining % 60);
        lastStatusLog = millis();
      }
    }
  }

  // Initialize Hardware & Systems
  stepperManager.begin();
  touchManager.begin();
  automationManager.begin();

  // Setup / Provisioning Mode
  if (!storage.isSetupComplete() || strlen(config.wifi_ssid) == 0) {
    Serial.println("Initial setup required. Starting BLE provisioning...");
    bleProvisioning.begin();
  } else {
    // Normal Boot
    Serial.println("Booting normal operation mode...");
    // Setup OTA Updates
    ArduinoOTA.setHostname("StepperMote-OTA");

    ArduinoOTA.onStart([]() {
      String type;
      if (ArduinoOTA.getCommand() == U_FLASH) {
        type = "sketch";
      } else { // U_SPIFFS
        type = "filesystem";
      }
      Serial.println("Start updating " + type);
    });

    ArduinoOTA.onEnd([]() { Serial.println("\nEnd"); });

    ArduinoOTA.onProgress([](unsigned int progress, unsigned int total) {
      // FIX: Original formula `progress / (total / 100)` causes integer
      // division-by-zero when total < 100. Use multiply-first form.
      unsigned int pct = (total > 0) ? (progress * 100 / total) : 0;
      Serial.printf("Progress: %u%%\r", pct);
    });

    ArduinoOTA.onError([](ota_error_t error) {
      Serial.printf("Error[%u]: ", error);
      String errStr = "Unknown";
      if (error == OTA_AUTH_ERROR) {
        Serial.println("Auth Failed");
        errStr = "Auth Failed";
      } else if (error == OTA_BEGIN_ERROR) {
        Serial.println("Begin Failed");
        errStr = "Begin Failed";
      } else if (error == OTA_CONNECT_ERROR) {
        Serial.println("Connect Failed");
        errStr = "Connect Failed";
      } else if (error == OTA_RECEIVE_ERROR) {
        Serial.println("Receive Failed");
        errStr = "Receive Failed";
      } else if (error == OTA_END_ERROR) {
        Serial.println("End Failed");
        errStr = "End Failed";
      }
      
      diagnosticsManager.logErrorEvent("OTA", errStr.c_str());
    });

    network.onConnected([]() {
      static bool otaStarted = false;
      if (!otaStarted) {
        // Explicitly start mDNS (helps with discovery on newer ESP32 cores)
        if (MDNS.begin("StepperMote-OTA")) {
            Serial.println("[OTA] mDNS responder started");
        } else {
            Serial.println("[OTA] Error setting up MDNS responder!");
        }
        
        ArduinoOTA.begin();
        otaStarted = true;
        Serial.println("[OTA] ArduinoOTA Started");
      }
      // Lazy-init SinricPro on first WiFi connection (avoids blocking
      // WebSocket setup before WiFi is ready during boot).
      static bool sinricStarted = false;
      if (!sinricStarted) {
        sinricManager.begin();
        sinricStarted = true;
      }
    });

    network.begin(); // Will connect to WiFi and then MQTT
    mqttManager.begin();

    // OTA Rollback: Mark this firmware as valid after successful initialization.
    // If a future OTA update causes a crash-on-boot, the bootloader will
    // automatically revert to this known-good partition.
    esp_ota_mark_app_valid_cancel_rollback();
    Serial.println("[OTA] Firmware marked as valid (rollback protection active)");
  }
}

void loop() {
  ArduinoOTA.handle(); // Handle OTA updates if they are queued

  // Update Core Inputs/Outputs
  stepperManager.update();
  touchManager.update();

  // Network & Protocols
  if (bleProvisioning.isActive()) {
    bleProvisioning.update();

    // Check if provisioning just completed
    if (bleProvisioning.hasReceivedCredentials()) {
      Serial.println("Credentials received via BLE!");

      // Save WiFi credentials
      strlcpy(config.wifi_ssid, bleProvisioning.getReceivedSSID().c_str(),
              sizeof(config.wifi_ssid));
      strlcpy(config.wifi_pass, bleProvisioning.getReceivedPass().c_str(),
              sizeof(config.wifi_pass));

      // Enable MQTT with defaults from secrets.h
      config.mqtt_enabled = (strlen(SECRET_MQTT_BROKER) > 0);
      strlcpy(config.mqtt_broker, SECRET_MQTT_BROKER,
              sizeof(config.mqtt_broker));
      config.mqtt_port = SECRET_MQTT_PORT;
      strlcpy(config.mqtt_user, SECRET_MQTT_USER, sizeof(config.mqtt_user));
      strlcpy(config.mqtt_pass, SECRET_MQTT_PASS, sizeof(config.mqtt_pass));

      // Attempt WiFi connection to verify credentials before reboot
      bleProvisioning.setStatus(BLE_STATUS_CONNECTING);
      WiFi.begin(config.wifi_ssid, config.wifi_pass);
      int attempts = 0;
      while (WiFi.status() != WL_CONNECTED && attempts < 20) {
        delay(500);
        attempts++;
      }

      if (WiFi.status() == WL_CONNECTED) {
        bleProvisioning.setStatus(BLE_STATUS_CONNECTED);
        Serial.printf("WiFi connected! IP: %s\n",
                      WiFi.localIP().toString().c_str());
        // FIX: Force-save all config fields (WiFi + MQTT credentials) BEFORE
        // marking setup complete. This ensures MQTT credentials survive reboot.
        storage.save();
        storage.markSetupComplete(); // Bypasses throttle internally to guarantee persistence
      } else {
        bleProvisioning.setStatus(BLE_STATUS_WIFI_FAIL);
        Serial.println("WiFi connection failed - will clear credentials and "
                       "stay in setup mode after reboot.");
        memset(config.wifi_ssid, 0, sizeof(config.wifi_ssid));
        memset(config.wifi_pass, 0, sizeof(config.wifi_pass));
        storage.save(); // Save empty credentials, setup is NOT complete
      }

      delay(1000);
      bleProvisioning.stop();

      Serial.println("Rebooting...");
      delay(1000);
      ESP.restart();
    }

    // Process BLE Jog Commands
    if (bleProvisioning.hasJogCommand()) {
      stepperManager.jog(bleProvisioning.getJogCommand());
      bleProvisioning.clearJogCommand();
    }

    // Process BLE Config (for Top/Bottom calibration)
    if (bleProvisioning.hasConfigUpdate()) {
      String json = bleProvisioning.getConfigJson();
      bleProvisioning.clearConfigUpdate();

      // Parse config locally during BLE provisioning.
      // mqttManager is NOT initialized in provisioning mode, so we handle
      // the subset of commands that make sense during setup (save_top/save_bottom).
      JsonDocument doc;
      if (!deserializeJson(doc, json)) {
        if (doc.containsKey("cmd")) {
          const char *cmd = doc["cmd"];
          if (strcmp(cmd, "save_top") == 0) {
            config.stepper_top_position = stepperManager.getCurrentAbsolutePosition();
            storage.save();
            Serial.printf("[BLE] Saved top position: %ld\n", config.stepper_top_position);
          } else if (strcmp(cmd, "save_bottom") == 0) {
            config.stepper_bottom_position = stepperManager.getCurrentAbsolutePosition();
            storage.save();
            Serial.printf("[BLE] Saved bottom position: %ld\n", config.stepper_bottom_position);
          }
        }
      }
    }

    // Process WiFi Scans
    if (bleProvisioning.hasScanRequest()) {
      bleProvisioning.clearScanRequest();
      bleProvisioning.setStatus(BLE_STATUS_SCANNING);

      // Ensure WiFi is in Station Mode to scan
      WiFi.mode(WIFI_STA);
      WiFi.disconnect();
      delay(100);

      int n = WiFi.scanNetworks();
      if (n >= 0) {
        String results = "[";
        for (int i = 0; i < min(n, 10); i++) {
          if (i > 0)
            results += ",";
          results += "{\"s\":\"" + WiFi.SSID(i) +
                     "\",\"r\":" + String(WiFi.RSSI(i)) + ",\"e\":" +
                     String(WiFi.encryptionType(i) == WIFI_AUTH_OPEN ? 0 : 1) +
                     "}";
        }
        results += "]";
        bleProvisioning.sendScanResults(results);
      } else {
        bleProvisioning.setStatus(BLE_STATUS_WIFI_FAIL);
      }
      WiFi.scanDelete();
    }

  } else if (rescueModeActive) {
    // ============================================================
    // RESCUE MODE — BLE active for device recovery
    // ============================================================
    bleProvisioning.update();

    // Process rescue commands
    processRescueCommands();

    // Process WiFi credential updates (same flow as initial setup)
    if (bleProvisioning.hasReceivedCredentials()) {
      handleRescueWiFiCredentials();
    }

    // Process WiFi scans in rescue mode
    if (bleProvisioning.hasScanRequest()) {
      handleWiFiScanRequest();
    }

    // Process jog commands even in rescue mode (for recalibration)
    if (bleProvisioning.hasJogCommand()) {
      stepperManager.jog(bleProvisioning.getJogCommand());
      bleProvisioning.clearJogCommand();
    }

    // Process config updates in rescue mode
    if (bleProvisioning.hasConfigUpdate()) {
      handleBLEConfigUpdate();
    }

  } else {
    // ============================================================
    // NORMAL MODE
    // ============================================================
    // RESPONSIVENESS FIX: Process incoming commands (MQTT, SinricPro) FIRST,
    // before network maintenance. network.update() can block up to 1.5s for
    // connectivity checks — if commands are processed after, the user perceives
    // the blind as unresponsive during that window.
    mqttManager.update();
    sinricManager.update();
    automationManager.update();
    storage.flushPending();  // Flush any deferred EEPROM writes

    // Network maintenance AFTER command processing.
    // Disable blocking HTTP connectivity checks while the stepper is moving.
    // The internet check has a 1.5s timeout which freezes the main loop and causes
    // visible stepper stuttering. Stepper motion is time-critical (20kHz timer).
    bool allowBlocking = !stepperManager.isMoving();
    network.update(allowBlocking);
    if (allowBlocking) {
      network.fetchWeatherIfNeeded();
    }

    // --- MQTT Connection Watchdog ---
    if (mqttManager.isEnabled() && network.isConnected() && mqttManager.hasEverConnected()) {
      if (mqttManager.isConnected()) {
        mqttDisconnectedSince = 0;
      } else {
        if (mqttDisconnectedSince == 0) {
          mqttDisconnectedSince = millis();
        } else if (millis() - mqttDisconnectedSince > MQTT_WATCHDOG_TIMEOUT_MS) {
          Serial.println("[Watchdog] MQTT disconnected for 30+ minutes. Rebooting...");
          // Clear crash counter before intentional reboot to prevent false safe mode
          config.boot_crash_count = 0;
          storage.save();
          delay(100);
          ESP.restart();
        }
      }
    }

    // --- Auto-BLE Rescue Trigger ---
    // If WiFi stays disconnected for RESCUE_AUTO_TRIGGER_MS after boot,
    // automatically activate BLE rescue mode for remote recovery.
    if (!network.isConnected() && storage.isSetupComplete()) {
      if (rescueAutoTriggerStart == 0) {
        rescueAutoTriggerStart = millis();
      } else if (millis() - rescueAutoTriggerStart >= RESCUE_AUTO_TRIGGER_MS) {
        Serial.println("\n[RESCUE] WiFi failed for 60+ seconds — activating BLE rescue mode!");
        diagnosticsManager.logGeneralEvent("RESCUE", "Auto-activated BLE rescue (WiFi timeout)");
        activateRescueMode();
      }
    } else {
      // WiFi is connected — reset the auto-trigger timer
      rescueAutoTriggerStart = 0;
    }
  }
}

// ============================================================
// Rescue Mode — Activation / Deactivation
// ============================================================

void activateRescueMode() {
  if (rescueModeActive) return; // Already active

  Serial.println("\n=============================================");
  Serial.println("[RESCUE] Entering BLE Rescue Mode");
  Serial.println("=============================================");

  // Stop network services to free resources for BLE
  // WiFi and BLE can coexist on ESP32-C6, but we pause reconnection
  // to avoid interference.
  WiFi.disconnect(true);
  delay(100);

  rescueModeActive = true;
  bleProvisioning.beginRescue();

  // Enable serial tee to capture all output for BLE streaming
  serialTee.enable();

  Serial.println("[RESCUE] BLE advertising as: " + String(BLEProvisioning::getDeviceName()));
  Serial.println("[RESCUE] Waiting for web app connection...");
  Serial.printf("[RESCUE] Free heap: %d bytes\n", ESP.getFreeHeap());
}

void deactivateRescueMode() {
  if (!rescueModeActive) return;

  Serial.println("[RESCUE] Deactivating rescue mode...");
  serialTee.disable();
  bleProvisioning.stop();
  rescueModeActive = false;
  rescueAutoTriggerStart = 0;

  // Resume normal WiFi operation
  delay(500);
  network.begin();
  mqttManager.begin();

  Serial.println("[RESCUE] Normal operation resumed");
}

// ============================================================
// Rescue Mode — Command Processing
// ============================================================

void processRescueCommands() {
  if (!bleProvisioning.hasCommand()) return;

  String cmdJson = bleProvisioning.getCommand();
  bleProvisioning.clearCommand();

  JsonDocument doc;
  DeserializationError err = deserializeJson(doc, cmdJson);
  if (err) {
    Serial.printf("[RESCUE] Invalid command JSON: %s\n", err.c_str());
    return;
  }

  const char *cmd = doc["cmd"] | "";

  // ----- Reboot -----
  if (strcmp(cmd, "reboot") == 0) {
    Serial.println("[RESCUE] Reboot command received");
    bleProvisioning.setStatus(BLE_STATUS_IDLE);
    delay(500);
    bleProvisioning.stop();
    delay(200);
    ESP.restart();
  }

  // ----- Factory Reset -----
  else if (strcmp(cmd, "factoryReset") == 0) {
    Serial.println("[RESCUE] Factory reset command received!");
    bleProvisioning.setStatus(BLE_STATUS_FACTORY_RESET);
    diagnosticsManager.logGeneralEvent("RESCUE", "Factory reset via BLE");
    delay(500);
    storage.reset();
    delay(200);
    bleProvisioning.stop();
    delay(200);
    ESP.restart();
  }

  // ----- Set WiFi Credentials (alternative to SSID/Pass characteristics) -----
  else if (strcmp(cmd, "setWifi") == 0) {
    const char *ssid = doc["ssid"] | "";
    const char *pass = doc["pass"] | "";

    if (strlen(ssid) == 0) {
      Serial.println("[RESCUE] setWifi: SSID is empty, ignoring");
      return;
    }

    Serial.printf("[RESCUE] Setting WiFi: SSID='%s'\n", ssid);

    // FIX: Preserve old credentials so we can restore them on failure.
    // Previously, bad credentials were saved before testing, which could
    // permanently brick the device's internet connectivity after a failed
    // rescue WiFi update (requiring factory reset to recover).
    char oldSSID[sizeof(config.wifi_ssid)];
    char oldPass[sizeof(config.wifi_pass)];
    strlcpy(oldSSID, config.wifi_ssid, sizeof(oldSSID));
    strlcpy(oldPass, config.wifi_pass, sizeof(oldPass));

    // Apply new credentials to config (in RAM only — NOT saved to flash yet)
    strlcpy(config.wifi_ssid, ssid, sizeof(config.wifi_ssid));
    strlcpy(config.wifi_pass, pass, sizeof(config.wifi_pass));

    // Try connecting
    bleProvisioning.setStatus(BLE_STATUS_CONNECTING);
    WiFi.mode(WIFI_STA);
    WiFi.begin(config.wifi_ssid, config.wifi_pass);

    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 20) {
      delay(500);
      attempts++;
      Serial.print(".");
    }
    Serial.println();

    if (WiFi.status() == WL_CONNECTED) {
      // SUCCESS: Now safe to persist the verified credentials
      storage.save();
      bleProvisioning.setStatus(BLE_STATUS_CONNECTED);
      Serial.printf("[RESCUE] WiFi connected! IP: %s\n", WiFi.localIP().toString().c_str());

      // Send success response via diagnostics channel
      String resp = "{\"wifiStatus\":\"connected\",\"ip\":\"" +
                    WiFi.localIP().toString() + "\",\"rssi\":" +
                    String(WiFi.RSSI()) + "}";
      bleProvisioning.sendDiagnostics(resp);
    } else {
      // FAILURE: Restore previous (known-good) credentials
      strlcpy(config.wifi_ssid, oldSSID, sizeof(config.wifi_ssid));
      strlcpy(config.wifi_pass, oldPass, sizeof(config.wifi_pass));
      // Don't save — the original credentials are still in flash

      bleProvisioning.setStatus(BLE_STATUS_WIFI_FAIL);
      Serial.println("[RESCUE] WiFi connection failed — previous credentials restored");

      String resp = "{\"wifiStatus\":\"failed\"}";
      bleProvisioning.sendDiagnostics(resp);
    }
  }

  // ----- Get Diagnostics -----
  else if (strcmp(cmd, "getDiagnostics") == 0) {
    Serial.println("[RESCUE] Building diagnostics...");
    JsonDocument diagDoc;
    diagnosticsManager.buildDiagnosticsJson(diagDoc, true);

    // Add rescue-specific info
    diagDoc["rescueMode"] = true;
    diagDoc["rescueReason"] = "WiFi failure auto-trigger";

    String diagJson;
    serializeJson(diagDoc, diagJson);
    bleProvisioning.sendDiagnostics(diagJson);
    Serial.printf("[RESCUE] Sent diagnostics (%d bytes)\n", diagJson.length());
  }

  // ----- OTA Update via URL -----
  else if (strcmp(cmd, "otaUrl") == 0) {
    const char *url = doc["url"] | "";
    if (strlen(url) == 0) {
      Serial.println("[RESCUE] otaUrl: URL is empty");
      return;
    }
    Serial.printf("[RESCUE] Starting HTTP OTA from: %s\n", url);
    performHTTPOTA(url);
  }

  // ----- Get Config -----
  else if (strcmp(cmd, "getConfig") == 0) {
    Serial.println("[RESCUE] Sending device config...");

    JsonDocument cfgDoc;
    cfgDoc["firmware"] = FIRMWARE_VERSION;
    cfgDoc["deviceId"] = config.mqtt_device_id;
    cfgDoc["setupComplete"] = config.setup_complete;

    // WiFi
    JsonObject wifi = cfgDoc["wifi"].to<JsonObject>();
    wifi["ssid"] = config.wifi_ssid;
    wifi["connected"] = (WiFi.status() == WL_CONNECTED);
    wifi["rssi"] = WiFi.RSSI();
    wifi["ip"] = WiFi.localIP().toString();

    // MQTT
    JsonObject mqtt = cfgDoc["mqtt"].to<JsonObject>();
    mqtt["enabled"] = config.mqtt_enabled;
    mqtt["broker"] = config.mqtt_broker;
    mqtt["port"] = config.mqtt_port;

    // Stepper
    JsonObject stepper = cfgDoc["stepper"].to<JsonObject>();
    stepper["topPos"] = config.stepper_top_position;
    stepper["bottomPos"] = config.stepper_bottom_position;
    stepper["currentPos"] = stepperManager.getCurrentAbsolutePosition();
    stepper["currentPct"] = stepperManager.getCurrentPercentage();
    stepper["isCalibrated"] = stepperManager.isCalibrated();
    stepper["openSpeed"] = config.stepper_open_speed;
    stepper["closeSpeed"] = config.stepper_close_speed;
    stepper["acceleration"] = config.stepper_acceleration;
    stepper["relaxSteps"] = config.stepper_relax_steps;
    stepper["stopDelayMs"] = config.stepper_stop_delay_ms;

    // Automation
    JsonObject automation = cfgDoc["automation"].to<JsonObject>();
    automation["sunset"] = config.sunset_enabled;
    automation["presence"] = config.presence_enabled;
    automation["morning"] = config.morning_enabled;
    automation["night"] = config.night_enabled;
    automation["temperature"] = config.temp_enabled;

    // TWT
    cfgDoc["twtEnabled"] = config.twt_enabled;

    String cfgJson;
    serializeJson(cfgDoc, cfgJson);
    bleProvisioning.sendDiagnostics(cfgJson);
    Serial.printf("[RESCUE] Sent config (%d bytes)\n", cfgJson.length());
  }

  // ----- Set Config Field -----
  else if (strcmp(cmd, "setConfig") == 0) {
    const char *key = doc["key"] | "";
    Serial.printf("[RESCUE] setConfig: key='%s'\n", key);

    if (strcmp(key, "openSpeed") == 0) {
      config.stepper_open_speed = doc["value"] | 2000;
    } else if (strcmp(key, "closeSpeed") == 0) {
      config.stepper_close_speed = doc["value"] | 2000;
    } else if (strcmp(key, "acceleration") == 0) {
      config.stepper_acceleration = doc["value"] | 2000;
    } else if (strcmp(key, "relaxSteps") == 0) {
      config.stepper_relax_steps = doc["value"] | 128;
    } else if (strcmp(key, "stopDelayMs") == 0) {
      config.stepper_stop_delay_ms = doc["value"] | 3000;
    } else if (strcmp(key, "sunset") == 0) {
      config.sunset_enabled = doc["value"] | false;
    } else if (strcmp(key, "presence") == 0) {
      config.presence_enabled = doc["value"] | false;
    } else if (strcmp(key, "morning") == 0) {
      config.morning_enabled = doc["value"] | false;
    } else if (strcmp(key, "night") == 0) {
      config.night_enabled = doc["value"] | false;
    } else if (strcmp(key, "temperature") == 0) {
      config.temp_enabled = doc["value"] | false;
    } else if (strcmp(key, "mqttBroker") == 0) {
      strlcpy(config.mqtt_broker, doc["value"] | "", sizeof(config.mqtt_broker));
    } else if (strcmp(key, "mqttPort") == 0) {
      config.mqtt_port = doc["value"] | 1883;
    } else if (strcmp(key, "twtEnabled") == 0) {
      config.twt_enabled = doc["value"] | false;
    } else {
      Serial.printf("[RESCUE] Unknown config key: %s\n", key);
      return;
    }

    storage.save();
    Serial.printf("[RESCUE] Config '%s' updated and saved\n", key);

    // Send acknowledgment
    String ack = "{\"configSaved\":true,\"key\":\"" + String(key) + "\"}";
    bleProvisioning.sendDiagnostics(ack);
  }

  // ----- Exit Rescue Mode -----
  else if (strcmp(cmd, "exitRescue") == 0) {
    Serial.println("[RESCUE] Exit rescue command — returning to normal mode");
    deactivateRescueMode();
  }

  // ----- Unknown Command -----
  else {
    Serial.printf("[RESCUE] Unknown command: '%s'\n", cmd);
  }
}

// ============================================================
// Rescue Mode — WiFi Credential Handling
// ============================================================

void handleRescueWiFiCredentials() {
  Serial.println("[RESCUE] WiFi credentials received via BLE characteristics");

  // FIX: Preserve old credentials so we can restore them on failure.
  // Previously, bad credentials were saved before testing, which could
  // permanently brick the device's internet connectivity.
  char oldSSID[sizeof(config.wifi_ssid)];
  char oldPass[sizeof(config.wifi_pass)];
  strlcpy(oldSSID, config.wifi_ssid, sizeof(oldSSID));
  strlcpy(oldPass, config.wifi_pass, sizeof(oldPass));

  // Apply new credentials to config (in RAM only — NOT saved to flash yet)
  strlcpy(config.wifi_ssid, bleProvisioning.getReceivedSSID().c_str(),
          sizeof(config.wifi_ssid));
  strlcpy(config.wifi_pass, bleProvisioning.getReceivedPass().c_str(),
          sizeof(config.wifi_pass));

  bleProvisioning.clearCredentials();

  // Try connecting
  bleProvisioning.setStatus(BLE_STATUS_CONNECTING);
  WiFi.mode(WIFI_STA);
  WiFi.begin(config.wifi_ssid, config.wifi_pass);

  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 20) {
    delay(500);
    attempts++;
    Serial.print(".");
  }
  Serial.println();

  if (WiFi.status() == WL_CONNECTED) {
    // SUCCESS: Now safe to persist the verified credentials
    storage.save();
    bleProvisioning.setStatus(BLE_STATUS_CONNECTED);
    Serial.printf("[RESCUE] WiFi connected! IP: %s\n",
                  WiFi.localIP().toString().c_str());

    // If setup wasn't complete, mark it now
    if (!storage.isSetupComplete()) {
      storage.markSetupComplete();
    }

    // Auto-exit rescue mode after successful connection
    delay(2000);
    Serial.println("[RESCUE] WiFi connected — rebooting to normal mode...");
    bleProvisioning.stop();
    delay(500);
    ESP.restart();
  } else {
    // FAILURE: Restore previous (known-good) credentials
    strlcpy(config.wifi_ssid, oldSSID, sizeof(config.wifi_ssid));
    strlcpy(config.wifi_pass, oldPass, sizeof(config.wifi_pass));
    // Don't save — the original credentials are still in flash

    bleProvisioning.setStatus(BLE_STATUS_WIFI_FAIL);
    Serial.println("[RESCUE] WiFi connection failed — previous credentials restored, staying in rescue mode");
  }
}

// ============================================================
// Rescue Mode — WiFi Scan  (shared with initial setup)
// ============================================================

void handleWiFiScanRequest() {
  bleProvisioning.clearScanRequest();
  bleProvisioning.setStatus(BLE_STATUS_SCANNING);

  WiFi.mode(WIFI_STA);
  WiFi.disconnect();
  delay(100);

  int n = WiFi.scanNetworks();
  if (n >= 0) {
    String results = "[";
    for (int i = 0; i < min(n, 15); i++) { // Up to 15 networks
      if (i > 0) results += ",";
      results += "{\"s\":\"" + WiFi.SSID(i) +
                 "\",\"r\":" + String(WiFi.RSSI(i)) + ",\"e\":" +
                 String(WiFi.encryptionType(i) == WIFI_AUTH_OPEN ? 0 : 1) + "}";
    }
    results += "]";
    bleProvisioning.sendScanResults(results);
  } else {
    bleProvisioning.setStatus(BLE_STATUS_WIFI_FAIL);
  }
  WiFi.scanDelete();
}

// ============================================================
// BLE Config Update Handler (shared)
// ============================================================

void handleBLEConfigUpdate() {
  String json = bleProvisioning.getConfigJson();
  bleProvisioning.clearConfigUpdate();

  JsonDocument doc;
  if (!deserializeJson(doc, json)) {
    if (doc.containsKey("cmd")) {
      const char *cmd = doc["cmd"];
      if (strcmp(cmd, "save_top") == 0) {
        config.stepper_top_position = stepperManager.getCurrentAbsolutePosition();
        storage.save();
        Serial.printf("[BLE] Saved top position: %ld\n", config.stepper_top_position);
      } else if (strcmp(cmd, "save_bottom") == 0) {
        config.stepper_bottom_position = stepperManager.getCurrentAbsolutePosition();
        storage.save();
        Serial.printf("[BLE] Saved bottom position: %ld\n", config.stepper_bottom_position);
      }
    }
  }
}

// ============================================================
// HTTP OTA Update — Downloads firmware from URL
// ============================================================

void performHTTPOTA(const char *url) {
  Serial.printf("[OTA] Starting HTTP OTA from: %s\n", url);
  bleProvisioning.sendOTAProgress(0, "Starting download...");
  diagnosticsManager.logGeneralEvent("OTA", "HTTP OTA started via BLE rescue");

  // Determine if HTTPS
  bool isHttps = (strncmp(url, "https://", 8) == 0);

  esp_http_client_config_t httpConfig = {};
  httpConfig.url = url;
  httpConfig.timeout_ms = 30000;
  httpConfig.keep_alive_enable = true;

  if (isHttps) {
    // Skip certificate validation for OTA (user-provided URL)
    httpConfig.skip_cert_common_name_check = true;
    httpConfig.crt_bundle_attach = NULL;
  }

  bleProvisioning.sendOTAProgress(5, "Connecting to server...");

  esp_https_ota_config_t ota_config = {};
  ota_config.http_config = &httpConfig;

  esp_https_ota_handle_t ota_handle = NULL;
  esp_err_t err = esp_https_ota_begin(&ota_config, &ota_handle);

  if (err != ESP_OK || ota_handle == NULL) {
    Serial.printf("[OTA] Begin failed: %s\n", esp_err_to_name(err));
    bleProvisioning.setStatus(BLE_STATUS_OTA_FAIL);
    bleProvisioning.sendOTAProgress(0, "Download failed - connection error");
    diagnosticsManager.logErrorEvent("OTA", "HTTP OTA begin failed");
    return;
  }

  int imageSize = esp_https_ota_get_image_size(ota_handle);
  Serial.printf("[OTA] Image size: %d bytes\n", imageSize);
  bleProvisioning.sendOTAProgress(10, "Downloading firmware...");

  int lastPct = 10;
  while (true) {
    err = esp_https_ota_perform(ota_handle);
    if (err != ESP_ERR_HTTPS_OTA_IN_PROGRESS) break;

    // Calculate and report progress
    if (imageSize > 0) {
      int read = esp_https_ota_get_image_len_read(ota_handle);
      int pct = 10 + (read * 85 / imageSize); // 10-95% range
      if (pct != lastPct && pct % 5 == 0) {
        lastPct = pct;
        char msg[32];
        snprintf(msg, sizeof(msg), "Downloading... %d%%", pct);
        bleProvisioning.sendOTAProgress((uint8_t)pct, msg);
      }
    }
    yield(); // Allow BLE stack to process
  }

  if (err == ESP_OK) {
    // Verify and finalize
    err = esp_https_ota_finish(ota_handle);
    if (err == ESP_OK) {
      Serial.println("[OTA] Update successful! Rebooting...");
      bleProvisioning.sendOTAProgress(100, "Update complete! Rebooting...");
      bleProvisioning.setStatus(BLE_STATUS_OTA_SUCCESS);
      diagnosticsManager.logGeneralEvent("OTA", "HTTP OTA success via BLE rescue");
      delay(2000);
      bleProvisioning.stop();
      delay(200);
      ESP.restart();
    } else {
      Serial.printf("[OTA] Finish failed: %s\n", esp_err_to_name(err));
      bleProvisioning.setStatus(BLE_STATUS_OTA_FAIL);
      bleProvisioning.sendOTAProgress(0, "Verification failed");
      diagnosticsManager.logErrorEvent("OTA", "HTTP OTA finish failed");
    }
  } else {
    esp_https_ota_abort(ota_handle);
    Serial.printf("[OTA] Download failed: %s\n", esp_err_to_name(err));
    bleProvisioning.setStatus(BLE_STATUS_OTA_FAIL);
    bleProvisioning.sendOTAProgress(0, "Download failed");
    diagnosticsManager.logErrorEvent("OTA", "HTTP OTA download failed");
  }
}
