// ============================================================================
// PEC Feedback Client — implementation
// See pec_feedback.h for architecture notes.
// ============================================================================

#include "pec_feedback.h"

#include <cstring>

// millis() shim
// - On Arduino: pulled in from <Arduino.h>
// - In UNIT_TEST (native env): the test file defines millis() once and owns
//   the shared clock.  We just forward-declare it here.
// - Otherwise: fallback stub.
#ifdef ARDUINO
#include <Arduino.h>
#elif defined(UNIT_TEST)
#include <cstdint>
extern uint32_t millis();   // provided by test_spoof_pec.cpp
#else
#include <cstdint>
static uint32_t s_fallback_millis_pec = 0;
static inline uint32_t millis() { return s_fallback_millis_pec; }
#endif

// ---------------------------------------------------------------------------
// LRU cache — 4 entries, static memory only
// ---------------------------------------------------------------------------

struct PecEntry {
    uint32_t track_id;
    uint8_t  num_classes;
    float    probs[MODEL_MAX_OUTPUT_CLASSES]; // converted from q8 on insert
    uint32_t ts_ms;
    bool     valid;
};

static PecEntry s_pec_cache[PEC_FEEDBACK_CACHE_SIZE];

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

static PecEntry* pec_find(uint32_t track_id) {
    for (uint8_t i = 0; i < PEC_FEEDBACK_CACHE_SIZE; i++) {
        if (s_pec_cache[i].valid && s_pec_cache[i].track_id == track_id) {
            return &s_pec_cache[i];
        }
    }
    return nullptr;
}

static uint8_t pec_evict_slot(uint32_t now_ms) {
    // Prefer invalid slot
    for (uint8_t i = 0; i < PEC_FEEDBACK_CACHE_SIZE; i++) {
        if (!s_pec_cache[i].valid) return i;
    }
    // Prefer expired slot
    for (uint8_t i = 0; i < PEC_FEEDBACK_CACHE_SIZE; i++) {
        if ((now_ms - s_pec_cache[i].ts_ms) >= PEC_FEEDBACK_ENTRY_TTL_MS) return i;
    }
    // Evict oldest by ts_ms
    uint8_t oldest = 0;
    uint32_t oldest_ts = s_pec_cache[0].ts_ms;
    for (uint8_t i = 1; i < PEC_FEEDBACK_CACHE_SIZE; i++) {
        if (s_pec_cache[i].ts_ms < oldest_ts) {
            oldest_ts = s_pec_cache[i].ts_ms;
            oldest = i;
        }
    }
    return oldest;
}

// Convert a q8 byte array to float probabilities.
// q8 encoding: prob_f = q8_value / 255.0f
static void q8_to_float(const uint8_t* q8, float* out, uint8_t n) {
    for (uint8_t i = 0; i < n; i++) {
        out[i] = (float)q8[i] / 255.0f;
    }
}

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

void pec_feedback_init() {
    memset(s_pec_cache, 0, sizeof(s_pec_cache));
}

void pec_feedback_handle_packet(const uint8_t* buf, size_t len) {
    if (buf == nullptr || len < PEC_FEEDBACK_MSG_MIN_LEN) return;

    // Parse fields manually from the raw buffer to avoid packed-struct
    // alignment/ABI differences between Arduino-ESP32 GCC and desktop compilers.
    //
    // Wire layout starting at buf[0]:
    //   [0]      msg_type    (already checked by caller's switch)
    //   [1]      num_classes
    //   [2..17]  class_probs_q8[16]
    //   [18..21] source_track_id (uint32_t LE)
    const uint8_t* p = buf + 1;                          // skip msg_type

    uint8_t  num_classes = p[0];
    const uint8_t* probs_q8 = p + 1;                     // 16 bytes at p[1..16]
    uint32_t track_id;
    memcpy(&track_id, p + 1 + MODEL_MAX_OUTPUT_CLASSES, sizeof(uint32_t));

    // Validate num_classes in range
    if (num_classes == 0 || num_classes > MODEL_MAX_OUTPUT_CLASSES) return;

    uint32_t now = millis();

    PecEntry* existing = pec_find(track_id);
    if (!existing) {
        uint8_t slot = pec_evict_slot(now);
        existing = &s_pec_cache[slot];
        existing->valid    = true;
        existing->track_id = track_id;
    }

    existing->ts_ms      = now;
    existing->num_classes = num_classes;
    q8_to_float(probs_q8, existing->probs, num_classes);

    // Zero remaining slots for clean reads
    for (uint8_t i = num_classes; i < MODEL_MAX_OUTPUT_CLASSES; i++) {
        existing->probs[i] = 0.0f;
    }
}

bool pec_feedback_get_latest(uint32_t track_id,
                             float*   probs_out,
                             uint8_t* n_out,
                             uint32_t max_age_ms) {
    if (probs_out == nullptr || n_out == nullptr) return false;

    uint32_t now = millis();
    PecEntry* e = pec_find(track_id);
    if (!e) return false;

    uint32_t age = now - e->ts_ms;
    if (age >= PEC_FEEDBACK_ENTRY_TTL_MS || age > max_age_ms) {
        if (age >= PEC_FEEDBACK_ENTRY_TTL_MS) e->valid = false;
        return false;
    }

    memcpy(probs_out, e->probs, e->num_classes * sizeof(float));
    *n_out = e->num_classes;
    return true;
}
