#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 模式并重启
//
// 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"
//   }
//
// 接线方式（V2 版本：红黄互换 + 黄绿互换）：
//   ESP32 3.3V  -> 原灯板 + / 原电池正极
//   ESP32 IO2   -> 220Ω -> L1 控制点 = 绿灯
//   ESP32 IO3   -> 220Ω -> L2 控制点 = 红灯
//   ESP32 IO4   -> 220Ω -> L3 控制点 = 黄灯
// =====================================================

// =====================================================
// BLE 配置
// =====================================================
const char* BLE_DEVICE_NAME = "AI-Light";
#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"

// =====================================================
// 引脚定义（V2 版本：红黄互换 + 黄绿互换）
// =====================================================
const int GREEN_PIN = 2;
const int RED_PIN = 3;
const int YELLOW_PIN = 4;
const int WIFI_LED_PIN = 8;
const int BTN_BOOT_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;

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


// =====================================================
// 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");
}

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

String getConfigJson() {
  JsonDocument doc;
  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;
  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>());

  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(RED_PIN, 0);
  writeLed(YELLOW_PIN, 0);
  writeLed(GREEN_PIN, 0);
}

void setOnly(int red, int yellow, int green) {
  writeLed(RED_PIN, constrain(red, 0, RED_MAX));
  writeLed(YELLOW_PIN, constrain(yellow, 0, YELLOW_MAX));
  writeLed(GREEN_PIN, 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(BTN_BOOT_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 ? (WiFi.status() == WL_CONNECTED) : bleDeviceConnected;
  if (connected) {
    digitalWrite(WIFI_LED_PIN, LOW);
    return;
  }
  if (millis() - lastStatusLedToggle >= 500) {
    statusLedState = !statusLedState;
    digitalWrite(WIFI_LED_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) { setMode(c->getValue()); }
  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(SERVICE_UUID);

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

void 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) break;
    Serial.println("\nFailed, retrying..."); delay(2000);
  }

  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("\nWiFi failed. Restarting...");
    delay(1000);
    ESP.restart();
  }

  Serial.printf("\nWiFi OK. IP: %s\n", WiFi.localIP().toString().c_str());
  digitalWrite(WIFI_LED_PIN, LOW);
}

void 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;
    }
    Serial.print("."); delay(500);
    if (++attempts > 60) { Serial.println("\nMQTT failed!"); attempts = 0; }
  }
}

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


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

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

  ledcAttach(RED_PIN, PWM_FREQ, PWM_RESOLUTION);
  ledcAttach(YELLOW_PIN, PWM_FREQ, PWM_RESOLUTION);
  ledcAttach(GREEN_PIN, PWM_FREQ, PWM_RESOLUTION);
  pinMode(WIFI_LED_PIN, OUTPUT);
  pinMode(BTN_BOOT_PIN, INPUT_PULLUP);
  digitalWrite(WIFI_LED_PIN, HIGH);
  allOff();

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

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

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

  setupBLE();

  if (useMQTT && isConfigComplete()) {
    connectWiFi();
    connectMQTT();
  } else if (useMQTT) {
    Serial.println("Config incomplete. Use BLE to configure.");
  }

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


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

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

  if (useMQTT) {
    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);
}