// ============================================================================
// auth_gate.cpp — Ed25519 auth-token verification gate
//
// See auth_gate.h for full API, token wire format, and audit log spec.
// ============================================================================

#include "auth_gate.h"
#include "ota_signed.h"   // ota_signed_get_pubkey, ED25519_PUBKEY_BYTES

// ---------------------------------------------------------------------------
// Ed25519 verify shim
//
// On device: calls ota_signed_verify_chunk() indirectly, or the raw
// ed25519_verify() from the embedded ref10 port — whichever Batch B exposes.
//
// The Batch B ota_signed.cpp provides:
//   bool ota_signed_verify_chunk(const uint8_t* chunk, size_t len,
//                                const uint8_t* sig64)
// which internally uses the cached public key.  We replicate the same pattern
// here but operate on token message bytes instead of OTA chunks.
//
// If Batch B is not yet linked, the `extern` below will be resolved by the
// integrator who merges Batch B.  Under UNIT_TEST a software mock is
// substituted.
// ---------------------------------------------------------------------------

#ifndef UNIT_TEST
// Forward declaration — resolved by ota_signed.cpp (Batch B) at link time.
// The function verifies sig64 over msg[0..len-1] using the NVS-cached key.
extern bool ota_signed_verify_chunk(const uint8_t* chunk, size_t len,
                                    const uint8_t* sig64);
// NVS helpers (from ESP-IDF nvs_flash)
#  include "nvs_flash.h"
#  include "nvs.h"
// millis() is Arduino; fall back to esp_timer on bare ESP-IDF
#  include "Arduino.h"
#else
// ---- Host-side stubs -------------------------------------------------------
#  include <cstdlib>
#  include <cstring>
#  include <cstdio>

// Test-injectable Ed25519 verify mock.
// Tests set these to control pass/fail.
bool g_ed25519_verify_result = true;
// Also expose the pubkey for tests to populate.  External linkage so the
// test TU's `extern` declarations resolve here.
uint8_t g_test_pubkey[ED25519_PUBKEY_BYTES] = {};
bool    g_test_pubkey_set = false;

// Stub for ota_signed_get_pubkey
bool ota_signed_get_pubkey(uint8_t* pk32_out) {
    if (!g_test_pubkey_set) return false;
    memcpy(pk32_out, g_test_pubkey, ED25519_PUBKEY_BYTES);
    return true;
}
// Stub for ota_signed_verify_chunk
bool ota_signed_verify_chunk(const uint8_t* /*chunk*/, size_t /*len*/,
                              const uint8_t* /*sig64*/) {
    return g_ed25519_verify_result;
}
// NVS stub: simple in-memory ring buffer (no persistence in tests)
static uint8_t  g_nvs_audit_buf[AUTHGATE_AUDIT_MAX_DUMP] = {};
static int      g_nvs_audit_used = 0;  // bytes written so far
static int      g_nvs_ring_head  = 0;  // next entry index to write (0-based)

#endif // UNIT_TEST

#include <cstring>
#include <cstdlib>

// ---------------------------------------------------------------------------
// SHA-256 stub for token_hash_first8 (first 8 bytes of SHA-256 of raw token)
//
// On device: use mbedTLS / ROM SHA-256.
// In tests: use a trivial FNV-1a-based hash to produce deterministic 8 bytes.
// ---------------------------------------------------------------------------
#ifndef UNIT_TEST
#  include "mbedtls/sha256.h"
static void compute_token_hash(const uint8_t* token, size_t len, uint8_t out8[8]) {
    uint8_t digest[32];
    mbedtls_sha256_context ctx;
    mbedtls_sha256_init(&ctx);
    mbedtls_sha256_starts(&ctx, 0);
    mbedtls_sha256_update(&ctx, token, len);
    mbedtls_sha256_finish(&ctx, digest);
    mbedtls_sha256_free(&ctx);
    memcpy(out8, digest, 8);
}
#else
static void compute_token_hash(const uint8_t* token, size_t len, uint8_t out8[8]) {
    // FNV-1a 64-bit over the token bytes
    uint64_t h = 14695981039346656037ULL;
    for (size_t i = 0; i < len; i++) {
        h ^= (uint64_t)token[i];
        h *= 1099511628211ULL;
    }
    memcpy(out8, &h, 8);
}
#endif

// ---------------------------------------------------------------------------
// Cached public key (loaded lazily from NVS via ota_signed_get_pubkey)
// ---------------------------------------------------------------------------
static uint8_t s_cached_pk[ED25519_PUBKEY_BYTES];
static bool    s_pk_loaded = false;

static bool load_pubkey_once() {
    if (s_pk_loaded) return true;
    if (!ota_signed_get_pubkey(s_cached_pk)) return false;
    s_pk_loaded = true;
    return true;
}

void authgate_reload_pubkey() {
    s_pk_loaded = false;
    load_pubkey_once();
}

// ---------------------------------------------------------------------------
// NVS audit ring buffer helpers
//
// On device: read/write "pdl_audit" / "gate_log" as a raw blob (max 1024B).
// In tests:  write to the g_nvs_audit_buf array.
// ---------------------------------------------------------------------------
#ifdef UNIT_TEST
static void nvs_audit_write_entry(const uint8_t entry[AUTHGATE_AUDIT_ENTRY_BYTES]) {
    int byte_off = g_nvs_ring_head * AUTHGATE_AUDIT_ENTRY_BYTES;
    memcpy(g_nvs_audit_buf + byte_off, entry, AUTHGATE_AUDIT_ENTRY_BYTES);
    g_nvs_ring_head = (g_nvs_ring_head + 1) % AUTHGATE_AUDIT_RING_SIZE;
    if (g_nvs_audit_used < AUTHGATE_AUDIT_MAX_DUMP)
        g_nvs_audit_used += AUTHGATE_AUDIT_ENTRY_BYTES;
}
static int nvs_audit_read_all(uint8_t* out, size_t max_len) {
    size_t copy = (g_nvs_audit_used < (int)max_len) ? (size_t)g_nvs_audit_used : max_len;
    memcpy(out, g_nvs_audit_buf, copy);
    return (int)copy;
}
#else
// On-device: serialize the ring into a flat NVS blob.
// We store [ring_head_idx (1B)][data (1024B)] = 1025 bytes max.
// For simplicity store as a flat array updated atomically.
static void nvs_audit_write_entry(const uint8_t entry[AUTHGATE_AUDIT_ENTRY_BYTES]) {
    nvs_handle_t h;
    if (nvs_open(AUTH_AUDIT_NVS_NS, NVS_READWRITE, &h) != ESP_OK) return;

    uint8_t blob[AUTHGATE_AUDIT_MAX_DUMP + 1] = {};
    size_t  blen = sizeof(blob);
    uint8_t ring_head = 0;
    // Read existing
    nvs_get_blob(h, AUTH_AUDIT_NVS_KEY, blob, &blen);
    if (blen >= 1) ring_head = blob[0];
    if (blen < AUTHGATE_AUDIT_MAX_DUMP + 1) {
        // First use — zero-initialize
        memset(blob, 0, sizeof(blob));
    }
    // Write new entry at ring_head
    uint16_t byte_off = (uint16_t)(1 + ring_head * AUTHGATE_AUDIT_ENTRY_BYTES);
    memcpy(blob + byte_off, entry, AUTHGATE_AUDIT_ENTRY_BYTES);
    ring_head = (ring_head + 1) % AUTHGATE_AUDIT_RING_SIZE;
    blob[0] = ring_head;
    nvs_set_blob(h, AUTH_AUDIT_NVS_KEY, blob, AUTHGATE_AUDIT_MAX_DUMP + 1);
    nvs_commit(h);
    nvs_close(h);
}
static int nvs_audit_read_all(uint8_t* out, size_t max_len) {
    nvs_handle_t h;
    if (nvs_open(AUTH_AUDIT_NVS_NS, NVS_READONLY, &h) != ESP_OK) return 0;
    uint8_t blob[AUTHGATE_AUDIT_MAX_DUMP + 1] = {};
    size_t blen = sizeof(blob);
    esp_err_t err = nvs_get_blob(h, AUTH_AUDIT_NVS_KEY, blob, &blen);
    nvs_close(h);
    if (err != ESP_OK || blen < 1) return 0;
    size_t copy = (blen - 1) < max_len ? (blen - 1) : max_len;
    memcpy(out, blob + 1, copy);
    return (int)copy;
}
#endif

// ---------------------------------------------------------------------------
// Public audit API
// ---------------------------------------------------------------------------

void authgate_audit_append(uint64_t       ts_ms,
                           uint32_t       sensor_id,
                           const uint8_t  bssid[6],
                           uint32_t       payload_id,
                           uint8_t        result,
                           const uint8_t* token,
                           size_t         token_len) {
    uint8_t entry[AUTHGATE_AUDIT_ENTRY_BYTES] = {};
    // [0..7]   ts_ms LE
    memcpy(entry + 0, &ts_ms, 8);
    // [8..11]  sensor_id LE
    memcpy(entry + 8, &sensor_id, 4);
    // [12..17] bssid
    if (bssid) memcpy(entry + 12, bssid, 6);
    // [18..21] payload_id LE
    memcpy(entry + 18, &payload_id, 4);
    // [22]     result
    entry[22] = result;
    // [23..30] token_hash_first8
    if (token && token_len >= AUTHGATE_TOKEN_LEN) {
        uint8_t h8[8];
        compute_token_hash(token, token_len, h8);
        memcpy(entry + 23, h8, 8);
    }
    // [31]     reserved = 0
    nvs_audit_write_entry(entry);
}

int authgate_audit_dump(uint8_t* out_buf, size_t max_len) {
    if (!out_buf || max_len == 0) return 0;
    return nvs_audit_read_all(out_buf, max_len);
}

// ---------------------------------------------------------------------------
// Core verification
// ---------------------------------------------------------------------------

bool authgate_verify(uint32_t         sensor_id,
                     const uint8_t    bssid[6],
                     uint32_t         payload_id,
                     uint64_t         now_ms,
                     const uint8_t*   token,
                     size_t           token_len) {
    // (a) Length check
    if (!token || token_len != AUTHGATE_TOKEN_LEN) {
        // Can't append audit without a valid token; do a minimal log
        authgate_audit_append(now_ms, sensor_id, bssid, payload_id,
                              AUDIT_RESULT_FAIL_SIG, nullptr, 0);
        return false;
    }

    // Parse token fields
    uint32_t tok_sensor_id;
    uint32_t tok_payload_id;
    uint64_t tok_expiry_ms;

    memcpy(&tok_sensor_id,  token + 0,  4);
    // token[4..9] = target_bssid — covered by Ed25519 signature, not re-checked here
    memcpy(&tok_payload_id, token + 10, 4);
    memcpy(&tok_expiry_ms,  token + 14, 8);
    // token[22..23] = reserved (ignored)
    const uint8_t* sig64 = token + AUTHGATE_SIG_OFFSET;

    // (b) Sensor ID check
    if (tok_sensor_id != sensor_id) {
        authgate_audit_append(now_ms, sensor_id, bssid, payload_id,
                              AUDIT_RESULT_FAIL_ID, token, token_len);
        return false;
    }

    // (c) Payload ID check
    if (tok_payload_id != payload_id) {
        authgate_audit_append(now_ms, sensor_id, bssid, payload_id,
                              AUDIT_RESULT_FAIL_ID, token, token_len);
        return false;
    }

    // (d) Expiry check
    if (tok_expiry_ms <= now_ms) {
        authgate_audit_append(now_ms, sensor_id, bssid, payload_id,
                              AUDIT_RESULT_FAIL_EXPIRY, token, token_len);
        return false;
    }

    // (e) Ed25519 signature verification
    // The signature covers token[0..AUTHGATE_TOKEN_MSG_LEN-1] which includes
    // sensor_id, target_bssid, payload_id, expiry_ms, and reserved bytes.
    // The BSSID binding is enforced cryptographically by the signature rather
    // than by a separate plaintext comparison (avoids aliasing the bssid param).
    if (!load_pubkey_once()) {
        authgate_audit_append(now_ms, sensor_id, bssid, payload_id,
                              AUDIT_RESULT_FAIL_NOPK, token, token_len);
        return false;
    }

    // ota_signed_verify_chunk uses its internally cached key; to ensure it
    // uses our freshly loaded key we call ota_signed_set_pubkey() if needed.
    // Under UNIT_TEST the stub ignores the key and returns g_ed25519_verify_result.
    // On device, ota_signed.cpp caches the key from NVS just as we do.
    bool sig_ok = ota_signed_verify_chunk(token, AUTHGATE_TOKEN_MSG_LEN, sig64);

    uint8_t result = sig_ok ? AUDIT_RESULT_PASS : AUDIT_RESULT_FAIL_SIG;
    authgate_audit_append(now_ms, sensor_id, bssid, payload_id,
                          result, token, token_len);
    return sig_ok;
}
