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

#include <mbedtls/gcm.h>
#include <mbedtls/sha256.h>
#include <Preferences.h>
#include <esp_ota_ops.h>
#include <esp_partition.h>
#include <nvs_flash.h>
#include <esp_log.h>
#include "version.h"
#include "ecdh_helper.h"

static const char *TAG = "server";

#define RELAY1_PIN  1
#define RELAY2_PIN  2
#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 MAX_CLIENTS 2

#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 VERSION_UUID        "beb5483e-36e1-4688-b7f5-ea07361f26db"
#define OTA_UUID            "beb5483e-36e1-4688-b7f5-ea07361f26da"
#define ECDH_UUID           "beb5483e-36e1-4688-b7f5-ea07361f26dc"

#define OTA_CHUNK_SIZE 500
#define OTA_CMD_START  0x10
#define OTA_CMD_END    0x11
#define OTA_STATUS_ACK 0x00
#define OTA_STATUS_NAK 0x01
#define OTA_STATUS_COMPLETE 0x10
#define OTA_STATUS_ERROR    0x11
#define ESP_UUID_LEN  16
#define NVS_NAMESPACE "syl"
#define NVS_PSK_NAME  "psk"
#define NVS_UUID_NAME "uuid"
#define FIRMWARE_VERSION FIRMWARE_VERSION_SERVER

NimBLECharacteristic *pCharacteristic;
NimBLECharacteristic *pUuidChar;
NimBLECharacteristic *pEcdhChar;
bool relay1State = false;
bool relay2State = false;

bool connActive[MAX_CLIENTS] = {false, false};
bool bleRelay1[MAX_CLIENTS] = {false, false};
bool bleRelay2[MAX_CLIENTS] = {false, false};
uint32_t lastSeq[MAX_CLIENTS] = {0, 0};
uint16_t connHandles[MAX_CLIENTS] = {0xFFFF, 0xFFFF};
int connectedCount = 0;

struct ClientState {
    CryptoECDH ecdh;
    uint8_t aesKey[AES_KEY_LEN];
    bool ecdhReady;
};
ClientState clients[MAX_CLIENTS];

uint8_t psk[PSK_LEN];
uint8_t deviceUuid[ESP_UUID_LEN];
bool pskProvisioned = false;
bool uuidProvisioned = false;
Preferences preferences;

#define STABLE_LED_TIMEOUT_MS 60000
uint32_t connectionStableMs = 0;
uint32_t relayFlashUntil = 0;

bool otaInProgress = false;
esp_ota_handle_t otaHandle = 0;
const esp_partition_t *otaPartition = nullptr;
uint32_t otaExpectedSize = 0;
uint32_t otaWritten = 0;
uint8_t otaExpectedSha[32];
mbedtls_sha256_context otaShaCtx;
bool otaShaStarted = false;
bool otaReady = false;
bool otaEnded = false;
bool otaConfirmed = false;
bool otaEarlyConfirm = false;
uint32_t otaEndedAt = 0;
#define OTA_CONFIRM_TIMEOUT_MS 10000

static int findClientIdx(uint16_t connHandle) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (connActive[i] && connHandles[i] == connHandle) return i;
    }
    return -1;
}

class MyServerCallbacks : public NimBLEServerCallbacks {
    void onConnect(NimBLEServer *pServer, NimBLEConnInfo &connInfo) override {
        uint16_t handle = connInfo.getConnHandle();
        int slot = -1;
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (!connActive[i]) {
                slot = i;
                break;
            }
        }
        if (slot >= 0) {
            connHandles[slot] = handle;
            connActive[slot] = true;
            bleRelay1[slot] = false;
            bleRelay2[slot] = false;
            lastSeq[slot] = 0;
            clients[slot].ecdhReady = false;
            memset(clients[slot].aesKey, 0, AES_KEY_LEN);
            clients[slot].ecdh.free();
            connectedCount++;
            ESP_LOGI(TAG, "BLE: connected (slot=%d, clients=%d/%d, mac=%s)",
                     slot, connectedCount, MAX_CLIENTS,
                     connInfo.getAddress().toString().c_str());
        } else {
            ESP_LOGW(TAG, "BLE: connected but no free slot (clients=%d/%d)",
                     connectedCount, MAX_CLIENTS);
        }
        pServer->startAdvertising();
    }
    void onDisconnect(NimBLEServer *pServer, NimBLEConnInfo &connInfo, int reason) override {
        uint16_t handle = connInfo.getConnHandle();
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (connActive[i] && connHandles[i] == handle) {
                connActive[i] = false;
                bleRelay1[i] = false;
                bleRelay2[i] = false;
                lastSeq[i] = 0;
                clients[i].ecdhReady = false;
                memset(clients[i].aesKey, 0, AES_KEY_LEN);
                clients[i].ecdh.free();
                connHandles[i] = 0xFFFF;
                connectedCount--;
                ESP_LOGI(TAG, "BLE: disconnected (slot=%d, clients=%d/%d, mac=%s)",
                         i, connectedCount, MAX_CLIENTS,
                         connInfo.getAddress().toString().c_str());
                break;
            }
        }
    }
};

static bool decryptAesGcm(const uint8_t *packet, int idx, uint8_t *outCmd) {
    if (idx < 0 || idx >= MAX_CLIENTS || !connActive[idx]) return false;
    if (!clients[idx].ecdhReady) {
        ESP_LOGW(TAG, "AES: ECDH not done for slot=%d", idx);
        return false;
    }

    uint32_t seq = ((uint32_t)packet[0] << 24) |
                   ((uint32_t)packet[1] << 16) |
                   ((uint32_t)packet[2] << 8)  |
                   (uint32_t)packet[3];

    if (seq <= lastSeq[idx]) {
        ESP_LOGW(TAG, "Replay rejected (seq=%lu, last=%lu, slot=%d)", seq, lastSeq[idx], idx);
        return false;
    }

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

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

    ret = mbedtls_gcm_auth_decrypt(&gcm,
        1,                              // ciphertext length
        nonce, 12,                      // iv + iv_len
        NULL, 0,                        // aad + aad_len
        packet + 13, 16,                // tag + tag_len
        packet + 12,                    // ciphertext
        outCmd);                        // plaintext
    mbedtls_gcm_free(&gcm);

    if (ret != 0) {
        ESP_LOGE(TAG, "AES-GCM decrypt failed (slot=%d, ret=%d)", idx, ret);
        return false;
    }

    lastSeq[idx] = seq;
    return true;
}

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 otaReset() {
    otaInProgress = false;
    otaHandle = 0;
    otaPartition = nullptr;
    otaExpectedSize = 0;
    otaWritten = 0;
    otaReady = false;
    otaEnded = false;
    otaConfirmed = false;
    otaEarlyConfirm = false;
    if (otaShaStarted) {
        mbedtls_sha256_free(&otaShaCtx);
        otaShaStarted = false;
    }
}

static void sendOtaNotify(NimBLECharacteristic *otaChar, uint8_t status, uint32_t seq) {
    uint8_t buf[5];
    buf[0] = status;
    buf[1] = (seq >> 24) & 0xFF;
    buf[2] = (seq >> 16) & 0xFF;
    buf[3] = (seq >> 8) & 0xFF;
    buf[4] = seq & 0xFF;
    otaChar->setValue(buf, 5);
    bool ok = otaChar->notify();
    if (!ok) {
        ESP_LOGW(TAG, "OTA notify() failed (status=0x%02X seq=%lu)", status, seq);
    }
}

static void abortOta(NimBLECharacteristic *otaChar, const char *reason) {
    ESP_LOGE(TAG, "OTA: abort — %s", reason);
    if (otaHandle != 0) {
        esp_ota_end(otaHandle);
    }
    otaReset();
    sendOtaNotify(otaChar, OTA_STATUS_ERROR, 0);
}

class OtaCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
    void onWrite(NimBLECharacteristic *pChar, NimBLEConnInfo &connInfo) override {
        std::string value = pChar->getValue();
        if (value.empty()) return;
        const uint8_t *data = (const uint8_t *)value.data();
        size_t len = value.length();

        if (len == 5 && data[0] == OTA_CMD_START) {
            if (otaInProgress) {
                esp_ota_end(otaHandle);
                otaReset();
            }
            otaExpectedSize = ((uint32_t)data[1] << 24) | ((uint32_t)data[2] << 16) |
                              ((uint32_t)data[3] << 8)  | (uint32_t)data[4];
            otaPartition = esp_ota_get_next_update_partition(nullptr);
            if (otaPartition == nullptr) {
                ESP_LOGE(TAG, "OTA: no update partition found");
                sendOtaNotify(pChar, OTA_STATUS_ERROR, 0);
                return;
            }
            esp_err_t err = esp_ota_begin(otaPartition, otaExpectedSize, &otaHandle);
            if (err != ESP_OK) {
                ESP_LOGE(TAG, "OTA: begin failed (%d)", err);
                sendOtaNotify(pChar, OTA_STATUS_ERROR, 0);
                return;
            }
            otaInProgress = true;
            otaWritten = 0;
            mbedtls_sha256_init(&otaShaCtx);
            mbedtls_sha256_starts_ret(&otaShaCtx, 0);
            otaShaStarted = true;
            digitalWrite(RELAY1_PIN, LOW);
            digitalWrite(RELAY2_PIN, LOW);
            relay1State = false;
            relay2State = false;
            for (int i = 0; i < MAX_CLIENTS; i++) {
                bleRelay1[i] = false;
                bleRelay2[i] = false;
            }
            ESP_LOGI(TAG, "OTA: started, expecting %lu bytes", otaExpectedSize);
            return;
        }

        if (len == 32 && otaInProgress) {
            memcpy(otaExpectedSha, data, 32);
            uint8_t computedSha[32];
            mbedtls_sha256_finish_ret(&otaShaCtx, computedSha);
            mbedtls_sha256_free(&otaShaCtx);
            otaShaStarted = false;

            ESP_LOGD(TAG, "OTA: SHA256 verify (written=%lu expected=%lu)", otaWritten, otaExpectedSize);
            if (memcmp(computedSha, otaExpectedSha, 32) != 0) {
                ESP_LOGE(TAG, "OTA: SHA256 mismatch — aborting");
                esp_ota_abort(otaHandle);
                otaReset();
                sendOtaNotify(pChar, OTA_STATUS_ERROR, 0);
                return;
            }

            ESP_LOGI(TAG, "OTA: SHA256 OK, waiting for 0xFF confirm");
            otaReady = true;
            sendOtaNotify(pChar, OTA_STATUS_COMPLETE, 0);
            return;
        }

        if (!otaInProgress) {
            ESP_LOGD(TAG, "OTA: write ignored (not in progress, len=%d data0=0x%02X)", len, data[0]);
            return;
        }

        if (len == 1 && data[0] == 0xFF) {
            ESP_LOGD(TAG, "OTA: 0xFF received (ready=%d ended=%d confirmed=%d inProgress=%d)",
                     otaReady, otaEnded, otaConfirmed, otaInProgress);
            if (!otaEnded) {
                ESP_LOGD(TAG, "OTA: 0xFF arrived early, deferring");
                otaEarlyConfirm = true;
                return;
            }
            ESP_LOGI(TAG, "OTA: 0xFF confirm accepted");
            otaConfirmed = true;
            return;
        }

        if (data[0] < 0x10) {
            uint32_t seq = ((uint32_t)data[0] << 24) | ((uint32_t)data[1] << 16) |
                           ((uint32_t)data[2] << 8)  | (uint32_t)data[3];
            size_t dataLen = len - 4;
            mbedtls_sha256_update_ret(&otaShaCtx, data + 4, dataLen);
            esp_err_t err = esp_ota_write(otaHandle, data + 4, dataLen);
            if (err != ESP_OK) {
                ESP_LOGE(TAG, "OTA: write error at seq %lu (%d)", seq, err);
                sendOtaNotify(pChar, OTA_STATUS_NAK, seq);
                abortOta(pChar, "write error");
                return;
            }
            otaWritten += dataLen;
            sendOtaNotify(pChar, OTA_STATUS_ACK, seq);
        }
    }
};

class EcdhCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
    void onWrite(NimBLECharacteristic *pChar, NimBLEConnInfo &connInfo) override {
        if (otaInProgress) return;

        int idx = findClientIdx(connInfo.getConnHandle());
        if (idx < 0) return;

        if (!pskProvisioned) {
            ESP_LOGW(TAG, "ECDH: write rejected — no PSK (slot=%d)", idx);
            return;
        }

        std::string value = pChar->getValue();
        if (value.length() != ECDH_PKT_SIZE) {
            ESP_LOGW(TAG, "ECDH: bad packet length %d (expected %d)", (int)value.length(), ECDH_PKT_SIZE);
            return;
        }

        const uint8_t *pkt = (const uint8_t *)value.data();

        if (!CryptoECDH::verifyPubkey(psk, pkt, pkt + ECDH_PUBKEY_LEN)) {
            ESP_LOGW(TAG, "ECDH: pubkey HMAC mismatch (slot=%d)", idx);
            return;
        }

        uint8_t peerPubkey[ECDH_PUBKEY_LEN];
        memcpy(peerPubkey, pkt, ECDH_PUBKEY_LEN);

        clients[idx].ecdh.free();

        if (!clients[idx].ecdh.init() || !clients[idx].ecdh.generateKeys()) {
            ESP_LOGE(TAG, "ECDH: keygen failed (slot=%d)", idx);
            return;
        }

        if (!clients[idx].ecdh.computeSharedSecret(peerPubkey, ECDH_PUBKEY_LEN)) {
            ESP_LOGE(TAG, "ECDH: shared secret failed (slot=%d)", idx);
            clients[idx].ecdh.free();
            return;
        }

        memcpy(clients[idx].aesKey, clients[idx].ecdh.aes_key, AES_KEY_LEN);

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

        pChar->setValue(resp, ECDH_PKT_SIZE);
        if (!pChar->notify()) {
            ESP_LOGW(TAG, "ECDH: notify failed (slot=%d)", idx);
        }

        clients[idx].ecdhReady = true;
        ESP_LOGI(TAG, "ECDH: handshake complete (slot=%d)", idx);
    }
};

class MyCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
    void onWrite(NimBLECharacteristic *pChar, NimBLEConnInfo &connInfo) override {
        if (otaInProgress) return;

        int idx = findClientIdx(connInfo.getConnHandle());
        if (idx < 0) {
            ESP_LOGW(TAG, "BLE: write from unknown conn handle %d", connInfo.getConnHandle());
            return;
        }

        std::string value = pChar->getValue();
        if (value.length() != AES_PKT_SIZE) {
            ESP_LOGW(TAG, "BLE: bad packet length %d (expected %d), slot=%d",
                     (int)value.length(), AES_PKT_SIZE, idx);
            return;
        }

        const uint8_t *pkt = (const uint8_t *)value.data();

        uint8_t cmd;
        if (!decryptAesGcm(pkt, idx, &cmd)) {
            return;
        }

        uint32_t seq = lastSeq[idx];

        if (cmd == 0x01) {
            bleRelay1[idx] = true;
            ESP_LOGI(TAG, "BLE: relay1 ON (seq=%lu, slot=%d)", seq, idx);
        } else if (cmd == 0x00) {
            bleRelay1[idx] = false;
            ESP_LOGI(TAG, "BLE: relay1 OFF (seq=%lu, slot=%d)", seq, idx);
        } else if (cmd == 0x02) {
            bleRelay2[idx] = true;
            ESP_LOGI(TAG, "BLE: relay2 ON (seq=%lu, slot=%d)", seq, idx);
        } else if (cmd == 0x03) {
            bleRelay2[idx] = false;
            ESP_LOGI(TAG, "BLE: relay2 OFF (seq=%lu, slot=%d)", seq, idx);
        }
    }
};

void setup() {
    Serial.begin(115200);
    delay(2000);
    ESP_LOGI(TAG, "=== AtomS3 Lite BLE Relay v%s ===", FIRMWARE_VERSION);

    FastLED.addLeds<WS2812, LED_PIN, GRB>(leds, NUM_LEDS);

    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);
    preferences.end();

    if (pskLen == PSK_LEN) {
        pskProvisioned = true;
        ESP_LOGI(TAG, "PSK loaded from NVS");
    } else {
        pskProvisioned = false;
        ESP_LOGW(TAG, "No PSK — connect USB and run SET_KEY");
    }

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

    for (int i = 0; i < MAX_CLIENTS; i++) {
        clients[i].ecdhReady = false;
        memset(clients[i].aesKey, 0, AES_KEY_LEN);
    }

    pinMode(TOUCH1_PIN, INPUT_PULLUP);
    pinMode(TOUCH2_PIN, INPUT_PULLUP);
    pinMode(BUTTON_PIN, INPUT);
    pinMode(RELAY1_PIN, OUTPUT);
    pinMode(RELAY2_PIN, OUTPUT);

    leds[0] = CRGB::Red;
    FastLED.show();

    NimBLEDevice::init("SylDa_Server");
    NimBLEDevice::setMTU(248);

    NimBLEServer *pServer = NimBLEDevice::createServer();
    pServer->setCallbacks(new MyServerCallbacks());
    pServer->advertiseOnDisconnect(true);

    NimBLEService *pService = pServer->createService(SERVICE_UUID);
    pCharacteristic = pService->createCharacteristic(
        CHARACTERISTIC_UUID,
        NIMBLE_PROPERTY::WRITE_NR
    );
    pCharacteristic->setCallbacks(new MyCharacteristicCallbacks());

    pEcdhChar = pService->createCharacteristic(
        ECDH_UUID,
        NIMBLE_PROPERTY::WRITE_NR | NIMBLE_PROPERTY::NOTIFY
    );
    pEcdhChar->setCallbacks(new EcdhCharacteristicCallbacks());

    pUuidChar = pService->createCharacteristic(
        UUID_CHAR_UUID,
        NIMBLE_PROPERTY::READ
    );
    pUuidChar->setValue(deviceUuid, ESP_UUID_LEN);

    NimBLECharacteristic *pVersionChar = pService->createCharacteristic(
        VERSION_UUID,
        NIMBLE_PROPERTY::READ
    );
    pVersionChar->setValue(FIRMWARE_VERSION);

    NimBLECharacteristic *pOtaChar = pService->createCharacteristic(
        OTA_UUID,
        NIMBLE_PROPERTY::WRITE_NR | NIMBLE_PROPERTY::NOTIFY
    );
    pOtaChar->setCallbacks(new OtaCharacteristicCallbacks());

    pService->start();

    NimBLEAdvertising *pAdvertising = NimBLEDevice::getAdvertising();
    pAdvertising->setName("SylDa_Server");

    NimBLEAdvertisementData scanResp;
    scanResp.setCompleteServices(NimBLEUUID(SERVICE_UUID));

    uint8_t mfgPayload[6];
    mfgPayload[0] = 0xFF;
    mfgPayload[1] = 0xFF;
    memcpy(mfgPayload + 2, deviceUuid, 4);
    scanResp.setManufacturerData(std::string((char*)mfgPayload, 6));

    pAdvertising->setScanResponseData(scanResp);
    pAdvertising->enableScanResponse(true);

    pAdvertising->setMinInterval(32);
    pAdvertising->setMaxInterval(64);
    bool advOk = pAdvertising->start();
    ESP_LOGI(TAG, "BLE: advertising started");

    if (pskProvisioned) {
        ESP_LOGD(TAG, "PSK present");
    } else {
        ESP_LOGW(TAG, "No PSK — use USB SET_KEY to provision");
    }

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

void loop() {
    if (!otaInProgress && 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();
                } else {
                    Serial.println("NO_KEY");
                }
            } else if (serialBuf == "ID?") {
                Serial.println("SylDa_Server");
            } else if (serialBuf == "VERSION?") {
                Serial.flush();
                Serial.print("BUILD_NR:");
                Serial.println(FIRMWARE_VERSION);
            } 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;
                    Serial.println("OK");
                    Serial.println("PSK stored in NVS");
                } else {
                    Serial.println("ERROR: invalid key format");
                }
            } 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;
                    pUuidChar->setValue(deviceUuid, ESP_UUID_LEN);

                    uint8_t mfgPayload[6];
                    mfgPayload[0] = 0xFF;
                    mfgPayload[1] = 0xFF;
                    memcpy(mfgPayload + 2, deviceUuid, 4);
                    NimBLEAdvertisementData scanResp;
                    scanResp.setCompleteServices(NimBLEUUID(SERVICE_UUID));
                    scanResp.setManufacturerData(std::string((char*)mfgPayload, 6));
                    NimBLEDevice::getAdvertising()->setScanResponseData(scanResp);
                    Serial.println("OK");
                    Serial.println("UUID stored in NVS");
                } else {
                    Serial.println("ERROR: invalid UUID format");
                }
            }
            serialBuf = "";
        }
    }

    if (otaReady) {
        ESP_LOGD(TAG, "OTA: processing otaReady (written=%lu expected=%lu)",
                 otaWritten, otaExpectedSize);
        esp_err_t err = esp_ota_end(otaHandle);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "OTA: esp_ota_end failed (%d) — aborting", err);
            otaReset();
        } else {
            ESP_LOGI(TAG, "OTA: esp_ota_end OK, waiting for 0xFF confirm");
            otaEnded = true;
            otaEndedAt = millis();
        }
        otaReady = false;
    }

    if (otaEnded && otaEarlyConfirm) {
        ESP_LOGD(TAG, "OTA: processing deferred 0xFF confirm");
        otaConfirmed = true;
        otaEarlyConfirm = false;
    }

    if (otaEnded && !otaConfirmed && (millis() - otaEndedAt > OTA_CONFIRM_TIMEOUT_MS)) {
        ESP_LOGW(TAG, "OTA: 0xFF confirm timeout (%lu ms) — discarding update",
                 millis() - otaEndedAt);
        otaReset();
    }

    if (otaConfirmed) {
        ESP_LOGI(TAG, "OTA: confirmed, setting boot partition");
        esp_err_t err = esp_ota_set_boot_partition(otaPartition);
        if (err != ESP_OK) {
            ESP_LOGE(TAG, "OTA: set boot partition failed (%d)", err);
        } else {
            leds[0] = CRGB::Green;
            FastLED.show();
            ESP_LOGI(TAG, "OTA: done, rebooting...");
            delay(500);
            esp_restart();
        }
        otaReset();
    }

    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;
    #define DEBOUNCE_MS 50
    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;

    if (!otaInProgress) {
        relay1State = !t1Active || btnActive ||
                      (connActive[0] && bleRelay1[0]) || (connActive[1] && bleRelay1[1]);
        relay2State = !t2Active ||
                      (connActive[0] && bleRelay2[0]) || (connActive[1] && bleRelay2[1]);
    } else {
        relay1State = false;
        relay2State = false;
    }

    CRGB ledColor;

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

        if (connectedCount > 0) {
            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;

    digitalWrite(RELAY1_PIN, relay1State ? HIGH : LOW);
    digitalWrite(RELAY2_PIN, relay2State ? HIGH : LOW);

    FastLED.show();

    static bool lastRelay1 = false;
    static bool lastRelay2 = false;
    if (relay1State != lastRelay1 || relay2State != lastRelay2) {
        bool anyBle1 = (connActive[0] && bleRelay1[0]) || (connActive[1] && bleRelay1[1]);
        bool anyBle2 = (connActive[0] && bleRelay2[0]) || (connActive[1] && bleRelay2[1]);
        ESP_LOGI(TAG, "Relays: R1=%s R2=%s (t1=%s t2=%s btn=%s ble1=%s ble2=%s)",
            relay1State ? "ON" : "OFF",
            relay2State ? "ON" : "OFF",
            !t1Active ? "YES" : "no",
            !t2Active ? "YES" : "no",
            btnActive ? "YES" : "no",
            anyBle1 ? "ON" : "OFF",
            anyBle2 ? "ON" : "OFF");
        lastRelay1 = relay1State;
        lastRelay2 = relay2State;
    }

    static uint32_t heartbeat = 0;
    if (!otaInProgress && millis() - heartbeat > 60000) {
        bool anyBle1 = (connActive[0] && bleRelay1[0]) || (connActive[1] && bleRelay1[1]);
        bool anyBle2 = (connActive[0] && bleRelay2[0]) || (connActive[1] && bleRelay2[1]);
        ESP_LOGD(TAG, "[heartbeat] R1=%s R2=%s t1=%s t2=%s btn=%s ble1=%s ble2=%s clients=%d",
            relay1State ? "ON" : "OFF",
            relay2State ? "ON" : "OFF",
            !t1Active ? "YES" : "no",
            !t2Active ? "YES" : "no",
            btnActive ? "YES" : "no",
            anyBle1 ? "ON" : "OFF",
            anyBle2 ? "ON" : "OFF",
            connectedCount);
        heartbeat = millis();
    }

    delay(50);
}