#include <Arduino.h>
#include <BLEDevice.h>
#include <BLEServer.h>
#include <BLEUtils.h>
#include <BLE2902.h>

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

#define SENSOR_PIN 38
#define RELAY_PIN  1
#define LED_PIN    46

#define SERVICE_UUID        "4fafc201-1fb5-459e-8fcc-c5c9c331914b"
#define CHARACTERISTIC_UUID "beb5483e-36e1-4688-b7f5-ea07361f26d8"

#define HMAC_PKT_SIZE 37   // 1 cmd + 4 seq + 32 hmac
#define HMAC_KEY_LEN  32
#define NVS_NAMESPACE "hmac"
#define NVS_KEY_NAME  "key"

BLECharacteristic *pCharacteristic;
bool deviceConnected = false;
bool bleCommand = false;
bool relayState = false;
uint32_t lastSeq = 0;

uint8_t hmacKey[HMAC_KEY_LEN];
bool keyProvisioned = false;
Preferences preferences;

class MyServerCallbacks : public BLEServerCallbacks {
    void onConnect(BLEServer *pServer) override {
        deviceConnected = true;
        lastSeq = 0;
        Serial.println("BLE: device connected");
    }
    void onDisconnect(BLEServer *pServer) override {
        deviceConnected = false;
        bleCommand = false;
        lastSeq = 0;
        Serial.println("BLE: device disconnected");
        BLEDevice::startAdvertising();
    }
};

static bool verifyHmac(const uint8_t *packet) {
    uint8_t cmd = packet[0];
    uint32_t seq = ((uint32_t)packet[1] << 24) |
                   ((uint32_t)packet[2] << 16) |
                   ((uint32_t)packet[3] << 8)  |
                   (uint32_t)packet[4];

    if (seq <= lastSeq) {
        Serial.printf("BLE: replay rejected (seq=%lu, last=%lu)\n", seq, lastSeq);
        return false;
    }

    uint8_t hmacData[5];
    hmacData[0] = cmd;
    hmacData[1] = (seq >> 24) & 0xFF;
    hmacData[2] = (seq >> 16) & 0xFF;
    hmacData[3] = (seq >> 8)  & 0xFF;
    hmacData[4] = seq & 0xFF;

    uint8_t computed[32];
    const mbedtls_md_info_t *mdInfo = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
    mbedtls_md_hmac(mdInfo, hmacKey, HMAC_KEY_LEN, hmacData, 5, computed);

    if (memcmp(computed, packet + 5, 32) != 0) {
        Serial.println("BLE: HMAC mismatch");
        return false;
    }

    lastSeq = seq;
    return true;
}

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) {
    for (int i = 0; i < HMAC_KEY_LEN; i++) {
        out[i] = (uint8_t)strtoul(hex.substring(i * 2, i * 2 + 2).c_str(), nullptr, 16);
    }
}

class MyCharacteristicCallbacks : public BLECharacteristicCallbacks {
    void onWrite(BLECharacteristic *pChar) override {
        std::string value = pChar->getValue();
        if (value.length() != HMAC_PKT_SIZE) {
            Serial.printf("BLE: bad packet length %d (expected %d)\n",
                          (int)value.length(), HMAC_PKT_SIZE);
            return;
        }

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

        if (!verifyHmac(pkt)) {
            return;
        }

        uint8_t cmd = pkt[0];
        uint32_t seq = lastSeq;

        if (cmd == 0x01) {
            bleCommand = true;
            Serial.printf("BLE: relay ON (seq=%lu)\n", seq);
        } else {
            bleCommand = false;
            Serial.printf("BLE: relay OFF (seq=%lu)\n", seq);
        }
    }
};

void setup() {
    delay(2000);
    Serial.begin(115200);
    delay(500);
    Serial.println("=== Atom VoiceS3R BLE Relay ===");

    preferences.begin(NVS_NAMESPACE, true);
    size_t len = preferences.getBytes(NVS_KEY_NAME, hmacKey, HMAC_KEY_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 provisioned — connect USB and run SET_KEY");
    }

    pinMode(SENSOR_PIN, INPUT_PULLUP);
    pinMode(RELAY_PIN, OUTPUT);
    pinMode(LED_PIN, OUTPUT);

    for (int i = 0; i < 3; i++) {
        digitalWrite(LED_PIN, HIGH);
        delay(200);
        digitalWrite(LED_PIN, LOW);
        delay(200);
    }
    Serial.println("LED blink done, starting BLE...");

    BLEDevice::init("AtomS3R-Relay");

    BLEServer *pServer = BLEDevice::createServer();
    pServer->setCallbacks(new MyServerCallbacks());

    BLEService *pService = pServer->createService(SERVICE_UUID);
    pCharacteristic = pService->createCharacteristic(
        CHARACTERISTIC_UUID,
        BLECharacteristic::PROPERTY_WRITE_NR | BLECharacteristic::PROPERTY_NOTIFY
    );
    pCharacteristic->setAccessPermissions(ESP_GATT_PERM_WRITE);
    pCharacteristic->addDescriptor(new BLE2902());
    pCharacteristic->setCallbacks(new MyCharacteristicCallbacks());
    pCharacteristic->setValue("0");
    pService->start();

    BLEAdvertising *pAdvertising = BLEDevice::getAdvertising();
    pAdvertising->addServiceUUID(SERVICE_UUID);
    pAdvertising->setScanResponse(true);
    pAdvertising->setMinPreferred(0x06);
    pAdvertising->setMinPreferred(0x12);
    BLEDevice::startAdvertising();

    Serial.println("BLE: advertising started as 'AtomS3R-Relay'");

    if (keyProvisioned) {
        Serial.print("HMAC Key: ");
        for (int i = 0; i < HMAC_KEY_LEN; i++) {
            Serial.printf("%02X", hmacKey[i]);
        }
        Serial.println();
    } else {
        Serial.println("No HMAC key — use USB SET_KEY to provision");
    }

    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) {
                    Serial.print("HMAC_KEY:");
                    for (int i = 0; i < HMAC_KEY_LEN; i++) {
                        Serial.printf("%02X", hmacKey[i]);
                    }
                    Serial.println();
                } else {
                    Serial.println("NO_KEY");
                }
            } else if (serialBuf.startsWith("SET_KEY:")) {
                String hex = serialBuf.substring(8);
                hex.toUpperCase();
                if (hex.length() == 64 && isHexKey(hex)) {
                    parseHexKey(hex, hmacKey);
                    preferences.begin(NVS_NAMESPACE, false);
                    preferences.putBytes(NVS_KEY_NAME, hmacKey, HMAC_KEY_LEN);
                    preferences.end();
                    keyProvisioned = true;
                    Serial.println("OK");
                    Serial.println("HMAC key stored in NVS");
                } else {
                    Serial.println("ERROR: invalid key format");
                }
            }
            serialBuf = "";
        }
    }

    int sensorState = digitalRead(SENSOR_PIN);

    bool touchActive = (sensorState == LOW);
    bool bleActive = (deviceConnected && bleCommand);

    relayState = touchActive || bleActive;

    digitalWrite(RELAY_PIN, relayState ? HIGH : LOW);

    static bool lastRelay = false;
    if (relayState != lastRelay) {
        Serial.printf("Relay %s (touch=%s, ble=%s)\n",
            relayState ? "ON" : "OFF",
            sensorState == LOW ? "YES" : "no",
            deviceConnected ? "connected" : "disconnected");
        lastRelay = relayState;
    }

    static uint32_t heartbeat = 0;
    if (millis() - heartbeat > 5000) {
        Serial.printf("[heartbeat] relay=%s touch=%s ble=%s\n",
            relayState ? "ON" : "OFF",
            sensorState == LOW ? "YES" : "no",
            deviceConnected ? "connected" : "disconnected");
        heartbeat = millis();
    }

    delay(50);
}