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

// =====================================================
// ESP32-C3 SuperMini + 原玩具公共正极灯板：BLE + MQTT 双模式
//
// 功能：
//   - 开机始终启动 BLE，支持灯效控制 + WiFi/MQTT 配置
//   - MQTT 模式下同时连接 WiFi/MQTT
//   - 运行长按 BOOT 按钮 3 秒 → 切换 BLE/MQTT 模式并重启
//   - 连续短按 BOOT 按钮 3 下 → 恢复出厂设置并重启
//
// BLE 配置：
//   Service:  b8b7e001-7a6b-4f4f-9a8b-11c0ffee0001
//   Mode:     b8b7e002-... (读写/通知 - 灯效模式)
//   Config:   b8b7e003-... (读写 - WiFi/MQTT 配置 JSON)
//
// 配置 JSON 格式（写入 Config 特征或 MQTT config topic）：
//   {
//     "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": "agent/status",
//     "mqtt_status": "agentLight/status",
//     "mqtt_topic_config": "agent/status/config",
//     "pin_red": 4,
//     "pin_green": 3,
//     "pin_yellow": 2
//   }
//
// 接线方式（默认引脚，可通过配置修改）：
//   ESP32 3.3V  -> 原灯板 + / 原电池正极
//   ESP32 IO4   -> 220Ω -> L3 控制点 = 红灯（默认）
//   ESP32 IO3   -> 220Ω -> L2 控制点 = 绿灯（默认）
//   ESP32 IO2   -> 220Ω -> L1 控制点 = 黄灯（默认）
//   ESP32 IO8   -> 状态 LED（固定）
//   ESP32 IO9   -> BOOT 按钮（固定）
// =====================================================

// =====================================================
// BLE 配置
// =====================================================
const char* BLE_DEVICE_NAME = "AI-Light";
const char* FW_VERSION = "1.1.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;
const unsigned long TRIPLE_PRESS_WINDOW_MS = 1500;
const int TRIPLE_PRESS_COUNT = 3;

// 非阻塞连接超时
const unsigned long WIFI_CONNECT_TIMEOUT_MS = 10000;
const unsigned long MQTT_CONNECT_TIMEOUT_MS = 5000;
const unsigned long MQTT_RECONNECT_INTERVAL_MS = 3000;

// =====================================================
// 连接状态机
// =====================================================
enum ConnState {
  CONN_IDLE,
  CONN_WIFI_CONNECTING,
  CONN_MQTT_CONNECTING,
  CONN_CONNECTED,
  CONN_WIFI_FAILED,
  CONN_MQTT_FAILED
};

ConnState connState = CONN_IDLE;
unsigned long connStateStart = 0;
unsigned long lastMqttReconnectAttempt = 0;
int connRetryCount = 0;
const int MAX_CONN_RETRIES = 3;

// =====================================================
// 全局状态
// =====================================================
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 = "agent/status";
String cfgMqttStatus = "agentLight/status";
String cfgMqttTopicConfig = "agent/status/config";


// =====================================================
// 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", "agent/status");
  cfgMqttStatus = preferences.getString("mqtt_status", "agentLight/status");
  cfgMqttTopicConfig = preferences.getString("mqtt_topic_cfg", "agent/status/config");
  redPin = preferences.getUInt("pin_red", 4);
  greenPin = preferences.getUInt("pin_green", 3);
  yellowPin = preferences.getUInt("pin_yellow", 2);
}

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["mqtt_topic_config"] = cfgMqttTopicConfig;
  doc["comm_mode"] = useMQTT ? 1 : 0;
  doc["pin_red"] = redPin;
  doc["pin_green"] = greenPin;
  doc["pin_yellow"] = yellowPin;
  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("mqtt_topic_config"))
    preferences.putString("mqtt_topic_cfg", doc["mqtt_topic_config"].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>());

  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);
}


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

unsigned long bootPressStart = 0;
bool bootWasPressed = false;
bool switchTriggered = false;
unsigned long lastPressEnd = 0;
int pressCount = 0;

void factoryReset() {
  Serial.println("Factory reset! Clearing all config...");
  allOff();
  for (int i = 0; i < 5; i++) {
    setOnly(255, 0, 0); delay(200);
    allOff(); delay(200);
  }
  preferences.clear();
  Serial.println("NVS cleared. Restarting...");
  delay(500);
  ESP.restart();
}

void checkBootButton() {
  bool pressed = (digitalRead(BUTTON_PIN) == LOW);

  if (pressed && !bootWasPressed) {
    bootPressStart = millis();
    bootWasPressed = true;
    switchTriggered = false;
  }

  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) {
    bootWasPressed = false;
    unsigned long pressDuration = millis() - bootPressStart;
    if (pressDuration < LONG_PRESS_MS) {
      if (millis() - lastPressEnd > TRIPLE_PRESS_WINDOW_MS) {
        pressCount = 0;
      }
      pressCount++;
      lastPressEnd = millis();
      if (pressCount >= TRIPLE_PRESS_COUNT) {
        pressCount = 0;
        factoryReset();
      }
    }
  }
}


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

void updateStatusLed() {
  if (useMQTT) {
    if (wifiConnected) {
      digitalWrite(STATUS_PIN, LOW);
    } else {
      digitalWrite(STATUS_PIN, HIGH);
    }
  } else {
    if (bleDeviceConnected) {
      digitalWrite(STATUS_PIN, LOW);
    } else {
      digitalWrite(STATUS_PIN, HIGH);
    }
  }
}


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

class ServerCallbacks : public BLEServerCallbacks {
  void onConnect(BLEServer* s) {
    bleDeviceConnected = true;
    Serial.println("BLE connected.");
    if (currentMode == "init") setMode("traffic");
  }
  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();
    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);

  if (String(topic) == cfgMqttTopicConfig) {
    saveConfigFromJson(message);
    return;
  }

  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 setMode(c);
  }
}

void mqttSubscribe() {
  mqttClient.subscribe(cfgMqttTopic.c_str());
  mqttClient.subscribe(cfgMqttTopicConfig.c_str());
}

void updateConnection() {
  unsigned long now = millis();

  switch (connState) {
    case CONN_IDLE:
      break;

    case CONN_WIFI_CONNECTING:
      if (WiFi.status() == WL_CONNECTED) {
        Serial.printf("\nWiFi OK. IP: %s\n", WiFi.localIP().toString().c_str());
        wifiConnected = true;
        digitalWrite(STATUS_PIN, LOW);
        mqttClient.setServer(cfgMqttBroker.c_str(), cfgMqttPort);
        mqttClient.setCallback(mqttCallback);
        connState = CONN_MQTT_CONNECTING;
        connStateStart = now;
        Serial.println("MQTT connecting...");
      } else if (now - connStateStart > WIFI_CONNECT_TIMEOUT_MS) {
        Serial.println("\nWiFi connect timeout.");
        WiFi.disconnect(true);
        WiFi.mode(WIFI_OFF);
        wifiConnected = false;
        connState = CONN_WIFI_FAILED;
      }
      break;

    case CONN_MQTT_CONNECTING:
      if (mqttClient.connect(cfgMqttClient.c_str(), cfgMqttUser.c_str(), cfgMqttPass.c_str())) {
        Serial.println("MQTT OK.");
        mqttSubscribe();
        setMode("traffic");
        connState = CONN_CONNECTED;
        connRetryCount = 0;
      } else if (now - connStateStart > MQTT_CONNECT_TIMEOUT_MS) {
        Serial.println("MQTT connect timeout.");
        connState = CONN_MQTT_FAILED;
      }
      break;

    case CONN_CONNECTED:
      break;

    case CONN_WIFI_FAILED:
      if (connRetryCount >= MAX_CONN_RETRIES) {
        Serial.println("Max WiFi retries reached. Falling back to BLE mode...");
        WiFi.disconnect(true);
        WiFi.mode(WIFI_OFF);
        wifiConnected = false;
        connState = CONN_IDLE;
        setMode("init");
        setupBLE();
        Serial.println("BLE advertising. Waiting for connection...");
      }
      break;

    case CONN_MQTT_FAILED:
      if (currentMode != "error") setMode("error");
      break;
  }
}

void startConnection() {
  Serial.printf("Starting WiFi connection (attempt %d/%d)...\n", connRetryCount + 1, MAX_CONN_RETRIES);
  WiFi.mode(WIFI_STA);
  WiFi.begin(cfgWifiSsid.c_str(), cfgWifiPass.c_str());
  connState = CONN_WIFI_CONNECTING;
  connStateStart = millis();
}

void handleMqttReconnect() {
  if (connState != CONN_CONNECTED) return;
  if (mqttClient.connected()) return;

  unsigned long now = millis();
  if (now - lastMqttReconnectAttempt < MQTT_RECONNECT_INTERVAL_MS) return;
  lastMqttReconnectAttempt = now;

  Serial.println("MQTT reconnecting...");
  if (mqttClient.connect(cfgMqttClient.c_str(), cfgMqttUser.c_str(), cfgMqttPass.c_str())) {
    Serial.println("MQTT reconnected.");
    mqttSubscribe();
    publishStatus();
  }
}


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

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);

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

  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);

  if (useMQTT && isConfigComplete()) {
    startConnection();
    Serial.println("WiFi/MQTT mode. Long press BOOT (3s) to switch.");
  } else {
    setupBLE();
    Serial.println("BLE advertising. Waiting for connection...");
  }
}


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

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

  if (useMQTT && isConfigComplete()) {
    updateConnection();

    if (connState == CONN_CONNECTED) {
      if (WiFi.status() != WL_CONNECTED) {
        Serial.println("WiFi lost, restarting connection...");
        wifiConnected = false;
        connState = CONN_IDLE;
        startConnection();
      } else {
        handleMqttReconnect();
        mqttClient.loop();
      }
    } else if (connState == CONN_WIFI_FAILED) {
      if (connRetryCount < MAX_CONN_RETRIES) {
        static unsigned long lastRetry = 0;
        if (millis() - lastRetry > 5000) {
          lastRetry = millis();
          connRetryCount++;
          startConnection();
        }
      }
    } else if (connState == CONN_MQTT_FAILED) {
      static unsigned long lastMqttRetry = 0;
      if (millis() - lastMqttRetry > MQTT_RECONNECT_INTERVAL_MS) {
        lastMqttRetry = millis();
        connState = CONN_MQTT_CONNECTING;
        connStateStart = millis();
        Serial.println("MQTT retrying...");
      }
    }
  }

  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);
}