#include <WiFi.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <BLEDevice.h>
#include <BLEServer.h>
#include <BLEUtils.h>
#include <BLE2902.h>
#include <Preferences.h>
#include <HTTPClient.h>
#include <Update.h>

// =====================================================
// ESP32-C3 SuperMini + 原玩具公共正极灯板：BLE + MQTT 双模式
//
// 功能：
//   - 开机始终启动 BLE，支持灯效控制 + WiFi/MQTT 配置
//   - MQTT 模式下同时连接 WiFi/MQTT
//   - 运行时长按 BOOT 按钮 3 秒 → 切换 BLE/MQTT 模式并重启
//
// BLE 配置：
//   Service:  b8b7e001-7a6b-4f4f-9a8b-11c0ffee0001
//   Mode:     b8b7e002-... (读写/通知 - 灯效模式)
//   Config:   b8b7e003-... (读写 - WiFi/MQTT 配置 JSON)
//
// 配置 JSON 格式（写入 Config 特征）：
//   {
//     "wifi_ssid": "xxx",
//     "wifi_pass": "xxx",
//     "mqtt_broker": "192.168.1.100",
//     "mqtt_port": 1883,
//     "mqtt_user": "user",
//     "mqtt_pass": "pass",
//     "mqtt_client": "AI-Light",
//     "mqtt_topic": "opencode/status",
//     "mqtt_status": "openCodeLight/status",
//     "ota_url": "https://example.com/firmware.bin"
//   }
//
// 接线方式（V3 版本）：
//   ESP32 3.3V  -> 原灯板 + / 原电池正极
//   ESP32 IO2   -> 220Ω -> L1 控制点 = 黄灯
//   ESP32 IO3   -> 220Ω -> L2 控制点 = 绿灯
//   ESP32 IO4   -> 220Ω -> L3 控制点 = 红灯
// =====================================================

// =====================================================
// BLE 配置
// =====================================================
const char* BLE_DEVICE_NAME = "AI-Light";
const char* FW_VERSION = "1.0.0";
#define SERVICE_UUID        "b8b7e001-7a6b-4f4f-9a8b-11c0ffee0001"
#define MODE_CHAR_UUID      "b8b7e002-7a6b-4f4f-9a8b-11c0ffee0001"
#define CONFIG_CHAR_UUID    "b8b7e003-7a6b-4f4f-9a8b-11c0ffee0001"

// =====================================================
// 引脚定义（红绿黄可通过配置动态修改，状态灯和按钮固定）
// =====================================================
int redPin = 4;
int greenPin = 3;
int yellowPin = 2;
const int STATUS_PIN = 8;
const int BUTTON_PIN = 9;

const int PWM_FREQ = 5000;
const int PWM_RESOLUTION = 8;

const int RED_MAX = 255;
const int YELLOW_MAX = 220;
const int GREEN_MAX = 220;

const unsigned long NORMAL_MODE_TIMEOUT_MS = 5UL * 60UL * 1000UL;
const unsigned long TRAFFIC_MODE_TIMEOUT_MS = 10UL * 60UL * 1000UL;
const unsigned long LONG_PRESS_MS = 3000;

// =====================================================
// 全局状态
// =====================================================
String currentMode = "init";
unsigned long modeStart = 0;
bool useMQTT = true;

WiFiClient wifiClient;
PubSubClient mqttClient(wifiClient);

BLEServer* pServer = nullptr;
BLECharacteristic* pModeCharacteristic = nullptr;
BLECharacteristic* pConfigCharacteristic = nullptr;
bool bleDeviceConnected = false;
bool wifiConnected = false;

Preferences preferences;

String cfgWifiSsid = "";
String cfgWifiPass = "";
String cfgMqttBroker = "";
uint16_t cfgMqttPort = 1883;
String cfgMqttUser = "";
String cfgMqttPass = "";
String cfgMqttClient = "AI-Light";
String cfgMqttTopic = "opencode/status";
String cfgMqttStatus = "openCodeLight/status";
String cfgOtaUrl = "";
bool otaInProgress = false;


// =====================================================
// NVS 配置读写
// =====================================================

void loadConfig() {
  cfgWifiSsid = preferences.getString("wifi_ssid", "");
  cfgWifiPass = preferences.getString("wifi_pass", "");
  cfgMqttBroker = preferences.getString("mqtt_broker", "");
  cfgMqttPort = preferences.getUInt("mqtt_port", 1883);
  cfgMqttUser = preferences.getString("mqtt_user", "");
  cfgMqttPass = preferences.getString("mqtt_pass", "");
  cfgMqttClient = preferences.getString("mqtt_client", "AI-Light");
  cfgMqttTopic = preferences.getString("mqtt_topic", "opencode/status");
  cfgMqttStatus = preferences.getString("mqtt_status", "openCodeLight/status");
  redPin = preferences.getUInt("pin_red", 4);
  greenPin = preferences.getUInt("pin_green", 3);
  yellowPin = preferences.getUInt("pin_yellow", 2);
  cfgOtaUrl = preferences.getString("ota_url", "");
}

bool isConfigComplete() {
  return cfgWifiSsid.length() > 0 && cfgMqttBroker.length() > 0;
}

String getConfigJson() {
  JsonDocument doc;
  doc["fw_version"] = FW_VERSION;
  doc["wifi_ssid"] = cfgWifiSsid;
  doc["mqtt_broker"] = cfgMqttBroker;
  doc["mqtt_port"] = cfgMqttPort;
  doc["mqtt_user"] = cfgMqttUser;
  doc["mqtt_client"] = cfgMqttClient;
  doc["mqtt_topic"] = cfgMqttTopic;
  doc["mqtt_status"] = cfgMqttStatus;
  doc["comm_mode"] = useMQTT ? 1 : 0;
  doc["pin_red"] = redPin;
  doc["pin_green"] = greenPin;
  doc["pin_yellow"] = yellowPin;
  doc["ota_url"] = cfgOtaUrl;
  String out;
  serializeJson(doc, out);
  return out;
}

void saveConfigFromJson(const String& json) {
  JsonDocument doc;
  DeserializationError err = deserializeJson(doc, json);
  if (err) {
    Serial.print("Config JSON parse error: ");
    Serial.println(err.c_str());
    return;
  }

  if (doc.containsKey("wifi_ssid"))
    preferences.putString("wifi_ssid", doc["wifi_ssid"].as<String>());
  if (doc.containsKey("wifi_pass"))
    preferences.putString("wifi_pass", doc["wifi_pass"].as<String>());
  if (doc.containsKey("mqtt_broker"))
    preferences.putString("mqtt_broker", doc["mqtt_broker"].as<String>());
  if (doc.containsKey("mqtt_port"))
    preferences.putUInt("mqtt_port", doc["mqtt_port"].as<uint16_t>());
  if (doc.containsKey("mqtt_user"))
    preferences.putString("mqtt_user", doc["mqtt_user"].as<String>());
  if (doc.containsKey("mqtt_pass"))
    preferences.putString("mqtt_pass", doc["mqtt_pass"].as<String>());
  if (doc.containsKey("mqtt_client"))
    preferences.putString("mqtt_client", doc["mqtt_client"].as<String>());
  if (doc.containsKey("mqtt_topic"))
    preferences.putString("mqtt_topic", doc["mqtt_topic"].as<String>());
  if (doc.containsKey("mqtt_status"))
    preferences.putString("mqtt_status", doc["mqtt_status"].as<String>());
  if (doc.containsKey("pin_red"))
    preferences.putUInt("pin_red", doc["pin_red"].as<uint8_t>());
  if (doc.containsKey("pin_green"))
    preferences.putUInt("pin_green", doc["pin_green"].as<uint8_t>());
  if (doc.containsKey("pin_yellow"))
    preferences.putUInt("pin_yellow", doc["pin_yellow"].as<uint8_t>());
  if (doc.containsKey("ota_url"))
    preferences.putString("ota_url", doc["ota_url"].as<String>());

  Serial.println("Config saved. Restarting...");
  delay(500);
  ESP.restart();
}


// =====================================================
// 基础工具函数：公共正极反相输出
// =====================================================

void writeLed(int pin, int value) {
  value = constrain(value, 0, 255);
  int pwmValue = 255 - value;
  ledcWrite(pin, pwmValue);
}

void allOff() {
  writeLed(redPin, 0);
  writeLed(yellowPin, 0);
  writeLed(greenPin, 0);
}

void setOnly(int red, int yellow, int green) {
  writeLed(redPin, constrain(red, 0, RED_MAX));
  writeLed(yellowPin, constrain(yellow, 0, YELLOW_MAX));
  writeLed(greenPin, constrain(green, 0, GREEN_MAX));
}

int triWave(unsigned long t, unsigned long period, int maxValue) {
  unsigned long x = t % period;
  if (x < period / 2) return map(x, 0, period / 2, 0, maxValue);
  return map(x, period / 2, period, maxValue, 0);
}

int fadeInOutBrightness(
  unsigned long t, unsigned long fadeIn, unsigned long hold,
  unsigned long fadeOut, unsigned long offTime, int maxValue
) {
  unsigned long total = fadeIn + hold + fadeOut + offTime;
  unsigned long x = t % total;
  if (x < fadeIn) return map(x, 0, fadeIn, 0, maxValue);
  x -= fadeIn;
  if (x < hold) return maxValue;
  x -= hold;
  if (x < fadeOut) return map(x, 0, fadeOut, maxValue, 0);
  return 0;
}

void fadeToStatic(int targetRed, int targetYellow, int targetGreen, int fadeMs = 80) {
  allOff();
  int steps = 12;
  int delayPerStep = max(1, fadeMs / steps);
  for (int i = 0; i <= steps; i++) {
    float p = (float)i / steps;
    setOnly(targetRed * p, targetYellow * p, targetGreen * p);
    delay(delayPerStep);
  }
}


// =====================================================
// 模式处理
// =====================================================

bool isValidMode(String mode) {
  return (
    mode == "red" || mode == "yellow" || mode == "green" ||
    mode == "busy" || mode == "error" || mode == "thinking" ||
    mode == "ai" || mode == "success" || mode == "traffic" ||
    mode == "alarm" || mode == "init" || mode == "off" || mode == "idle"
  );
}

void publishStatus() {
  if (useMQTT && mqttClient.connected()) {
    mqttClient.publish(cfgMqttStatus.c_str(), currentMode.c_str(), true);
  }
}

void notifyMode() {
  if (pModeCharacteristic) {
    pModeCharacteristic->setValue(currentMode.c_str());
    if (bleDeviceConnected) pModeCharacteristic->notify();
  }
}

void setMode(String mode) {
  mode.trim();
  mode.toLowerCase();
  if (!isValidMode(mode)) {
    Serial.print("Unknown mode: ");
    Serial.println(mode);
    return;
  }
  if (mode == "idle") mode = "traffic";

  if (mode == currentMode) return;

  currentMode = mode;
  modeStart = millis();

  Serial.print("Mode: ");
  Serial.println(currentMode);

  if (mode == "red") fadeToStatic(RED_MAX, 0, 0, 80);
  else if (mode == "yellow") fadeToStatic(0, YELLOW_MAX, 0, 80);
  else if (mode == "green") fadeToStatic(0, 0, GREEN_MAX, 80);
  else if (mode == "success") setOnly(0, 0, GREEN_MAX);
  else if (mode == "off") allOff();

  publishStatus();
  notifyMode();
}

void autoTimeoutCheck() {
  unsigned long elapsed = millis() - modeStart;
  if (currentMode == "off") return;
  if (currentMode == "traffic") {
    if (elapsed >= TRAFFIC_MODE_TIMEOUT_MS) setMode("off");
    return;
  }
  if (elapsed >= NORMAL_MODE_TIMEOUT_MS) setMode("traffic");
}


// =====================================================
// 灯效模式
// =====================================================

void updateBusy() {
  unsigned long t = millis() - modeStart;
  int y = fadeInOutBrightness(t, 80, 500, 120, 500, YELLOW_MAX);
  setOnly(0, y, 0);
}

void updateError() {
  unsigned long t = millis() - modeStart;
  int r = fadeInOutBrightness(t, 40, 180, 80, 180, RED_MAX);
  setOnly(r, 0, 0);
}

void updateThinking() {
  unsigned long t = millis() - modeStart;
  const unsigned long period = 1050;
  unsigned long x = t % period;
  int g = 0, y = 0, r = 0;
  if (x < 350) {
    g = map(x, 0, 350, GREEN_MAX, 70);
    y = map(x, 0, 350, 20, YELLOW_MAX);
  } else if (x < 700) {
    unsigned long p = x - 350;
    g = map(p, 0, 350, 70, 0);
    y = map(p, 0, 350, YELLOW_MAX, 70);
    r = map(p, 0, 350, 20, RED_MAX);
  } else {
    unsigned long p = x - 700;
    g = map(p, 0, 350, 20, GREEN_MAX);
    y = map(p, 0, 350, 70, 0);
    r = map(p, 0, 350, RED_MAX, 70);
  }
  setOnly(r, y, g);
}

void updateAi() {
  unsigned long t = millis() - modeStart;
  const unsigned long period = 1800;
  unsigned long x = t % period;
  int g = triWave((x + 0) % period, period, 150);
  int y = triWave((x + period / 3) % period, period, 140);
  int r = triWave((x + 2 * period / 3) % period, period, 170);
  setOnly(r, y, g);
}

void updateSuccess() { setOnly(0, 0, GREEN_MAX); }

void updateAlarm() {
  unsigned long t = millis() - modeStart;
  const unsigned long phaseMs = 260;
  int phase = (t / phaseMs) % 2;
  unsigned long inside = t % phaseMs;
  int brightness;
  if (inside < 60) brightness = map(inside, 0, 60, 0, 255);
  else if (inside < 180) brightness = 255;
  else brightness = map(inside, 180, phaseMs, 255, 0);
  if (phase == 0) setOnly(brightness, 0, 0);
  else setOnly(0, min(brightness, YELLOW_MAX), 0);
}

void updateTraffic() {
  unsigned long t = (millis() - modeStart) % 15000;
  if (t < 5000) {
    setOnly(0, 0, GREEN_MAX);
  } else if (t < 6500) {
    unsigned long phase = (t - 5000) % 500;
    int g = 0;
    if (phase < 60) g = map(phase, 0, 60, 0, GREEN_MAX);
    else if (phase < 230) g = GREEN_MAX;
    else if (phase < 320) g = map(phase, 230, 320, GREEN_MAX, 0);
    setOnly(0, 0, g);
  } else if (t < 8500) {
    setOnly(0, YELLOW_MAX, 0);
  } else if (t < 13500) {
    setOnly(RED_MAX, 0, 0);
  } else {
    unsigned long phase = (t - 13500) % 500;
    int r = 0;
    if (phase < 60) r = map(phase, 0, 60, 0, RED_MAX);
    else if (phase < 230) r = RED_MAX;
    else if (phase < 320) r = map(phase, 230, 320, RED_MAX, 0);
    setOnly(r, 0, 0);
  }
}

void updateInit() {
  unsigned long t = millis() - modeStart;
  const unsigned long period = 2500;
  unsigned long x = t % period;
  int brightness;
  if (x < 800) brightness = map(x, 0, 800, 0, 200);
  else if (x < 1200) brightness = 200;
  else if (x < 2000) brightness = map(x, 1200, 2000, 200, 0);
  else brightness = 0;
  setOnly(brightness, brightness, brightness);
}

void breathingGreen(int times) {
  const unsigned long period = 2500;
  for (int i = 0; i < times; i++) {
    unsigned long start = millis();
    while (millis() - start < period) {
      unsigned long t = millis() - start;
      int brightness;
      if (t < 800) brightness = map(t, 0, 800, 0, GREEN_MAX);
      else if (t < 1200) brightness = GREEN_MAX;
      else if (t < 2000) brightness = map(t, 1200, 2000, GREEN_MAX, 0);
      else brightness = 0;
      setOnly(0, 0, brightness);
      delay(5);
    }
  }
  allOff();
}


// =====================================================
// BOOT 按钮处理
// =====================================================

unsigned long bootPressStart = 0;
bool bootWasPressed = false;
bool switchTriggered = false;

void checkBootButton() {
  bool pressed = (digitalRead(BUTTON_PIN) == LOW);
  if (pressed && !bootWasPressed) {
    bootPressStart = millis();
    bootWasPressed = true;
  }
  if (pressed && bootWasPressed && !switchTriggered) {
    if (millis() - bootPressStart >= LONG_PRESS_MS) {
      switchTriggered = true;
      Serial.println("BOOT long press -> switching mode...");
      allOff();
      for (int i = 0; i < 3; i++) {
        setOnly(100, 100, 100); delay(100);
        allOff(); delay(100);
      }
      useMQTT = !useMQTT;
      preferences.putUInt("comm_mode", useMQTT ? 1 : 0);
      delay(200);
      ESP.restart();
    }
  }
  if (!pressed) {
    bootWasPressed = false;
    switchTriggered = false;
  }
}


// =====================================================
// 状态 LED
// =====================================================

unsigned long lastStatusLedToggle = 0;
bool statusLedState = false;

void updateStatusLed() {
  bool connected = (useMQTT && wifiConnected) ? (WiFi.status() == WL_CONNECTED) : bleDeviceConnected;
  if (connected) {
    digitalWrite(STATUS_PIN, LOW);
    return;
  }
  if (millis() - lastStatusLedToggle >= 500) {
    statusLedState = !statusLedState;
    digitalWrite(STATUS_PIN, statusLedState ? LOW : HIGH);
    lastStatusLedToggle = millis();
  }
}


// =====================================================
// BLE
// =====================================================

class ServerCallbacks : public BLEServerCallbacks {
  void onConnect(BLEServer* s) { bleDeviceConnected = true; Serial.println("BLE connected."); }
  void onDisconnect(BLEServer* s) { bleDeviceConnected = false; Serial.println("BLE disconnected."); delay(100); BLEDevice::startAdvertising(); }
};

class ModeCharCallbacks : public BLECharacteristicCallbacks {
  void onWrite(BLECharacteristic* c) {
    String val = c->getValue();
    if (val == "ota") {
      performOTA();
      return;
    }
    setMode(val);
  }
  void onRead(BLECharacteristic* c) { c->setValue(currentMode.c_str()); }
};

class ConfigCharCallbacks : public BLECharacteristicCallbacks {
  void onWrite(BLECharacteristic* c) { saveConfigFromJson(c->getValue()); }
  void onRead(BLECharacteristic* c) { c->setValue(getConfigJson().c_str()); }
};

void setupBLE() {
  Serial.println("Starting BLE...");
  BLEDevice::init(BLE_DEVICE_NAME);
  pServer = BLEDevice::createServer();
  pServer->setCallbacks(new ServerCallbacks());
  BLEService* pService = pServer->createService(BLEUUID(SERVICE_UUID), 20);

  pModeCharacteristic = pService->createCharacteristic(MODE_CHAR_UUID,
    BLECharacteristic::PROPERTY_READ | BLECharacteristic::PROPERTY_WRITE | BLECharacteristic::PROPERTY_NOTIFY);
  pModeCharacteristic->setCallbacks(new ModeCharCallbacks());
  pModeCharacteristic->setValue(currentMode.c_str());
  pModeCharacteristic->addDescriptor(new BLE2902());

  pConfigCharacteristic = pService->createCharacteristic(CONFIG_CHAR_UUID,
    BLECharacteristic::PROPERTY_READ | BLECharacteristic::PROPERTY_WRITE);
  pConfigCharacteristic->setCallbacks(new ConfigCharCallbacks());

  pService->start();
  BLEAdvertising* pAdvertising = BLEDevice::getAdvertising();
  pAdvertising->addServiceUUID(SERVICE_UUID);
  pAdvertising->setScanResponse(true);
  pAdvertising->setMinPreferred(0x0c);
  pAdvertising->setMaxPreferred(0x18);
  BLEDevice::startAdvertising();
  Serial.println("BLE advertising.");
}


// =====================================================
// MQTT
// =====================================================

void mqttCallback(char* topic, byte* payload, unsigned int length) {
  String message = "";
  for (unsigned int i = 0; i < length; i++) message += (char)payload[i];
  Serial.print("MQTT: "); Serial.print(topic); Serial.print(" -> "); Serial.println(message);

  JsonDocument doc;
  if (deserializeJson(doc, message)) { setMode(message); return; }
  const char* code = doc["code"];
  if (code) {
    String c = String(code);
    if (c == "idle") setMode("idle");
    else if (c == "busy" || c == "running") setMode("busy");
    else if (c == "retry" || c == "permission") setMode("alarm");
    else if (c == "pending") setMode("yellow");
    else if (c == "reasoning") setMode("thinking");
    else if (c == "using_tool") setMode("ai");
    else if (c == "error") setMode("error");
    else if (c == "ota") {
      performOTA();
      return;
    }
    else setMode(c);
  }
}

bool connectWiFi() {
  Serial.print("Connecting WiFi");
  WiFi.mode(WIFI_STA);
  WiFi.begin(cfgWifiSsid.c_str(), cfgWifiPass.c_str());

  for (int retry = 1; retry <= 5; retry++) {
    Serial.printf("\nWiFi %d/5", retry);
    int attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 60) {
      delay(5); setOnly(100, 100, 100); updateStatusLed(); Serial.print("."); attempts++;
    }
    if (WiFi.status() == WL_CONNECTED) {
      Serial.printf("\nWiFi OK. IP: %s\n", WiFi.localIP().toString().c_str());
      digitalWrite(STATUS_PIN, LOW);
      return true;
    }
    Serial.println("\nFailed, retrying..."); delay(2000);
  }

  Serial.println("\nWiFi failed. Turning off WiFi.");
  WiFi.disconnect(true);
  WiFi.mode(WIFI_OFF);
  return false;
}

bool connectMQTT() {
  mqttClient.setServer(cfgMqttBroker.c_str(), cfgMqttPort);
  mqttClient.setCallback(mqttCallback);
  Serial.print("MQTT connecting");
  int attempts = 0;
  while (!mqttClient.connected()) {
    if (mqttClient.connect(cfgMqttClient.c_str(), cfgMqttUser.c_str(), cfgMqttPass.c_str())) {
      Serial.println("\nMQTT OK.");
      mqttClient.subscribe(cfgMqttTopic.c_str());
      allOff(); breathingGreen(3); publishStatus();
      return true;
    }
    Serial.print("."); delay(500);
    if (++attempts > 60) { Serial.println("\nMQTT failed!"); return false; }
  }
  return false;
}

void checkMQTTConnection() {
  if (useMQTT && !mqttClient.connected()) { Serial.println("MQTT reconnecting..."); connectMQTT(); }
}


// =====================================================
// OTA 固件更新
// =====================================================

void notifyOtaStatus(const String& status) {
  if (pModeCharacteristic && bleDeviceConnected) {
    pModeCharacteristic->setValue(status.c_str());
    pModeCharacteristic->notify();
  }
  if (useMQTT && mqttClient.connected()) {
    mqttClient.publish(cfgMqttStatus.c_str(), status.c_str(), true);
  }
}

void performOTA() {
  Serial.println("=== OTA Start ===");

  if (otaInProgress) {
    Serial.println("OTA already in progress");
    return;
  }

  if (cfgOtaUrl.length() == 0) {
    Serial.println("OTA URL not configured");
    notifyOtaStatus("ota:no_url");
    return;
  }

  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("WiFi not connected");
    notifyOtaStatus("ota:no_wifi");
    return;
  }

  otaInProgress = true;
  notifyOtaStatus("ota:downloading");

  HTTPClient http;
  http.begin(cfgOtaUrl);
  http.setTimeout(30000);

  int httpCode = http.GET();
  Serial.printf("HTTP GET: %d\n", httpCode);

  if (httpCode != 200) {
    Serial.printf("HTTP error: %d\n", httpCode);
    notifyOtaStatus("ota:error");
    http.end();
    otaInProgress = false;
    return;
  }

  int contentLength = http.getSize();
  Serial.printf("Firmware size: %d bytes\n", contentLength);

  if (contentLength <= 0) {
    Serial.println("Invalid content length");
    notifyOtaStatus("ota:error");
    http.end();
    otaInProgress = false;
    return;
  }

  if (!Update.begin(contentLength)) {
    Serial.println("Update.begin failed");
    notifyOtaStatus("ota:error");
    http.end();
    otaInProgress = false;
    return;
  }

  WiFiClient* stream = http.getStreamPtr();
  uint8_t buf[1024];
  int written = 0;
  unsigned long lastProgress = millis();

  while (http.connected() && written < contentLength) {
    size_t size = stream->available();
    if (size) {
      int bytesRead = stream->readBytes(buf, min(size, sizeof(buf)));
      size_t bytesWritten = Update.write(buf, bytesRead);
      written += bytesWritten;

      if (millis() - lastProgress >= 1000) {
        Serial.printf("Progress: %d/%d (%.1f%%)\n", written, contentLength, (float)written / contentLength * 100);
        lastProgress = millis();
      }
    }
    delay(1);
  }

  Serial.printf("Written: %d/%d\n", written, contentLength);

  if (Update.end(true)) {
    Serial.println("OTA success!");
    notifyOtaStatus("ota:success");
    delay(1000);
    ESP.restart();
  } else {
    Serial.printf("OTA error: %s\n", Update.errorString());
    notifyOtaStatus("ota:error");
  }

  http.end();
  otaInProgress = false;
}


// =====================================================
// 初始化
// =====================================================

void setup() {
  Serial.begin(115200);
  delay(500);

  preferences.begin("ai-light", false);
  useMQTT = (preferences.getUInt("comm_mode", 1) != 0);
  loadConfig();

  ledcAttach(redPin, PWM_FREQ, PWM_RESOLUTION);
  ledcAttach(yellowPin, PWM_FREQ, PWM_RESOLUTION);
  ledcAttach(greenPin, PWM_FREQ, PWM_RESOLUTION);
  pinMode(STATUS_PIN, OUTPUT);
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  digitalWrite(STATUS_PIN, HIGH);
  allOff();

  Serial.println();
  Serial.println("=== AI-Light ===");
  Serial.printf("Mode: %s\n", useMQTT ? "MQTT" : "BLE-only");
  Serial.printf("WiFi: %s\n", cfgWifiSsid.length() > 0 ? cfgWifiSsid.c_str() : "(not set)");
  Serial.printf("MQTT: %s\n", cfgMqttBroker.length() > 0 ? cfgMqttBroker.c_str() : "(not set)");
  Serial.printf("Pins: R=%d G=%d Y=%d\n", redPin, greenPin, yellowPin);

  currentMode = "init";
  modeStart = millis();

  if (useMQTT && isConfigComplete()) {
    wifiConnected = connectWiFi();
    if (wifiConnected) {
      connectMQTT();
      setMode("traffic");
      Serial.println("WiFi/MQTT mode. Long press BOOT (3s) to switch.");
      return;
    }
    Serial.println("WiFi failed. Entering BLE config mode.");
  }

  setupBLE();
  delay(100);

  setMode("traffic");
  Serial.println("BLE config mode. Long press BOOT (3s) to switch.");
}


// =====================================================
// 主循环
// =====================================================

void loop() {
  updateStatusLed();
  checkBootButton();

  if (useMQTT && wifiConnected) {
    if (WiFi.status() != WL_CONNECTED) {
      Serial.println("WiFi lost, reconnecting...");
      wifiConnected = connectWiFi();
      if (!wifiConnected) {
        Serial.println("WiFi failed. BLE config mode active.");
      }
    }
    if (wifiConnected) {
      checkMQTTConnection();
      mqttClient.loop();
    }
  }

  autoTimeoutCheck();

  if (currentMode == "busy") updateBusy();
  else if (currentMode == "error") updateError();
  else if (currentMode == "thinking") updateThinking();
  else if (currentMode == "ai") updateAi();
  else if (currentMode == "success") updateSuccess();
  else if (currentMode == "traffic") updateTraffic();
  else if (currentMode == "alarm") updateAlarm();
  else if (currentMode == "init") updateInit();
  else if (currentMode == "off") allOff();

  delay(5);
}