// SPDX-License-Identifier: GPL-3.0-or-later
#include <Arduino.h>
#include <NimBLEDevice.h>
#include <FastLED.h>

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

#include <esp_log.h>
#include "version.h"
#include "ecdh_helper.h"

static const char *TAG = "client";

static bool serialMuted = false;

static void muteSerial() {
    serialMuted = true;
    esp_log_level_set("*", ESP_LOG_NONE);
}

static void unmuteSerial() {
    serialMuted = false;
    esp_log_level_set("*", (esp_log_level_t)CONFIG_LOG_DEFAULT_LEVEL);
}

#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 ECDH_UUID           "beb5483e-36e1-4688-b7f5-ea07361f26dc"
#define ESP_UUID_LEN  16
#define NVS_NAMESPACE "syl"
#define NVS_PSK_NAME  "psk"
#define NVS_UUID_NAME "uuid"
#define NVS_SERVER_NAME "server"
#define NVS_ADDR_NAME   "addr"
#define DEBOUNCE_MS   50
#define SCAN_DURATION_MS   5000
#define SCAN_INTERVAL_MS   8000
#define RECONNECT_DELAY_MS 1000

uint8_t psk[PSK_LEN];
uint8_t deviceUuid[ESP_UUID_LEN];
uint8_t serverUuid[ESP_UUID_LEN];
bool pskProvisioned = false;
bool uuidProvisioned = false;
bool serverProvisioned = false;
bool allProvisioned = false;
bool serverAddrStored = false;
NimBLEAddress storedServerAddr;
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;

static uint8_t aesKey[AES_KEY_LEN];
static bool ecdhReady = false;

#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;
static NimBLEAddress foundAddr;
static bool scanning = false;
static uint32_t scanStarted = 0;
static uint32_t disconnectTime = 0;

static bool ecdhNotifyReceived = false;
static uint8_t ecdhNotifyBuf[ECDH_PKT_SIZE];
static size_t ecdhNotifyLen = 0;

class MyScanCallbacks : public NimBLEScanCallbacks {
    void onResult(const NimBLEAdvertisedDevice *advertisedDevice) override {
        if (scanning && advertisedDevice->isAdvertisingService(NimBLEUUID(SERVICE_UUID))) {
            std::string mfgData = advertisedDevice->getManufacturerData();
            if (mfgData.length() >= 6) {
                if (memcmp(mfgData.data() + 2, serverUuid, 4) == 0) {
                    foundAddr = advertisedDevice->getAddress();
                    serverFound = true;
                    ESP_LOGI(TAG, "BLE: server found at %s", foundAddr.toString().c_str());
                }
            }
        }
    }
    void onScanEnd(const NimBLEScanResults &results, int reason) override {
        scanning = false;
        ESP_LOGI(TAG, "BLE: scan ended (reason=%d, found=%d)", reason, serverFound);
    }
};

class MyClientCallbacks : public NimBLEClientCallbacks {
    void onConnect(NimBLEClient *client) override {
        connected = true;
        ESP_LOGI(TAG, "BLE: connected to server");
    }
    void onDisconnect(NimBLEClient *client, int reason) override {
        connected = false;
        disconnectTime = millis();
        relay1SendState = false;
        relay2SendState = false;
        pCmdChar = nullptr;
        ecdhReady = false;
        memset(aesKey, 0, AES_KEY_LEN);
        ESP_LOGW(TAG, "BLE: disconnected (reason=%d)", reason);
    }
};

static MyClientCallbacks clientCallbacks;

static bool isHex(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 parseHex(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 buildAesPacket(uint8_t cmd, uint8_t *packet) {
    uint32_t s = ++seq;
    packet[0] = (s >> 24) & 0xFF;
    packet[1] = (s >> 16) & 0xFF;
    packet[2] = (s >> 8)  & 0xFF;
    packet[3] = s & 0xFF;

    uint8_t nonceRandom[8];
    for (int i = 0; i < 8; i++) nonceRandom[i] = (uint8_t)(esp_random() & 0xFF);
    memcpy(packet + 4, nonceRandom, 8);

    uint8_t nonce[12];
    memcpy(nonce, packet, 4);
    memcpy(nonce + 4, nonceRandom, 8);

    mbedtls_gcm_context gcm;
    mbedtls_gcm_init(&gcm);
    int ret = mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, aesKey, AES_KEY_LEN * 8);
    if (ret != 0) {
        ESP_LOGE(TAG, "AES: setkey failed (ret=%d)", ret);
        mbedtls_gcm_free(&gcm);
        return;
    }

    uint8_t pt[1] = { cmd };
    mbedtls_gcm_crypt_and_tag(&gcm, MBEDTLS_GCM_ENCRYPT,
        1, nonce, 12, NULL, 0,
        pt, packet + 12, 16, packet + 13);
    mbedtls_gcm_free(&gcm);
}

static void disconnectClient() {
    if (pClient != nullptr) {
        if (pClient->isConnected()) {
            pClient->disconnect();
            delay(100);
        }
    }
    pCmdChar = nullptr;
    connected = false;
    relay1SendState = false;
    relay2SendState = false;
    ecdhReady = false;
    memset(aesKey, 0, AES_KEY_LEN);
}

static bool performEcdhHandshake(NimBLERemoteService *pSvc) {
    if (!pskProvisioned) {
        ESP_LOGE(TAG, "ECDH: no PSK");
        return false;
    }

    CryptoECDH ecdh;
    if (!ecdh.init()) {
        ESP_LOGE(TAG, "ECDH: init failed");
        return false;
    }
    if (!ecdh.generateKeys()) {
        ESP_LOGE(TAG, "ECDH: keygen failed");
        ecdh.free();
        return false;
    }

    NimBLERemoteCharacteristic *pEcdhChar = pSvc->getCharacteristic(ECDH_UUID);
    if (pEcdhChar == nullptr) {
        ESP_LOGW(TAG, "ECDH: characteristic not found");
        ecdh.free();
        return false;
    }

    if (!pEcdhChar->subscribe(true, [](NimBLERemoteCharacteristic *pChar, const uint8_t *data, size_t len, bool isNotify) {
            if (ecdhReady) return;
            if (len != ECDH_PKT_SIZE) {
                ESP_LOGW(TAG, "ECDH: bad notify length %d", len);
                return;
            }
            memcpy(ecdhNotifyBuf, data, len);
            ecdhNotifyLen = len;
            ecdhNotifyReceived = true;
            ESP_LOGI(TAG, "ECDH: notify received (%d bytes)", len);
        }, true)) {
        ESP_LOGW(TAG, "ECDH: subscribe failed");
        ecdh.free();
        return false;
    }

    uint8_t handshake[ECDH_PKT_SIZE];
    memcpy(handshake, ecdh.public_key, ECDH_PUBKEY_LEN);
    CryptoECDH::signPubkey(psk, ecdh.public_key, handshake + ECDH_PUBKEY_LEN);

    ecdhNotifyReceived = false;
    if (!pEcdhChar->writeValue(handshake, ECDH_PKT_SIZE, false)) {
        ESP_LOGW(TAG, "ECDH: write pubkey failed");
        ecdh.free();
        return false;
    }

    uint32_t start = millis();
    while (!ecdhNotifyReceived && (millis() - start < 5000)) {
        delay(10);
    }

    if (!ecdhNotifyReceived) {
        ESP_LOGW(TAG, "ECDH: notify timeout");
        ecdh.free();
        return false;
    }

    if (!CryptoECDH::verifyPubkey(psk, ecdhNotifyBuf, ecdhNotifyBuf + ECDH_PUBKEY_LEN)) {
        ESP_LOGW(TAG, "ECDH: server pubkey HMAC mismatch");
        ecdh.free();
        return false;
    }

    if (!ecdh.computeSharedSecret(ecdhNotifyBuf, ECDH_PUBKEY_LEN)) {
        ESP_LOGE(TAG, "ECDH: shared secret failed");
        ecdh.free();
        return false;
    }

    memcpy(aesKey, ecdh.aes_key, AES_KEY_LEN);
    ecdh.free();
    ecdhReady = true;
    ESP_LOGI(TAG, "ECDH: handshake complete");
    return true;
}

static bool connectToServer() {
    if (pClient == nullptr) {
        pClient = NimBLEDevice::createClient();
        pClient->setConnectTimeout(5000);
        pClient->setClientCallbacks(&clientCallbacks);
    }

    if (!pClient->connect(serverAddress)) {
        ESP_LOGE(TAG, "BLE: connect failed, retry later");
        return false;
    }

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

    NimBLERemoteCharacteristic *pUuidChar = pSvc->getCharacteristic(UUID_CHAR_UUID);
    if (pUuidChar == nullptr) {
        ESP_LOGW(TAG, "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) {
        ESP_LOGW(TAG, "BLE: server UUID mismatch — disconnecting");
        disconnectClient();
        return false;
    }

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

    ecdhReady = false;
    memset(aesKey, 0, AES_KEY_LEN);

    if (!performEcdhHandshake(pSvc)) {
        ESP_LOGE(TAG, "BLE: ECDH handshake failed");
        disconnectClient();
        return false;
    }

    seq = 0;
    relay1SendState = false;
    relay2SendState = false;
    ESP_LOGI(TAG, "BLE: server UUID verified, ECDH done, ready");
    return true;
}

void setup() {
    Serial.begin(115200);
    delay(2000);
    ESP_LOGI(TAG, "=== SylDa Client v%s ===", FIRMWARE_VERSION_CLIENT);

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

    preferences.begin(NVS_NAMESPACE, true);
    size_t pskLen = preferences.getBytes(NVS_PSK_NAME, psk, PSK_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 (pskLen == PSK_LEN) {
        pskProvisioned = true;
        ESP_LOGI(TAG, "PSK loaded from NVS");
    } else {
        pskProvisioned = false;
        ESP_LOGW(TAG, "No PSK — use USB SET_KEY to provision");
    }

    if (uuidLen == ESP_UUID_LEN) {
        uuidProvisioned = true;
        ESP_LOGD(TAG, "UUID loaded from NVS");
    } 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;
        ESP_LOGD(TAG, "Generated new UUID");
    }

    if (serverLen == ESP_UUID_LEN) {
        serverProvisioned = true;
        ESP_LOGD(TAG, "Server UUID loaded from NVS");
    } else {
        serverProvisioned = false;
        ESP_LOGW(TAG, "No server UUID — use USB SET_SERVER to provision");
    }

    allProvisioned = pskProvisioned && uuidProvisioned && serverProvisioned;

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

    if (allProvisioned) {
        NimBLEDevice::init("SylDa_Client");
        bleReady = true;

        preferences.begin(NVS_NAMESPACE, true);
        uint8_t addrBuf[6];
        size_t addrLen = preferences.getBytes(NVS_ADDR_NAME, addrBuf, 6);
        preferences.end();
        if (addrLen == 6) {
            ble_addr_t addr;
            addr.type = BLE_ADDR_PUBLIC;
            memcpy(addr.val, addrBuf, 6);
            storedServerAddr = NimBLEAddress(addr);
            serverAddrStored = true;
            ESP_LOGI(TAG, "BLE: stored server address %s", storedServerAddr.toString().c_str());
        }

        ESP_LOGI(TAG, "BLE: initialized, will scan for server");
    }

    ESP_LOGI(TAG, "Setup done.");
}

static void startScan() {
    if (scanning) return;
    NimBLEScan *pScan = NimBLEDevice::getScan();
    pScan->setScanCallbacks(new MyScanCallbacks());
    pScan->setActiveScan(true);
    pScan->setInterval(32);
    pScan->setWindow(16);
    serverFound = false;
    scanning = true;
    scanStarted = millis();
    pScan->start(SCAN_DURATION_MS / 100, false);
    ESP_LOGI(TAG, "BLE: scanning started");
}

void loop() {
    if (Serial.available()) {
        static String serialBuf;
        char c = Serial.read();
        serialBuf += c;
        if (c == '\n') {
            serialBuf.trim();
            if (serialBuf == "KEY?") {
                if (pskProvisioned && uuidProvisioned) {
                    Serial.print("DEVICE_ID:");
                    for (int i = 0; i < ESP_UUID_LEN; i++) Serial.printf("%02X", deviceUuid[i]);
                    Serial.println();
                    Serial.print("PSK:");
                    for (int i = 0; i < PSK_LEN; i++) Serial.printf("%02X", psk[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 == "ID?") {
                Serial.flush();
                muteSerial();
                Serial.println("SylDa_Client");
                delay(50);
                unmuteSerial();
            } 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 && isHex(hex)) {
                    parseHex(hex, psk, PSK_LEN);
                    preferences.begin(NVS_NAMESPACE, false);
                    preferences.putBytes(NVS_PSK_NAME, psk, PSK_LEN);
                    preferences.end();
                    pskProvisioned = true;
                    allProvisioned = pskProvisioned && uuidProvisioned && serverProvisioned;
                    Serial.println("OK");
                    Serial.println("PSK stored in NVS");
                    if (allProvisioned && !bleReady) {
                        NimBLEDevice::init("SylDa_Client");
                        bleReady = true;
                        ESP_LOGI(TAG, "BLE: initialized after key provisioning");
                    }
                } else {
                    ESP_LOGE(TAG, "Invalid key format");
                    Serial.println("ERROR");
                }
            } else if (serialBuf.startsWith("SET_ID:")) {
                String hex = serialBuf.substring(7);
                hex.toUpperCase();
                if (hex.length() == 32 && isHex(hex)) {
                    parseHex(hex, deviceUuid, ESP_UUID_LEN);
                    preferences.begin(NVS_NAMESPACE, false);
                    preferences.putBytes(NVS_UUID_NAME, deviceUuid, ESP_UUID_LEN);
                    preferences.end();
                    uuidProvisioned = true;
                    allProvisioned = pskProvisioned && uuidProvisioned && serverProvisioned;
                    Serial.println("OK");
                    Serial.println("UUID stored in NVS");
                } else {
                    ESP_LOGE(TAG, "Invalid UUID format");
                    Serial.println("ERROR");
                }
            } else if (serialBuf.startsWith("SET_SERVER:")) {
                String hex = serialBuf.substring(11);
                hex.toUpperCase();
                if (hex.length() == 32 && isHex(hex)) {
                    parseHex(hex, serverUuid, ESP_UUID_LEN);
                    preferences.begin(NVS_NAMESPACE, false);
                    preferences.putBytes(NVS_SERVER_NAME, serverUuid, ESP_UUID_LEN);
                    preferences.end();
                    serverProvisioned = true;
                    allProvisioned = pskProvisioned && uuidProvisioned && serverProvisioned;
                    Serial.println("OK");
                    Serial.println("Server UUID stored in NVS");
                    if (allProvisioned && !bleReady) {
                        NimBLEDevice::init("SylDa_Client");
                        bleReady = true;
                        ESP_LOGI(TAG, "BLE: initialized after server provisioning");
                    }
                } else {
                    ESP_LOGE(TAG, "Invalid server UUID format");
                    Serial.println("ERROR");
                }
            }
            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;

    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 && ecdhReady) {
        static uint32_t lastCmd = 0;
        if (millis() - lastCmd > 200) {
            if (relay1State != relay1SendState) {
                uint8_t pkt[AES_PKT_SIZE];
                buildAesPacket(relay1State ? 0x01 : 0x00, pkt);
                pCmdChar->writeValue(pkt, AES_PKT_SIZE, false);
                relay1SendState = relay1State;
                ESP_LOGD(TAG, "TX: relay1 %s (seq=%lu)", relay1State ? "ON" : "OFF", seq);
                lastCmd = millis();
            }
            if (relay2State != relay2SendState) {
                uint8_t pkt[AES_PKT_SIZE];
                buildAesPacket(relay2State ? 0x02 : 0x03, pkt);
                pCmdChar->writeValue(pkt, AES_PKT_SIZE, false);
                relay2SendState = relay2State;
                ESP_LOGD(TAG, "TX: relay2 %s (seq=%lu)", relay2State ? "ON" : "OFF", seq);
                lastCmd = millis();
            }
        }
    }

    static uint32_t lastScanAttempt = 0;
    static uint32_t statusLog = 0;

    if (bleReady && millis() - statusLog > 2000) {
        ESP_LOGI(TAG, "BLE: connected=%d scanning=%d ecdh=%d serverFound=%d",
                 connected, scanning, ecdhReady, serverFound);
        statusLog = millis();
    }

    if (!connected && bleReady) {
        if (serverFound && !scanning) {
            serverFound = false;
            serverAddress = foundAddr;
            ESP_LOGI(TAG, "BLE: connecting to %s type=%d", serverAddress.toString().c_str(), serverAddress.getType());
            if (connectToServer()) {
                connected = true;
                uint8_t addrBuf[6];
                memcpy(addrBuf, serverAddress.getVal(), 6);
                preferences.begin(NVS_NAMESPACE, false);
                preferences.putBytes(NVS_ADDR_NAME, addrBuf, 6);
                preferences.end();
                storedServerAddr = serverAddress;
                serverAddrStored = true;
                ESP_LOGI(TAG, "BLE: server address stored in NVS");
            } else {
                ESP_LOGW(TAG, "BLE: connectToServer failed");
            }
        } else if (!scanning) {
            if (millis() - lastScanAttempt > SCAN_INTERVAL_MS) {
                if (serverAddrStored) {
                    ESP_LOGI(TAG, "BLE: trying stored address %s", storedServerAddr.toString().c_str());
                    serverAddress = storedServerAddr;
                    if (connectToServer()) {
                        connected = true;
                        ESP_LOGI(TAG, "BLE: connected via stored address");
                    } else {
                        ESP_LOGI(TAG, "BLE: stored address failed, starting scan");
                        serverAddrStored = false;
                        startScan();
                    }
                } else {
                    startScan();
                }
                lastScanAttempt = millis();
            }
        } else {
            if (millis() - scanStarted > SCAN_DURATION_MS) {
                NimBLEDevice::getScan()->stop();
                scanning = false;
                ESP_LOGI(TAG, "BLE: scan timeout, no server found");
            }
        }
    }

    if (connected && pClient != nullptr && !pClient->isConnected()) {
        ESP_LOGW(TAG, "BLE: lost connection");
        disconnectClient();
    }

    static uint32_t heartbeat = 0;
    if (millis() - heartbeat > 60000) {
        ESP_LOGD(TAG, "[hb] t1=%s t2=%s btn=%s r1=%s r2=%s conn=%s ecdh=%s",
            !t1Active ? "YES" : "no",
            !t2Active ? "YES" : "no",
            btnActive ? "YES" : "no",
            relay1State ? "ON" : "OFF",
            relay2State ? "ON" : "OFF",
            connected ? "yes" : "no",
            ecdhReady ? "yes" : "no");
        heartbeat = millis();
    }

    delay(20);
}