// ============================================================================
// markov_scorer.cpp — T1 RID Behavioral Anomaly Scorer (Markov chain)
//
// Static allocation only.  No new/malloc in any hot path.
// Compiles under:
//   - Arduino-ESP32 (ARDUINO defined) — uses SPIFFS for file I/O
//   - PlatformIO native env (UNIT_TEST defined) — stubs SPIFFS with stdio
//
// Heavy STL banned.  Allowed headers: <cstdint>, <cstring>, <cmath>.
// ============================================================================

#include "markov_scorer.h"
#include <cstdint>
#include <cstring>
#include <cmath>

// ---------------------------------------------------------------------------
// File-I/O shim: SPIFFS (Arduino) vs stdio (native/UNIT_TEST)
// ---------------------------------------------------------------------------
#ifdef ARDUINO
#  include "LittleFS.h"
   // Wrapper type for consistent cross-env file handle usage
   struct MkFile {
       File f;
       bool ok() const { return (bool)f; }
       bool read_bytes(void* buf, size_t n) {
           return (size_t)f.read(reinterpret_cast<uint8_t*>(buf), n) == n;
       }
       void close() { f.close(); }
   };
   static MkFile mk_open(const char* path) {
       MkFile mf;
       mf.f = LittleFS.open(path, "r");
       return mf;
   }
#else
#  include <cstdio>
   struct MkFile {
       FILE* f = nullptr;
       bool ok() const { return f != nullptr; }
       bool read_bytes(void* buf, size_t n) {
           return std::fread(buf, 1, n, f) == n;
       }
       void close() { if (f) { std::fclose(f); f = nullptr; } }
   };
   static MkFile mk_open(const char* path) {
       MkFile mf;
       mf.f = std::fopen(path, "rb");
       return mf;
   }
#endif

// ---------------------------------------------------------------------------
// Q1.15 helpers
//   NLL values are stored as int16_t: q15 = round(nll * 32767 / NLL_SCALE)
//   NLL_SCALE caps the representable range at ~8.0 nats (plenty — benign
//   traffic sits well below 2.0, anomalous spikes to 6-8+).
// ---------------------------------------------------------------------------
static constexpr float MK_NLL_SCALE = 8.0f;  // maps [0..8] nats → int16 range

static inline float q15_to_nll(int16_t q) {
    // clamp at 0 (no negative NLL)
    if (q <= 0) return 0.0f;
    return (static_cast<float>(q) / 32767.0f) * MK_NLL_SCALE;
}

// Not needed in firmware hot path but kept for symmetry / test tooling.
static inline int16_t nll_to_q15(float nll) {
    if (nll <= 0.0f) return 0;
    float v = (nll / MK_NLL_SCALE) * 32767.0f;
    if (v > 32767.0f) return 32767;
    return static_cast<int16_t>(v + 0.5f);
}

// ---------------------------------------------------------------------------
// Transition matrix (static SRAM, row-major)
//   Row = "from" state, Col = "to" state, both are state_index().
//   Default fill = MK_UNKNOWN_NLL_Q15 (max int16, signals OOV).
// ---------------------------------------------------------------------------
static int16_t s_matrix[MK_N_STATES][MK_N_STATES];
static bool    s_model_loaded = false;

// State index: (event_type, iat_bucket) → flat index
static inline int state_index(uint8_t ev, uint8_t iat) {
    return static_cast<int>(ev) * MK_IAT_BUCKETS + static_cast<int>(iat);
}

// ---------------------------------------------------------------------------
// IAT bucket — mirrors Python _iat_bucket with 5-bucket scheme
//   0 : < 100 ms
//   1 : 100 – 500 ms
//   2 : 500 ms – 2000 ms
//   3 : 2000 – 10000 ms
//   4 : > 10000 ms
// ---------------------------------------------------------------------------
static const uint32_t IAT_EDGES_MS[4] = { 100u, 500u, 2000u, 10000u };

static inline uint8_t iat_bucket(uint32_t dt_ms) {
    for (int i = 0; i < 4; ++i) {
        if (dt_ms < IAT_EDGES_MS[i]) return static_cast<uint8_t>(i);
    }
    return 4u;
}

// ---------------------------------------------------------------------------
// Per-track state
// ---------------------------------------------------------------------------
struct TrackState {
    // Rolling window: last MK_WINDOW_LEN NLL values (as float, ~128 bytes/track)
    float  window[MK_WINDOW_LEN];
    uint8_t window_head;   // next write index (circular)
    uint8_t window_count;  // how many entries are valid

    // Last observed state index (for transition lookup)
    int8_t  last_state;    // -1 = no prior state (start of track)

    // Last event timestamp in ms
    uint32_t last_ts_ms;

    // Whether this slot is active (has been reset/initialized)
    bool active;
};

static TrackState s_tracks[MK_MAX_TRACKS];

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------
static void reset_track_slot(uint8_t slot) {
    TrackState& t = s_tracks[slot];
    std::memset(t.window, 0, sizeof(t.window));
    t.window_head  = 0;
    t.window_count = 0;
    t.last_state   = -1;
    t.last_ts_ms   = 0;
    t.active       = true;
}

static void push_nll(uint8_t slot, float nll) {
    TrackState& t = s_tracks[slot];
    t.window[t.window_head] = nll;
    t.window_head = static_cast<uint8_t>((t.window_head + 1u) % MK_WINDOW_LEN);
    if (t.window_count < MK_WINDOW_LEN) ++t.window_count;
}

static float rolling_nll(uint8_t slot) {
    const TrackState& t = s_tracks[slot];
    if (t.window_count == 0) return 0.0f;
    float sum = 0.0f;
    for (uint8_t i = 0; i < t.window_count; ++i) sum += t.window[i];
    return sum / static_cast<float>(t.window_count);
}

static void observe_on_slot(uint8_t slot, uint8_t event_type, uint32_t ts_ms) {
    if (slot >= MK_MAX_TRACKS) return;
    if (!s_model_loaded)       return;

    TrackState& t = s_tracks[slot];
    if (!t.active) reset_track_slot(slot);

    // Clamp out-of-vocabulary event types to produce a defined OOV NLL
    if (event_type >= MK_EVENT_COUNT) {
        push_nll(slot, q15_to_nll(MK_UNKNOWN_NLL_Q15));
        t.last_ts_ms = ts_ms;
        return;
    }

    // Compute inter-arrival time
    uint32_t dt_ms = (t.last_state < 0) ? 0u :
                     (ts_ms >= t.last_ts_ms ? ts_ms - t.last_ts_ms : 0u);

    uint8_t iat  = iat_bucket(dt_ms);
    int     to   = state_index(event_type, iat);
    int     from = t.last_state; // -1 means "START"

    float nll;
    if (from < 0) {
        // First event on track: treat as unconditional NLL from start state.
        // We encode the "from-START" row as row 0 in the matrix (by
        // build_baseline.py convention).  Use state_index(0,0)=0 as proxy.
        // This gives a reasonable prior for the first event.
        nll = q15_to_nll(s_matrix[0][to]);
    } else {
        nll = q15_to_nll(s_matrix[from][to]);
    }

    push_nll(slot, nll);
    t.last_state = static_cast<int8_t>(to);
    t.last_ts_ms = ts_ms;
}

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

bool markov_init(const char* baseline_path) {
    s_model_loaded = false;

    // Fill matrix with max-NLL (OOV sentinel) before loading
    for (int i = 0; i < MK_N_STATES; ++i)
        for (int j = 0; j < MK_N_STATES; ++j)
            s_matrix[i][j] = static_cast<int16_t>(MK_UNKNOWN_NLL_Q15);

    // Initialise all track slots
    for (int s = 0; s < MK_MAX_TRACKS; ++s) reset_track_slot(static_cast<uint8_t>(s));

    MkFile f = mk_open(baseline_path);
    if (!f.ok()) return false;

    // Read and validate header (8 bytes)
    uint8_t header[8];
    if (!f.read_bytes(header, 8)) { f.close(); return false; }

    // magic (LE uint32)
    uint32_t magic =  static_cast<uint32_t>(header[0])
                   | (static_cast<uint32_t>(header[1]) << 8)
                   | (static_cast<uint32_t>(header[2]) << 16)
                   | (static_cast<uint32_t>(header[3]) << 24);
    if (magic != MK_MAGIC) { f.close(); return false; }

    // version (LE uint16) — accept 1
    // uint16_t version = header[4] | (header[5] << 8);  // currently unused

    // n_states (LE uint16) — must match compile-time constant
    uint16_t n_states = static_cast<uint16_t>(header[6])
                      | static_cast<uint16_t>(header[7] << 8);
    if (n_states != static_cast<uint16_t>(MK_N_STATES)) { f.close(); return false; }

    // Read transition matrix: n_states * n_states int16_t values (LE)
    // Read row by row to avoid a large stack buffer
    for (int row = 0; row < MK_N_STATES; ++row) {
        uint8_t row_buf[MK_N_STATES * 2];
        if (!f.read_bytes(row_buf, MK_N_STATES * 2)) { f.close(); return false; }
        for (int col = 0; col < MK_N_STATES; ++col) {
            int16_t v = static_cast<int16_t>(
                            static_cast<uint16_t>(row_buf[col * 2])
                          | static_cast<uint16_t>(row_buf[col * 2 + 1] << 8));
            s_matrix[row][col] = v;
        }
    }

    f.close();
    s_model_loaded = true;
    return true;
}

void markov_observe_event(uint8_t event_type, uint32_t timestamp_ms) {
    observe_on_slot(0, event_type, timestamp_ms);
}

void markov_observe_event_track(uint8_t track_slot,
                                uint8_t event_type,
                                uint32_t timestamp_ms) {
    observe_on_slot(track_slot, event_type, timestamp_ms);
}

float markov_current_nll(void) {
    return rolling_nll(0);
}

float markov_current_nll_track(uint8_t track_slot) {
    if (track_slot >= MK_MAX_TRACKS) return 0.0f;
    return rolling_nll(track_slot);
}

bool markov_is_anomalous(float threshold) {
    return markov_current_nll() > threshold;
}

void markov_reset_track(uint32_t track_id) {
    uint8_t slot = static_cast<uint8_t>(track_id % MK_MAX_TRACKS);
    reset_track_slot(slot);
}

// ---------------------------------------------------------------------------
// Anomaly threshold + telemetry emission
// ---------------------------------------------------------------------------
static float s_anomaly_threshold = 4.0f;

void  markov_set_anomaly_threshold(float t) { if (t > 0.0f && t < MK_NLL_SCALE) s_anomaly_threshold = t; }
float markov_get_anomaly_threshold(void)    { return s_anomaly_threshold; }

#ifdef ARDUINO
#include "mesh_publish.h"
#include "config.h"

bool markov_publish_nll(uint32_t sensor_id, uint8_t track_slot, float threshold) {
    if (track_slot >= MK_MAX_TRACKS) return false;
    float nll = rolling_nll(track_slot);
    float t   = (threshold > 0.0f) ? threshold : s_anomaly_threshold;

    uint8_t pkt[16];
    std::memset(pkt, 0, sizeof(pkt));
    pkt[0] = (uint8_t)MSG_TYPE_BEHAVIORAL_NLL;
    pkt[1] = (uint8_t)(sensor_id        & 0xFF);
    pkt[2] = (uint8_t)((sensor_id >> 8) & 0xFF);
    pkt[3] = (uint8_t)((sensor_id >> 16)& 0xFF);
    pkt[4] = (uint8_t)((sensor_id >> 24)& 0xFF);
    uint32_t ts = (uint32_t)millis();
    pkt[5] = (uint8_t)(ts        & 0xFF);
    pkt[6] = (uint8_t)((ts >> 8) & 0xFF);
    pkt[7] = (uint8_t)((ts >> 16)& 0xFF);
    pkt[8] = (uint8_t)((ts >> 24)& 0xFF);
    pkt[9]  = track_slot;
    float nll_norm = nll / MK_NLL_SCALE; if (nll_norm < 0) nll_norm = 0; if (nll_norm > 1) nll_norm = 1;
    pkt[10] = (uint8_t)(nll_norm * 255.0f);
    float thr_norm = t / MK_NLL_SCALE; if (thr_norm < 0) thr_norm = 0; if (thr_norm > 1) thr_norm = 1;
    pkt[11] = (uint8_t)(thr_norm * 255.0f);
    pkt[12] = (nll > t) ? 1 : 0;
    return mesh_publish_raw(pkt, sizeof(pkt));
}
#else
bool markov_publish_nll(uint32_t, uint8_t, float) { return false; }
#endif
