// ============================================================================
// enrollment.cpp — First-boot self-registration
// ============================================================================

#include "enrollment.h"
#include "config.h"

// ---------------------------------------------------------------------------
// Platform shims — Arduino vs native/unit-test
// ---------------------------------------------------------------------------
#ifdef ARDUINO
#   include <Arduino.h>
#   include <esp_efuse.h>
#   include <esp_efuse_table.h>
#   include <nvs_flash.h>
#   include <nvs.h>

static uint32_t enroll_millis() { return millis(); }

// NVS helpers ----------------------------------------------------------------
static bool nvs_read_u8(const char* ns, const char* key, uint8_t* out) {
    nvs_handle_t h;
    if (nvs_open(ns, NVS_READONLY, &h) != ESP_OK) return false;
    esp_err_t err = nvs_get_u8(h, key, out);
    nvs_close(h);
    return (err == ESP_OK);
}

static bool nvs_write_u8(const char* ns, const char* key, uint8_t val) {
    nvs_handle_t h;
    if (nvs_open(ns, NVS_READWRITE, &h) != ESP_OK) return false;
    bool ok = (nvs_set_u8(h, key, val) == ESP_OK) && (nvs_commit(h) == ESP_OK);
    nvs_close(h);
    return ok;
}

static bool nvs_write_u32(const char* ns, const char* key, uint32_t val) {
    nvs_handle_t h;
    if (nvs_open(ns, NVS_READWRITE, &h) != ESP_OK) return false;
    bool ok = (nvs_set_u32(h, key, val) == ESP_OK) && (nvs_commit(h) == ESP_OK);
    nvs_close(h);
    return ok;
}

static bool nvs_write_i32(const char* ns, const char* key, int32_t val) {
    nvs_handle_t h;
    if (nvs_open(ns, NVS_READWRITE, &h) != ESP_OK) return false;
    bool ok = (nvs_set_i32(h, key, val) == ESP_OK) && (nvs_commit(h) == ESP_OK);
    nvs_close(h);
    return ok;
}

// Hardware UID: lower 32 bits of the default MAC (unique per ESP32-S3)
static uint32_t get_hw_uid() {
    uint8_t mac[6] = {0};
    esp_efuse_mac_get_default(mac);
    return ((uint32_t)mac[2] << 24) | ((uint32_t)mac[3] << 16)
         | ((uint32_t)mac[4] <<  8) |  (uint32_t)mac[5];
}

// Random 12-byte boot token (non-secure; used for replay detection only)
static void fill_boot_token(uint8_t* out12) {
    for (int i = 0; i < 12; i++) {
        out12[i] = (uint8_t)(esp_random() & 0xFF);
    }
}

extern void mesh_publish_raw(const uint8_t* data, size_t len);
static void radio_send(const uint8_t* buf, size_t len) {
    mesh_publish_raw(buf, len);
}

#else // ---- Native / unit-test stubs ----------------------------------------
#   include <cstdlib>
#   include <cstring>
#   include <cstdio>

static uint32_t g_stub_millis = 0;
static uint32_t enroll_millis() { return g_stub_millis; }
void enroll_stub_set_millis(uint32_t ms) { g_stub_millis = ms; }

// In-memory NVS substitute
struct StubNvsEntry { char key[32]; uint8_t type; union { uint8_t u8; uint32_t u32; int32_t i32; }; };
static StubNvsEntry g_nvs[32];
static int g_nvs_count = 0;

static StubNvsEntry* nvs_find(const char* key) {
    for (int i = 0; i < g_nvs_count; i++) {
        if (strcmp(g_nvs[i].key, key) == 0) return &g_nvs[i];
    }
    return nullptr;
}

static bool nvs_read_u8(const char*, const char* key, uint8_t* out) {
    auto* e = nvs_find(key);
    if (!e || e->type != 1) return false;
    *out = e->u8; return true;
}
static bool nvs_write_u8(const char*, const char* key, uint8_t val) {
    auto* e = nvs_find(key);
    if (!e) {
        if (g_nvs_count >= 32) return false;
        e = &g_nvs[g_nvs_count++];
        strncpy(e->key, key, 31);
    }
    e->type = 1; e->u8 = val; return true;
}
static bool nvs_write_u32(const char*, const char* key, uint32_t val) {
    auto* e = nvs_find(key);
    if (!e) {
        if (g_nvs_count >= 32) return false;
        e = &g_nvs[g_nvs_count++];
        strncpy(e->key, key, 31);
    }
    e->type = 2; e->u32 = val; return true;
}
static bool nvs_write_i32(const char*, const char* key, int32_t val) {
    auto* e = nvs_find(key);
    if (!e) {
        if (g_nvs_count >= 32) return false;
        e = &g_nvs[g_nvs_count++];
        strncpy(e->key, key, 31);
    }
    e->type = 3; e->i32 = val; return true;
}

static uint32_t get_hw_uid() { return 0xDEADBEEFu; }
static void fill_boot_token(uint8_t* out12) {
    for (int i = 0; i < 12; i++) out12[i] = (uint8_t)(i ^ 0xA5u);
}

// Capture last transmission for test inspection
uint8_t  g_last_tx_buf[ENROLLMENT_BROADCAST_SIZE];
size_t   g_last_tx_len = 0;
static void radio_send(const uint8_t* buf, size_t len) {
    size_t cp = len < ENROLLMENT_BROADCAST_SIZE ? len : ENROLLMENT_BROADCAST_SIZE;
    memcpy(g_last_tx_buf, buf, cp);
    g_last_tx_len = cp;
}

// Expose NVS stub reset for tests
void enroll_stub_reset_nvs() {
    g_nvs_count = 0;
    memset(g_nvs, 0, sizeof(g_nvs));
}
#endif

// ---------------------------------------------------------------------------
// MSG_TYPE guards — may already be defined in config.h after integration
// ---------------------------------------------------------------------------
#ifndef MSG_TYPE_ENROLLMENT
#   define MSG_TYPE_ENROLLMENT           0x17
#endif
#ifndef MSG_TYPE_ENROLLMENT_RESPONSE
#   define MSG_TYPE_ENROLLMENT_RESPONSE  0x18
#endif

#define ENROLL_NVS_NS         "pdl_enroll"
#define ENROLL_KEY_DONE       "done"
#define ENROLL_KEY_SENSOR_ID  "sensor_id"
#define ENROLL_KEY_LAT        "lat"
#define ENROLL_KEY_LON        "lon"

// FW version: parse "1.0.0" → major=1, minor=0
static uint16_t fw_version_word() {
    uint8_t major = 1, minor = 0;
#ifdef SENSOR_FW_VERSION
    // Best-effort parse; non-critical
    const char* s = SENSOR_FW_VERSION;
    major = (uint8_t)(s[0] - '0');
    if (s[1] == '.' && s[2] >= '0') minor = (uint8_t)(s[2] - '0');
#endif
    return ((uint16_t)major << 8) | minor;
}

// ---------------------------------------------------------------------------
// Module state
// ---------------------------------------------------------------------------
static bool     s_enrolled       = false;
static uint32_t s_sensor_id      = DEFAULT_SENSOR_ID;
static uint8_t  s_boot_token[12] = {0};
static uint32_t s_last_broadcast = 0;

// Build and transmit the 24-byte enrollment broadcast
static void send_enrollment_broadcast() {
    uint8_t pkt[ENROLLMENT_BROADCAST_SIZE];
    memset(pkt, 0, sizeof(pkt));

    pkt[0] = MSG_TYPE_ENROLLMENT;

    uint32_t uid = get_hw_uid();
    pkt[1] = (uint8_t)(uid      );
    pkt[2] = (uint8_t)(uid >>  8);
    pkt[3] = (uint8_t)(uid >> 16);
    pkt[4] = (uint8_t)(uid >> 24);

    uint16_t fwv = fw_version_word();
    pkt[5] = (uint8_t)(fwv >> 8);   // major
    pkt[6] = (uint8_t)(fwv     );   // minor

    memcpy(&pkt[7], s_boot_token, 12);
    // bytes [19..23] remain 0 (reserved)

    radio_send(pkt, sizeof(pkt));
    s_last_broadcast = enroll_millis();
}

// ---------------------------------------------------------------------------
bool enrollment_init() {
    // Generate a fresh boot token each power cycle
    fill_boot_token(s_boot_token);

    // Check NVS for prior enrollment
    uint8_t done = 0;
    if (nvs_read_u8(ENROLL_NVS_NS, ENROLL_KEY_DONE, &done) && done == 1) {
        s_enrolled  = true;
        // Retrieve stored sensor ID (best-effort; keep DEFAULT if missing)
        uint32_t stored_id = DEFAULT_SENSOR_ID;
#ifdef ARDUINO
        nvs_handle_t h;
        if (nvs_open(ENROLL_NVS_NS, NVS_READONLY, &h) == ESP_OK) {
            nvs_get_u32(h, ENROLL_KEY_SENSOR_ID, &stored_id);
            nvs_close(h);
        }
#else
        auto* e = nvs_find(ENROLL_KEY_SENSOR_ID);
        if (e && e->type == 2) stored_id = e->u32;
#endif
        s_sensor_id = stored_id;
        return true;   // already enrolled
    }

    // Not enrolled — broadcast immediately
    send_enrollment_broadcast();
    return false;
}

// ---------------------------------------------------------------------------
void enrollment_tick() {
    if (s_enrolled) return;
    uint32_t now = enroll_millis();
    if (now - s_last_broadcast >= ENROLLMENT_RETRY_MS) {
        send_enrollment_broadcast();
    }
}

// ---------------------------------------------------------------------------
void enrollment_handle_response(const uint8_t* buf, size_t len) {
    if (!buf || len < ENROLLMENT_RESPONSE_SIZE) return;
    if (buf[0] != MSG_TYPE_ENROLLMENT_RESPONSE) return;
    if (s_enrolled) return;   // idempotent

    // Parse fields (all little-endian)
    uint32_t assigned_id =
        ((uint32_t)buf[1]      ) |
        ((uint32_t)buf[2] <<  8) |
        ((uint32_t)buf[3] << 16) |
        ((uint32_t)buf[4] << 24);

    int32_t lat =
        (int32_t)(((uint32_t)buf[5]      ) |
                  ((uint32_t)buf[6] <<  8) |
                  ((uint32_t)buf[7] << 16) |
                  ((uint32_t)buf[8] << 24));

    int32_t lon =
        (int32_t)(((uint32_t)buf[9]       ) |
                  ((uint32_t)buf[10] <<  8) |
                  ((uint32_t)buf[11] << 16) |
                  ((uint32_t)buf[12] << 24));

    // bytes [13..15] reserved; bytes [16..27] = auth_pubkey_first16 (stored
    // for use by ota_signed — future: write to pdl_pki/ed25519_pk_hi16)

    // Persist to NVS
    nvs_write_u32(ENROLL_NVS_NS, ENROLL_KEY_SENSOR_ID, assigned_id);
    nvs_write_i32(ENROLL_NVS_NS, ENROLL_KEY_LAT,       lat);
    nvs_write_i32(ENROLL_NVS_NS, ENROLL_KEY_LON,       lon);
    nvs_write_u8 (ENROLL_NVS_NS, ENROLL_KEY_DONE,      1);

    s_sensor_id = assigned_id;
    s_enrolled  = true;
}

// ---------------------------------------------------------------------------
bool enrollment_is_complete() {
    return s_enrolled;
}

uint32_t enrollment_get_sensor_id() {
    return s_sensor_id;
}
