// =============================================================================
// PDL CONSULTING LLC — TRADE SECRET
// Copyright (c) 2026 PDL Consulting LLC. All rights reserved.
//
// src/talon_federated.cpp
// TALON Federated-Learning — T1 Receiver Side (v1)
//
// v1 scope: receive aggregated weight-delta chunks, reassemble in PSRAM,
// verify CRC32, apply atomically.  No on-device gradient computation.
//
// No malloc in hot path.  PSRAM allocation occurs once in talon_init().
// All reassembly state lives in static structs (< 40 bytes SRAM).
// =============================================================================

#include "talon_federated.h"

#include <cstring>
#include <cstdint>
#include <cstddef>

#ifndef UNIT_TEST
#  include <esp_heap_caps.h>
#  include <Arduino.h>     // millis()
#endif

// Pull in OTA signed verification if it has been compiled in.
// The guard lets the unit-test build omit the ESP32 NVS dependency.
#ifndef UNIT_TEST
#include "ota_signed.h"
#endif

// ---------------------------------------------------------------------------
// CRC-32 (ISO 3309 polynomial, same as zlib/ESP-IDF CRC).
// Pure software implementation — no table allocation, ~40 µs for 200 bytes.
// ---------------------------------------------------------------------------
static uint32_t crc32_update(uint32_t crc, const uint8_t* data, size_t len) {
    crc = ~crc;
    for (size_t i = 0; i < len; i++) {
        uint8_t b = data[i];
        for (int bit = 0; bit < 8; bit++) {
            uint32_t mask = (uint32_t)(-(int32_t)((crc ^ b) & 1u));
            crc = (crc >> 1) ^ (0xEDB88320u & mask);
            b >>= 1;
        }
    }
    return ~crc;
}

// ---------------------------------------------------------------------------
// Reassembly state (SRAM — 36 bytes)
// ---------------------------------------------------------------------------
struct TalonRoundState {
    uint16_t round_id;          // round_id from the most recent chunk header
    uint32_t total_chunks;      // expected total chunk count
    uint32_t received_chunks;   // number of distinct chunks received
    uint32_t expected_crc32;    // CRC32 of the complete delta (from header)
    uint32_t apply_after_ts;    // Unix epoch seconds; 0 = apply immediately
    uint32_t write_cursor;      // next byte offset in s_candidate_buf
    uint32_t bytes_per_chunk;   // actual payload bytes in each chunk (set from
                                //   first received chunk; used as reassembly stride)
    bool     crc_ok;            // set to true when reassembly CRC verified
    bool     poisoned;          // set to true on CRC mismatch; cleared on next round
};

// Bit-array tracking which chunk_idx values have been received.
// 512 bits = 64 bytes; covers TALON_MAX_CHUNKS = 512.
static uint8_t s_chunk_received[TALON_MAX_CHUNKS / 8u];

static TalonRoundState s_state;

// PSRAM candidate buffer pointer — allocated once in talon_init()
static uint8_t* s_candidate_buf = nullptr;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

static inline void chunk_mark_received(uint32_t idx) {
    if (idx < TALON_MAX_CHUNKS) {
        s_chunk_received[idx >> 3u] |= (uint8_t)(1u << (idx & 7u));
    }
}

static inline bool chunk_was_received(uint32_t idx) {
    if (idx >= TALON_MAX_CHUNKS) return false;
    return (s_chunk_received[idx >> 3u] & (uint8_t)(1u << (idx & 7u))) != 0u;
}

static void reset_round_state() {
    s_state.round_id         = 0u;
    s_state.total_chunks     = 0u;
    s_state.received_chunks  = 0u;
    s_state.expected_crc32   = 0u;
    s_state.apply_after_ts   = 0u;
    s_state.write_cursor     = 0u;
    s_state.bytes_per_chunk  = 0u;
    s_state.crc_ok           = false;
    s_state.poisoned         = false;
    memset(s_chunk_received, 0, sizeof(s_chunk_received));
    if (s_candidate_buf != nullptr) {
        // Zero the candidate slot to avoid stale bytes influencing CRC
        memset(s_candidate_buf, 0, TALON_CANDIDATE_BUF_SIZE);
    }
}

// Big-endian decode helpers (no alignment assumption)
static inline uint16_t be16(const uint8_t* p) {
    return (uint16_t)(((uint16_t)p[0] << 8) | p[1]);
}
static inline uint32_t be32(const uint8_t* p) {
    return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) |
           ((uint32_t)p[2] << 8)  |  (uint32_t)p[3];
}

// ---------------------------------------------------------------------------
// talon_init
// ---------------------------------------------------------------------------
void talon_init() {
    memset(&s_state, 0, sizeof(s_state));
    memset(s_chunk_received, 0, sizeof(s_chunk_received));

    if (s_candidate_buf == nullptr) {
#ifndef UNIT_TEST
        // Allocate from PSRAM on embedded target
        s_candidate_buf = (uint8_t*)heap_caps_malloc(
            TALON_CANDIDATE_BUF_SIZE, MALLOC_CAP_SPIRAM);
#else
        // Native test build: allocate from regular heap
        s_candidate_buf = (uint8_t*)malloc(TALON_CANDIDATE_BUF_SIZE);
#endif
    }

    if (s_candidate_buf != nullptr) {
        memset(s_candidate_buf, 0, TALON_CANDIDATE_BUF_SIZE);
    }
}

// ---------------------------------------------------------------------------
// talon_handle_chunk
// ---------------------------------------------------------------------------
void talon_handle_chunk(const uint8_t* buf, size_t len) {
    // Minimum: 1 msg_type + 28-byte header = 29 bytes
    if (buf == nullptr || len < (TALON_CHUNK_HEADER_BYTES + 1u)) return;
    if (buf[0] != 0x1Bu) return;   // not MSG_TYPE_FED_UPDATE
    if (s_candidate_buf == nullptr) return;

    const uint8_t* hdr = buf + 1u;   // header starts at byte 1

    uint16_t round_id      = be16(hdr + 0u);
    uint32_t total_chunks  = be32(hdr + 2u);
    uint32_t chunk_idx     = be32(hdr + 6u);
    uint32_t delta_crc32   = be32(hdr + 10u);
    uint32_t apply_ts      = be32(hdr + 14u);
    // bytes 18..26 = reserved; ignored

    // Guard: sanity-check chunk_idx and total_chunks
    if (total_chunks == 0u || total_chunks > TALON_MAX_CHUNKS) return;
    if (chunk_idx >= total_chunks) return;

    // If round_id changed, start a new round
    if (round_id != s_state.round_id || s_state.total_chunks == 0u) {
        reset_round_state();
        s_state.round_id       = round_id;
        s_state.total_chunks   = total_chunks;
        s_state.expected_crc32 = delta_crc32;
        s_state.apply_after_ts = apply_ts;
    }

    // Duplicate suppression
    if (chunk_was_received(chunk_idx)) return;

    // Poison state: this round has already failed CRC; reject until next round
    if (s_state.poisoned) return;

    // Payload bytes immediately follow the 28-byte header (buf[1..28] = hdr,
    // buf[29..] = payload)
    const uint8_t* payload = buf + 1u + TALON_CHUNK_HEADER_BYTES;
    size_t         payload_len = len - (1u + TALON_CHUNK_HEADER_BYTES);
    if (payload_len > TALON_CHUNK_MAX_PAYLOAD) {
        payload_len = TALON_CHUNK_MAX_PAYLOAD;
    }

    // Record the actual chunk payload size from the first chunk received.
    // All chunks in a round must be the same size (except possibly the last),
    // so we use this as the per-chunk stride for offset calculation.
    // This ensures the reassembled buffer is contiguous and the CRC is
    // computed over exactly the assembled data bytes (not padded to max stride).
    if (s_state.bytes_per_chunk == 0u && payload_len > 0u) {
        s_state.bytes_per_chunk = (uint32_t)payload_len;
    }
    uint32_t stride = (s_state.bytes_per_chunk > 0u) ? s_state.bytes_per_chunk
                                                      : TALON_CHUNK_MAX_PAYLOAD;

    // Write payload into the candidate buffer at the correct offset.
    // Offset = chunk_idx * actual_bytes_per_chunk (contiguous packing).
    uint32_t offset = chunk_idx * stride;
    if (offset + payload_len > TALON_CANDIDATE_BUF_SIZE) return;

    memcpy(s_candidate_buf + offset, payload, payload_len);

    chunk_mark_received(chunk_idx);
    s_state.received_chunks++;

    // Track the write_cursor to know total assembled bytes
    uint32_t end = offset + (uint32_t)payload_len;
    if (end > s_state.write_cursor) {
        s_state.write_cursor = end;
    }

    // Check if all chunks have been received
    if (s_state.received_chunks >= s_state.total_chunks) {
        // Verify CRC32 over assembled data (write_cursor bytes = contiguous delta)
        uint32_t computed = crc32_update(0u, s_candidate_buf, s_state.write_cursor);
        if (computed == s_state.expected_crc32) {
            s_state.crc_ok = true;
        } else {
            // CRC mismatch — poison this round; T2 will need to re-transmit
            s_state.crc_ok   = false;
            s_state.poisoned = true;
        }
    }
}

// ---------------------------------------------------------------------------
// talon_round_complete
// ---------------------------------------------------------------------------
bool talon_round_complete() {
    return s_state.crc_ok && !s_state.poisoned;
}

// ---------------------------------------------------------------------------
// talon_apply_round
// ---------------------------------------------------------------------------
bool talon_apply_round() {
    if (!talon_round_complete()) return false;
    if (s_candidate_buf == nullptr) return false;

    // Honour apply_after_ts: if T2 requested a deferred apply, refuse until
    // the device clock has advanced past that timestamp.
    // On embedded targets millis()/1000 gives a rough epoch offset.
    if (s_state.apply_after_ts != 0u) {
#ifndef UNIT_TEST
        extern unsigned long millis();
        uint32_t now_s = (uint32_t)(millis() / 1000u);
        if (now_s < s_state.apply_after_ts) return false;
#endif
        // Native test build: skip timestamp check (test controls timing)
    }

    // Signature verification: delegate to ota_signed_verify_chunk when
    // compiled for the embedded target (Ed25519).  In v1 we do not embed
    // a signature in the FED_UPDATE packet; ota_signed_verify_chunk is
    // called with nullptr for sig64, which returns false if a key has been
    // provisioned.  Fall back to CRC-only when no key is provisioned.
#ifndef UNIT_TEST
    {
        uint8_t pk[ED25519_PUBKEY_BYTES];
        bool has_key = ota_signed_get_pubkey(pk);
        if (has_key) {
            // v1 payload carries no Ed25519 signature — reject if key is set.
            // The integrator must either:
            //   a) not provision a key until the v2 signed FED_UPDATE is deployed, OR
            //   b) set a flag to bypass signature check for federated updates.
            // For safety in v1 we reject to avoid silently applying unsigned deltas.
            return false;
        }
        // No key provisioned → CRC-only apply (already verified in handle_chunk)
    }
#endif

    // Atomic apply: overwrite the active model manager candidate slot.
    // In v1 we overwrite model slot 0 (drone_classifier — most likely to
    // receive federated updates) if its size fits the assembled delta.
    // The model_manager_commit_slot validates the CRC it is passed against
    // its own internal CRC before marking the slot active; passing our
    // expected_crc32 ensures it re-validates.
#ifndef UNIT_TEST
    extern bool model_manager_commit_slot(uint8_t slot_id, uint32_t expected_crc32);
    bool applied = model_manager_commit_slot(0u, s_state.expected_crc32);
#else
    // Native test stub: always succeed (test inspects s_candidate_buf directly)
    bool applied = true;
#endif

    if (applied) {
        reset_round_state();
    }
    return applied;
}

// ---------------------------------------------------------------------------
// talon_current_round
// ---------------------------------------------------------------------------
uint16_t talon_current_round() {
    return s_state.round_id;
}

// ---------------------------------------------------------------------------
// v2 stubs — NOT implemented in v1
// ---------------------------------------------------------------------------
bool talon_collect_local_anchors() {
    // v2 deferred: requires access to labeled local detection window.
    return false;
}

bool talon_compute_gradient() {
    // v2 deferred: requires int8 backprop through model_manager inference path.
    return false;
}
