#include <WiFi.h>
#include <WebSocketMCP.h>
#include <ArduinoJson.h>

// ================================================================
// 1. CAU HINH PIN (7 thiet bi - Kich muc cao)
// ================================================================
#define DEN_SAN             23
#define PHONG_KHACH         22
#define PHONG_RIENG         21
#define DEN_DOC_SACH        19
#define QUAT                18
#define MAY_BOM              5
#define DEN_TIEC             4

const int allPins[] = {DEN_SAN, PHONG_KHACH, PHONG_RIENG, DEN_DOC_SACH, QUAT, MAY_BOM, DEN_TIEC};
const char* deviceNames[] = {"Den san", "Den phong khach", "Den phong rieng", "Den doc sach", "Quat", "May bom", "Den mo tiec"};
const int deviceCount = 7;

// ================================================================
// 2. CAU HINH HE THONG
// ================================================================
const char* ssid = "Wokwi-GUEST";
const char* password = "";
const char* mcpEndpoint = "wss://api.xiaozhi.me/mcp/?token=eyJhbGciOiJFUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySWQiOjgzOTEwNywiYWdlbnRJZCI6MTUxNzU0MSwiZW5kcG9pbnRJZCI6ImFnZW50XzE1MTc1NDEiLCJwdXJwb3NlIjoibWNwLWVuZHBvaW50IiwiaWF0IjoxNzc0MDAzMjYzLCJleHAiOjE4MDU1NjA4NjN9.0GfzygKNbl8B0UJ76RnsuCsOhvC3LRDZvAjcyrlDXdQpqoHAsSoSssVmlqVBWlHzLyyofg-CYXC1IfAuOg8Prw";

WebSocketMCP mcpClient;
bool isFailsafeTriggered = false;

// ================================================================
// 3. HAM HO TRO VA DANG KY TOOLS
// ================================================================

void shutDownAll() {
    for (int i = 0; i < deviceCount; i++) {
        digitalWrite(allPins[i], LOW); 
    }
    Serial.println("\n[AN TOAN] Mat ket noi WiFi! Da ngat toan bo thiet bi.");
}

void registerMcpTools() {
    auto controlFunc = [](const String& args, int pin) {
        JsonDocument doc;
        deserializeJson(doc, args);
        String state = doc["state"].as<String>();
        digitalWrite(pin, (state == "on") ? HIGH : LOW);
        return WebSocketMCP::ToolResponse("{\"success\":true}");
    };

    mcpClient.registerTool("yard_light", "Bat/Tat Den san", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})", [controlFunc](const String& a){ return controlFunc(a, DEN_SAN); });
    mcpClient.registerTool("living_room_light", "Bat/Tat Den phong khach", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})", [controlFunc](const String& a){ return controlFunc(a, PHONG_KHACH); });
    mcpClient.registerTool("private_room_light", "Bat/Tat Den phong rieng", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})", [controlFunc](const String& a){ return controlFunc(a, PHONG_RIENG); });
    mcpClient.registerTool("reading_light", "Bat/Tat Den doc sach", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})", [controlFunc](const String& a){ return controlFunc(a, DEN_DOC_SACH); });
    mcpClient.registerTool("fan_control", "Bat/Tat Quat", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})", [controlFunc](const String& a){ return controlFunc(a, QUAT); });
    mcpClient.registerTool("pump_control", "Bat/Tat May bom", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})", [controlFunc](const String& a){ return controlFunc(a, MAY_BOM); });
    mcpClient.registerTool("party_light", "Bat/Tat Den nhap nhay hoac Den mo tiec", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})", [controlFunc](const String& a){ return controlFunc(a, DEN_TIEC); });

    mcpClient.registerTool("get_detailed_status", "Kiem tra thiet bi nao dang bat hay tat", R"({"type":"object","properties":{}})", 
        [](const String& args) {
            JsonDocument res;
            int onCount = 0;
            String onList = "", offList = "";
            for (int i = 0; i < deviceCount; i++) {
                if (digitalRead(allPins[i]) == HIGH) {
                    onCount++; onList += String(deviceNames[i]) + ", ";
                } else {
                    offList += String(deviceNames[i]) + ", ";
                }
            }
            res["so_luong_dang_bat"] = onCount;
            res["danh_sach_bat"] = onList.length() > 0 ? onList : "Khong co";
            res["danh_sach_tat"] = offList.length() > 0 ? offList : "Khong co";
            String output; serializeJson(res, output);
            return WebSocketMCP::ToolResponse(output);
        }
    );

    mcpClient.registerTool("all_devices_control", "Dieu khien Bat hoac Tat tat ca cung luc", R"({"type":"object","properties":{"state":{"type":"string","enum":["on","off"]}},"required":["state"]})",
        [](const String& args) {
            JsonDocument doc; deserializeJson(doc, args);
            int val = (doc["state"].as<String>() == "on") ? HIGH : LOW;
            for (int i = 0; i < deviceCount; i++) digitalWrite(allPins[i], val);
            return WebSocketMCP::ToolResponse("{\"success\":true}");
        }
    );

    Serial.println("[MCP] He thong 7 thiet bi da san sang.");
}

void onConnectionStatus(bool connected) {
    if (connected) {
        Serial.println("[MCP] Da ket noi thanh cong voi Server Xiaozhi");
        registerMcpTools();
    } else {
        Serial.println("[MCP] Mat ket noi Server MCP");
    }
}

// ================================================================
// 4. SETUP VA LOOP
// ================================================================

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

    for (int i = 0; i < deviceCount; i++) {
        pinMode(allPins[i], OUTPUT);
        digitalWrite(allPins[i], LOW); 
    }

    WiFi.begin(ssid, password);
    Serial.print("Connecting to WiFi");
    while (WiFi.status() != WL_CONNECTED) {
        delay(500);
        Serial.print(".");
    }
    Serial.println("\n✅ WiFi Connected!");

    mcpClient.begin(mcpEndpoint, onConnectionStatus);
}

void loop() {
    if (WiFi.status() != WL_CONNECTED) {
        if (!isFailsafeTriggered) {
            shutDownAll();
            isFailsafeTriggered = true;
        }
        static unsigned long lastRetry = 0;
        if (millis() - lastRetry > 5000) {
            WiFi.reconnect();
            lastRetry = millis();
            Serial.println("[WIFI] Dang thu ket noi lai...");
        }
    } else {
        if (isFailsafeTriggered) {
            Serial.println("[WIFI] Co mang tro lai, he thong hoat dong binh thuong.");
            isFailsafeTriggered = false;
        }
        mcpClient.loop();
    }
    delay(10);
}