#include <Arduino.h>
#include <NimBLEDevice.h>
#include <FastLED.h>

#include <mbedtls/md.h>
#include <mbedtls/sha256.h>
#include <Preferences.h>

#include "version.h"

#define TOUCH1_PIN  8
#define TOUCH2_PIN  6
#define BUTTON_PIN  41
#define LED_PIN     35
#define NUM_LEDS    1

CRGB leds[NUM_LEDS];

#define SERVICE_UUID        "4fafc201-1fb5-459e-8fcc-c5c9c331914b"
#define CHARACTERISTIC_UUID "beb5483e-36e1-4688-b7f5-ea07361f26d8"
#define UUID_CHAR_UUID      "beb5483e-36e1-4688-b7f5-ea07361f26d9"
#define HMAC_PKT_SIZE 37
#define HMAC_KEY_LEN  32
#define ESP_UUID_LEN  16
#define NVS_NAMESPACE "hmac"
#define NVS_KEY_NAME  "key"
#define NVS_UUID_NAME "uuid"
#define NVS_SERVER_NAME "server"
#define DEBOUNCE_MS   50
#define SCAN_INTERVAL 6000
#define SCAN_TIMEOUT  5000

uint8_t hmacKey[HMAC_KEY_LEN];
uint8_t deviceUuid[ESP_UUID_LEN];
uint8_t serverUuid[ESP_UUID_LEN];
bool keyProvisioned = false;
bool uuidProvisioned = false;
bool serverProvisioned = false;
bool allProvisioned = false;
Preferences preferences;

static bool connected = false;
NimBLEClient *pClient = nullptr;
NimBLERemoteCharacteristic *pCmdChar = nullptr;

static bool relay1SendState = false;
static bool relay2SendState = false;
static uint32_t seq = 0;

#define STABLE_LED_TIMEOUT_MS 60000
uint32_t connectionStableMs = 0;
uint32_t relayFlashUntil = 0;
bool otaInProgress = false;
bool otaEnded = false;

static bool bleReady = false;
static bool serverFound = false;
static NimBLEAddress serverAddress;

class MyScanCallbacks : public NimBLEScanCallbacks {
    void onResult(const NimBLEAdvertisedDevice *advertisedDevice) override {
        if (advertisedDevice->haveName() && advertisedDevice->getName() == "SylDa_Server") {
            serverAddress = advertisedDevice->getAddress();
            serverFound = true;
            NimBLEDevice::getScan()->stop();
        }
    }
};

class MyClientCallbacks : public NimBLEClientCallbacks {
    void onConnect(NimBLEClient *client) override {
        connected = true;
        Serial.println("BLE: connected to server");
    }
    void onDisconnect(NimBLEClient *client, int reason) override {
        connected = false;
        relay1SendState = false;
        relay2SendState = false;
        pCmdChar = nullptr;
        Serial.printf("BLE: disconnected (reason=%d)\n", reason);
    }
};

static MyClientCallbacks clientCallbacks;

static bool isHexKey(const String &hex) {
    for (int i = 0; i < hex.length(); i++) {
        char c = hex.charAt(i);
        if (!((c >= '0' && c <= '9') || (c >= 'A' && c <= 'F'))) return false;
    }
    return true;
}

static void parseHexKey(const String &hex, uint8_t *out, int len) {
    for (int i = 0; i < len; i++) {
        out[i] = (uint8_t)strtoul(hex.substring(i * 2, i * 2 + 2).c_str(), nullptr, 16);
    }
}

static void buildHmacPacket(uint8_t cmd, uint8_t *packet) {
    uint32_t s = ++seq;
    packet[0] = cmd;
    packet[1] = (s >> 24) & 0xFF;
    packet[2] = (s >> 16) & 0xFF;
    packet[3] = (s >> 8)  & 0xFF;
    packet[4] = s & 0xFF;
    uint8_t hmacData[5];
    hmacData[0] = cmd;
    memcpy(hmacData + 1, packet + 1, 4);
    const mbedtls_md_info_t *mdInfo = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
    mbedtls_md_hmac(mdInfo, hmacKey, HMAC_KEY_LEN, hmacData, 5, packet + 5);
}

static void disconnectClient() {
    if (pClient != nullptr) {
        pClient->disconnect();
        delay(100);
        NimBLEDevice::deleteClient(pClient);
        pClient = nullptr;
    }
    pCmdChar = nullptr;
    connected = false;
    relay1SendState = false;
    relay2SendState = false;
}

static bool connectToServer() {
    pClient = NimBLEDevice::createClient();
    pClient->setClientCallbacks(&clientCallbacks);

    if (!pClient->connect(serverAddress)) {
        Serial.println("BLE: connect failed");
        NimBLEDevice::deleteClient(pClient);
        pClient = nullptr;
        return false;
    }

    NimBLERemoteService *pSvc = pClient->getService(SERVICE_UUID);
    if (pSvc == nullptr) {
        Serial.println("BLE: service not found");
        disconnectClient();
        return false;
    }

    NimBLERemoteCharacteristic *pUuidChar = pSvc->getCharacteristic(UUID_CHAR_UUID);
    if (pUuidChar == nullptr) {
        Serial.println("BLE: UUID characteristic not found");
        disconnectClient();
        return false;
    }

    std::string uuidVal = pUuidChar->readValue();
    if (uuidVal.length() != ESP_UUID_LEN ||
        memcmp(uuidVal.data(), serverUuid, ESP_UUID_LEN) != 0) {
        Serial.println("BLE: server UUID mismatch — disconnecting");
        disconnectClient();
        return false;
    }

    pCmdChar = pSvc->getCharacteristic(CHARACTERISTIC_UUID);
    if (pCmdChar == nullptr) {
        Serial.println("BLE: command characteristic not found");
        disconnectClient();
        return false;
    }

    seq = 0;
    relay1SendState = false;
    relay2SendState = false;
    Serial.println("BLE: server UUID verified, ready");
    return true;
}

void setup() {
    Serial.begin(115200);
    delay(2000);
    Serial.println("=== SylDa Client ===");

    FastLED.addLeds<WS2812, LED_PIN, GRB>(leds, NUM_LEDS);
    leds[0] = CRGB::Red;
    FastLED.show();

    preferences.begin(NVS_NAMESPACE, true);
    size_t len = preferences.getBytes(NVS_KEY_NAME, hmacKey, HMAC_KEY_LEN);
    size_t uuidLen = preferences.getBytes(NVS_UUID_NAME, deviceUuid, ESP_UUID_LEN);
    size_t serverLen = preferences.getBytes(NVS_SERVER_NAME, serverUuid, ESP_UUID_LEN);
    preferences.end();

    if (len == HMAC_KEY_LEN) {
        keyProvisioned = true;
        Serial.println("HMAC key loaded from NVS");
    } else {
        keyProvisioned = false;
        Serial.println("No HMAC key — use USB SET_KEY to provision");
    }

    if (uuidLen == ESP_UUID_LEN) {
        uuidProvisioned = true;
        Serial.print("UUID loaded from NVS: ");
        for (int i = 0; i < ESP_UUID_LEN; i++) Serial.printf("%02X", deviceUuid[i]);
        Serial.println();
    } else {
        uuidProvisioned = false;
        for (int i = 0; i < ESP_UUID_LEN; i++) deviceUuid[i] = (uint8_t)(esp_random() & 0xFF);
        preferences.begin(NVS_NAMESPACE, false);
        preferences.putBytes(NVS_UUID_NAME, deviceUuid, ESP_UUID_LEN);
        preferences.end();
        uuidProvisioned = true;
        Serial.print("Generated new UUID: ");
        for (int i = 0; i < ESP_UUID_LEN; i++) Serial.printf("%02X", deviceUuid[i]);
        Serial.println();
    }

    if (serverLen == ESP_UUID_LEN) {
        serverProvisioned = true;
        Serial.print("Server UUID loaded from NVS: ");
        for (int i = 0; i < ESP_UUID_LEN; i++) Serial.printf("%02X", serverUuid[i]);
        Serial.println();
    } else {
        serverProvisioned = false;
        Serial.println("No server UUID — use USB SET_SERVER to provision");
    }

    allProvisioned = keyProvisioned && uuidProvisioned && serverProvisioned;

    pinMode(TOUCH1_PIN, INPUT);
    pinMode(TOUCH2_PIN, INPUT);
    pinMode(BUTTON_PIN, INPUT);

    if (allProvisioned) {
        NimBLEDevice::init("SylDa_Client");

        NimBLEScan *pScan = NimBLEDevice::getScan();
        pScan->setScanCallbacks(new MyScanCallbacks());
        pScan->setInterval(97);
        pScan->setWindow(37);
        pScan->setActiveScan(false);
        pScan->setMaxResults(0);
        bleReady = true;
        Serial.println("BLE: initialized, will scan for SylDa_Server");
    }

    Serial.println("Setup done.");
}

void loop() {
    if (Serial.available()) {
        static String serialBuf;
        char c = Serial.read();
        serialBuf += c;
        if (c == '\n') {
            serialBuf.trim();
            if (serialBuf == "KEY?") {
                if (keyProvisioned && uuidProvisioned) {
                    Serial.print("DEVICE_ID:");
                    for (int i = 0; i < ESP_UUID_LEN; i++) Serial.printf("%02X", deviceUuid[i]);
                    Serial.println();
                    Serial.print("HMAC_KEY:");
                    for (int i = 0; i < HMAC_KEY_LEN; i++) Serial.printf("%02X", hmacKey[i]);
                    Serial.println();
                    Serial.print("SERVER_UUID:");
                    if (serverProvisioned) {
                        for (int i = 0; i < ESP_UUID_LEN; i++) Serial.printf("%02X", serverUuid[i]);
                    }
                    Serial.println();
                } else {
                    Serial.println("NO_KEY");
                }
            } else if (serialBuf == "VERSION?") {
                Serial.print("BUILD_NR:");
                Serial.println(FIRMWARE_VERSION_CLIENT);
            } else if (serialBuf.startsWith("SET_KEY:")) {
                String hex = serialBuf.substring(8);
                hex.toUpperCase();
                if (hex.length() == 64 && isHexKey(hex)) {
                    parseHexKey(hex, hmacKey, HMAC_KEY_LEN);
                    preferences.begin(NVS_NAMESPACE, false);
                    preferences.putBytes(NVS_KEY_NAME, hmacKey, HMAC_KEY_LEN);
                    preferences.end();
                    keyProvisioned = true;
                    allProvisioned = keyProvisioned && uuidProvisioned && serverProvisioned;
                    Serial.println("OK");
                    Serial.println("HMAC key stored in NVS");
                    if (allProvisioned && !bleReady) {
                        NimBLEDevice::init("SylDa_Client");
                        NimBLEScan *pScan = NimBLEDevice::getScan();
                        pScan->setScanCallbacks(new MyScanCallbacks());
                        pScan->setInterval(97);
                        pScan->setWindow(37);
                        pScan->setActiveScan(false);
                        pScan->setMaxResults(0);
                        bleReady = true;
                        Serial.println("BLE: initialized after provisioning");
                    }
                } else {
                    Serial.println("ERROR: invalid key format");
                }
            } else if (serialBuf.startsWith("SET_ID:")) {
                String hex = serialBuf.substring(7);
                hex.toUpperCase();
                if (hex.length() == 32 && isHexKey(hex)) {
                    parseHexKey(hex, deviceUuid, ESP_UUID_LEN);
                    preferences.begin(NVS_NAMESPACE, false);
                    preferences.putBytes(NVS_UUID_NAME, deviceUuid, ESP_UUID_LEN);
                    preferences.end();
                    uuidProvisioned = true;
                    allProvisioned = keyProvisioned && uuidProvisioned && serverProvisioned;
                    Serial.println("OK");
                    Serial.println("UUID stored in NVS");
                } else {
                    Serial.println("ERROR: invalid UUID format");
                }
            } else if (serialBuf.startsWith("SET_SERVER:")) {
                String hex = serialBuf.substring(11);
                hex.toUpperCase();
                if (hex.length() == 32 && isHexKey(hex)) {
                    parseHexKey(hex, serverUuid, ESP_UUID_LEN);
                    preferences.begin(NVS_NAMESPACE, false);
                    preferences.putBytes(NVS_SERVER_NAME, serverUuid, ESP_UUID_LEN);
                    preferences.end();
                    serverProvisioned = true;
                    allProvisioned = keyProvisioned && uuidProvisioned && serverProvisioned;
                    Serial.println("OK");
                    Serial.println("Server UUID stored in NVS");
                    if (allProvisioned && !bleReady) {
                        NimBLEDevice::init("SylDa_Client");
                        NimBLEScan *pScan = NimBLEDevice::getScan();
                        pScan->setScanCallbacks(new MyScanCallbacks());
                        pScan->setInterval(97);
                        pScan->setWindow(37);
                        pScan->setActiveScan(false);
                        pScan->setMaxResults(0);
                        bleReady = true;
                        Serial.println("BLE: initialized after provisioning");
                    }
                } else {
                    Serial.println("ERROR: invalid server UUID format");
                }
            }
            serialBuf = "";
        }
    }

    if (!allProvisioned) {
        delay(20);
        return;
    }

    static bool t1Stable = false;
    static bool t2Stable = false;
    static bool btnStable = false;
    static uint32_t t1Change = 0;
    static uint32_t t2Change = 0;
    static uint32_t btnChange = 0;
    bool t1Raw = (digitalRead(TOUCH1_PIN) == HIGH);
    bool t2Raw = (digitalRead(TOUCH2_PIN) == HIGH);
    bool btnRaw = (digitalRead(BUTTON_PIN) == LOW);
    if (t1Raw != t1Stable) { t1Change = millis(); t1Stable = t1Raw; }
    if (t2Raw != t2Stable) { t2Change = millis(); t2Stable = t2Raw; }
    if (btnRaw != btnStable) { btnChange = millis(); btnStable = btnRaw; }
    bool t1Active = (millis() - t1Change > DEBOUNCE_MS) && t1Stable;
    bool t2Active = (millis() - t2Change > DEBOUNCE_MS) && t2Stable;
    bool btnActive = (millis() - btnChange > DEBOUNCE_MS) && btnStable;

    static bool prevBtnActive = false;
    if (btnActive && !prevBtnActive) {
        Serial.printf("Button: pressed\n");
    }
    prevBtnActive = btnActive;

    bool relay1State = t1Active || btnActive;
    bool relay2State = t2Active;

    CRGB ledColor;

    if (otaInProgress) {
        ledColor = otaEnded
            ? CRGB::Yellow
            : (((millis() / 100) % 2) ? CRGB::Green : CRGB::Black);
    } else {
        if (connected) {
            if (connectionStableMs == 0) connectionStableMs = millis();
        } else {
            connectionStableMs = 0;
        }

        if (connected) {
            ledColor = (millis() - connectionStableMs >= STABLE_LED_TIMEOUT_MS)
                ? CRGB::Black
                : CRGB::Yellow;
        } else {
            ledColor = CRGB::Red;
        }

        static bool lastRelayActive = false;
        bool relayActive = relay1State || relay2State;
        if (relayActive != lastRelayActive) {
            relayFlashUntil = millis() + 150;
        }
        lastRelayActive = relayActive;

        if (millis() < relayFlashUntil) {
            ledColor = CRGB::Green;
        }
    }

    leds[0] = ledColor;
    FastLED.show();

    if (connected && pCmdChar != nullptr) {
        static uint32_t lastCmd = 0;
        if (millis() - lastCmd > 200) {
            if (relay1State != relay1SendState) {
                uint8_t pkt[HMAC_PKT_SIZE];
                buildHmacPacket(relay1State ? 0x01 : 0x00, pkt);
                pCmdChar->writeValue(pkt, HMAC_PKT_SIZE, false);
                relay1SendState = relay1State;
                Serial.printf("TX: relay1 %s (seq=%lu)\n", relay1State ? "ON" : "OFF", seq);
                lastCmd = millis();
            }
            if (relay2State != relay2SendState) {
                uint8_t pkt[HMAC_PKT_SIZE];
                buildHmacPacket(relay2State ? 0x02 : 0x03, pkt);
                pCmdChar->writeValue(pkt, HMAC_PKT_SIZE, false);
                relay2SendState = relay2State;
                Serial.printf("TX: relay2 %s (seq=%lu)\n", relay2State ? "ON" : "OFF", seq);
                lastCmd = millis();
            }
        }
    }

    static uint32_t lastScanStart = 0;
    static bool scanRunnable = false;

    if (!connected) {
        NimBLEScan *pScan = NimBLEDevice::getScan();
        if ((millis() - lastScanStart > SCAN_INTERVAL) && !pScan->isScanning()) {
            serverFound = false;
            pScan->start(SCAN_TIMEOUT / 1000, false);
            lastScanStart = millis();
            scanRunnable = true;
        }

        if (!pScan->isScanning() && scanRunnable) {
            scanRunnable = false;
            if (serverFound) {
                Serial.printf("BLE: found SylDa_Server (%s), connecting...\n",
                              serverAddress.toString().c_str());
                connectToServer();
            }
        }
    }

    if (connected && pClient != nullptr && !pClient->isConnected()) {
        Serial.println("BLE: lost connection");
        disconnectClient();
    }

    static uint32_t heartbeat = 0;
    if (millis() - heartbeat > 60000) {
        Serial.printf("[heartbeat] t1=%s t2=%s btn=%s r1=%s r2=%s connected=%s\n",
            t1Active ? "YES" : "no",
            t2Active ? "YES" : "no",
            btnActive ? "YES" : "no",
            relay1State ? "ON" : "OFF",
            relay2State ? "ON" : "OFF",
            connected ? "yes" : "no");
        heartbeat = millis();
    }

    delay(20);
}