/*  fullmode_t1.cpp — firmware glue: SPIFFS discovery + cactus_fullmode dispatch
 *
 *  Scans /models/ on SPIFFS for *.cactus files at init time, loads each via
 *  cactus_fm_load(), and exposes fullmode_t1_infer_all_resident() to mirror
 *  the model_manager API so results slot into correlator_feed() unchanged.
 *
 *  This file has ZERO changes to existing sources (model_manager, main, config).
 */

#include "fullmode_t1.h"
#include "cactus_fullmode.h"

#include <string.h>
#include <math.h>

#ifdef ARDUINO
#include <Arduino.h>
#include <LittleFS.h>
#include <esp_heap_caps.h>
#else
// Host-side stubs so this module (previously buildable ONLY on-device,
// because fullmode_t1_init()'s SPIFFS directory scan unconditionally used
// LittleFS/File/Serial/micros()) can be exercised by a native Unity test --
// same pattern used throughout this sweep. The SPIFFS scan itself has no
// native equivalent (there is no filesystem to scan), so under UNIT_TEST
// fullmode_t1_init() is a no-op state reset and tests populate slots
// directly via fullmode_t1_test_seed_slot() instead -- what's actually
// worth testing here is fullmode_t1_infer_all_resident()'s dispatch/
// softmax/argmax math, not the SPIFFS glue.
#include <cstdio>
#include <cstdarg>
#include <cstdint>
struct NativeSerialStub {
    void printf(const char* fmt, ...) {
        va_list args;
        va_start(args, fmt);
        vprintf(fmt, args);
        va_end(args);
    }
};
static NativeSerialStub Serial;
#ifndef STUB_MICROS_PROVIDED
static uint32_t stub_micros_val = 0;
static inline uint32_t micros() { return stub_micros_val; }
#endif
#endif // ARDUINO

static const char* TAG = "FM_T1";

// ── Slot tracking ──────────────────────────────────────────────────────────
struct FmSlotInfo {
    bool  active;
    int   cactus_slot;   // slot id from cactus_fm_load()
    char  path[48];
    int   input_dim;
    int   output_dim;
    uint32_t infer_count;
    uint64_t total_us;
};

static FmSlotInfo s_fm[FULLMODE_T1_MAX_SLOTS];
static int        s_fm_count = 0;
static bool       s_spiffs_ok = false;

// ── Helpers ────────────────────────────────────────────────────────────────

#ifdef ARDUINO
static bool ensure_spiffs() {
    if (s_spiffs_ok) return true;
    if (!LittleFS.begin(true)) {
        Serial.printf("[%s] SPIFFS mount failed\n", TAG);
        return false;
    }
    s_spiffs_ok = true;
    return true;
}

static inline bool ends_with_cactus(const char* name) {
    size_t n = strlen(name);
    return n > 7 && strcmp(name + n - 7, ".cactus") == 0;
}

// Read the declared input/output dims back from a loaded slot by probing
// a zero-vector. We just cache the dims from the SPIFFS file via a tiny
// dry-run approach: because we can't query cactus_fm dimensions directly
// through the public API without an infer call, we keep a dedicated
// two-field record. We set them when we parse via a trial infer of zeros
// before committing the slot.
//
// Simpler approach: attempt infer with increasing probe sizes from 1..512
// and record the first that succeeds. Works for any MLP <= 512 inputs.
static void probe_slot_dims(int cactus_slot, int* in_dim_out, int* out_dim_out) {
    static float probe_in[512];
    static float probe_out[32];
    memset(probe_in, 0, sizeof(probe_in));

    // Try common power-of-two input dims first (fast path)
    static const int candidates[] = { 128, 64, 256, 32, 512, 16, 8, 1 };
    for (size_t ci = 0; ci < sizeof(candidates) / sizeof(candidates[0]); ci++) {
        int idim = candidates[ci];
        // Try common output dims
        for (int odim = 32; odim >= 1; odim /= 2) {
            if (cactus_fm_infer(cactus_slot,
                                probe_in, (size_t)idim,
                                probe_out, (size_t)odim)) {
                *in_dim_out  = idim;
                *out_dim_out = odim;
                return;
            }
        }
        // Try odd output dims up to 16
        for (int odim = 1; odim <= 16; odim++) {
            if (cactus_fm_infer(cactus_slot,
                                probe_in, (size_t)idim,
                                probe_out, (size_t)odim)) {
                *in_dim_out  = idim;
                *out_dim_out = odim;
                return;
            }
        }
    }
    // Fallback: unknown dims
    *in_dim_out  = 0;
    *out_dim_out = 0;
}

// ── Public: fullmode_t1_init ───────────────────────────────────────────────
void fullmode_t1_init(void) {
    memset(s_fm, 0, sizeof(s_fm));
    s_fm_count = 0;

    if (!cactus_fm_init()) {
        Serial.printf("[%s] cactus_fm_init failed — no .cactus models will run\n", TAG);
        return;
    }

    if (!ensure_spiffs()) return;

    // Walk /models/ directory looking for *.cactus files
    File dir = LittleFS.open("/models");
    if (!dir || !dir.isDirectory()) {
        Serial.printf("[%s] /models/ not found on SPIFFS — no .cactus files\n", TAG);
        return;
    }

    File entry = dir.openNextFile();
    while (entry && s_fm_count < FULLMODE_T1_MAX_SLOTS) {
        if (!entry.isDirectory()) {
            const char* name = entry.name();
            if (ends_with_cactus(name)) {
                // Build full path (LittleFS.open on ESP32 Arduino 2.x needs the
                // full path; entry.name() already includes it on IDF >=4.x,
                // but we normalise to be safe)
                char full_path[56];
                if (name[0] == '/') {
                    strncpy(full_path, name, sizeof(full_path) - 1);
                } else {
                    snprintf(full_path, sizeof(full_path), "/models/%s", name);
                }
                full_path[sizeof(full_path) - 1] = '\0';

                int cslot = cactus_fm_load(full_path);
                if (cslot >= 0) {
                    FmSlotInfo& fi = s_fm[s_fm_count];
                    fi.active       = true;
                    fi.cactus_slot  = cslot;
                    strncpy(fi.path, full_path, sizeof(fi.path) - 1);
                    fi.path[sizeof(fi.path) - 1] = '\0';
                    fi.infer_count  = 0;
                    fi.total_us     = 0;

                    // Probe actual dims so infer can validate lengths
                    probe_slot_dims(cslot, &fi.input_dim, &fi.output_dim);

                    Serial.printf("[%s] .cactus slot %d: %s  in=%d out=%d\n",
                                  TAG, s_fm_count, full_path,
                                  fi.input_dim, fi.output_dim);
                    s_fm_count++;
                } else {
                    Serial.printf("[%s] Failed to load %s\n", TAG, full_path);
                }
            }
        }
        entry.close();
        entry = dir.openNextFile();
    }
    dir.close();

    Serial.printf("[%s] init complete: %d .cactus model(s) loaded, PSRAM %u KB used\n",
                  TAG, s_fm_count,
                  (unsigned)(cactus_fm_psram_used() / 1024));
}
#else  // !ARDUINO
// Native builds have no SPIFFS to scan -- tests populate slots directly via
// fullmode_t1_test_seed_slot() below instead. What's worth testing here is
// fullmode_t1_infer_all_resident()'s dispatch/softmax/argmax math, not the
// filesystem glue.
void fullmode_t1_init(void) {
    memset(s_fm, 0, sizeof(s_fm));
    s_fm_count = 0;
    cactus_fm_init();
}
#endif // ARDUINO

// ── Public: fullmode_t1_infer_all_resident ────────────────────────────────
int fullmode_t1_infer_all_resident(const float*          input_data,
                                   uint16_t              input_len,
                                   ModelInferenceResult* results,
                                   int                   max_results)
{
    int count = 0;

    // Static output scratch — max MODEL_MAX_OUTPUT_CLASSES floats
    // s_out holds raw model output (logits or probs); s_probs holds softmax result.
    static float s_out[MODEL_MAX_OUTPUT_CLASSES];
    static float s_probs[MODEL_MAX_OUTPUT_CLASSES];

    for (int i = 0; i < s_fm_count && count < max_results; i++) {
        FmSlotInfo& fi = s_fm[i];
        if (!fi.active) continue;
        if (fi.input_dim > 0 && (int)input_len != fi.input_dim) continue;

        ModelInferenceResult& r = results[count];
        memset(&r, 0, sizeof(r));
        r.valid = false;

        memset(s_out, 0, sizeof(s_out));

        uint32_t t0 = micros();
        bool ok = cactus_fm_infer(fi.cactus_slot,
                                  input_data, (size_t)input_len,
                                  s_out, MODEL_MAX_OUTPUT_CLASSES);
        uint32_t elapsed_us = micros() - t0;

        if (!ok) continue;

        uint8_t num_out = (fi.output_dim > 0 && fi.output_dim <= MODEL_MAX_OUTPUT_CLASSES)
                          ? (uint8_t)fi.output_dim
                          : (uint8_t)MODEL_MAX_OUTPUT_CLASSES;

        // Apply softmax over raw logits so confidence is a true probability in [0,1].
        // Models that already have a softmax output layer are unaffected (softmax(softmax(x))
        // still peaks at the same class and confidence stays near 1.0).
        {
            float mx = s_out[0];
            for (int c = 1; c < num_out; c++) if (s_out[c] > mx) mx = s_out[c];
            float sum = 0.0f;
            for (int c = 0; c < num_out; c++) {
                s_probs[c] = expf(s_out[c] - mx);
                sum += s_probs[c];
            }
            float inv = (sum > 0.0f) ? 1.0f / sum : 1.0f;
            for (int c = 0; c < num_out; c++) s_probs[c] *= inv;
        }

        // Find argmax in probability space
        float max_p = s_probs[0];
        int   max_i = 0;
        for (int c = 1; c < num_out; c++) {
            if (s_probs[c] > max_p) { max_p = s_probs[c]; max_i = c; }
        }

        // Encode slot_id with backend sentinel so correlator can distinguish
        r.slot_id     = (uint8_t)(MODEL_BACKEND_CACTUS | (uint8_t)i);
        r.class_label = (uint8_t)max_i;
        r.confidence  = max_p;
        r.num_classes = num_out;
        r.latency_us  = elapsed_us;
        r.valid       = true;

        for (int c = 0; c < num_out; c++)
            r.probabilities[c] = s_probs[c];

        fi.infer_count++;
        fi.total_us += elapsed_us;

        count++;
    }

    return count;
}

// ── Public: fullmode_t1_psram_used ────────────────────────────────────────
uint32_t fullmode_t1_psram_used(void) {
    return (uint32_t)cactus_fm_psram_used();
}

#ifdef UNIT_TEST
void fullmode_t1_test_reset(void) {
    memset(s_fm, 0, sizeof(s_fm));
    s_fm_count = 0;
}

int fullmode_t1_test_seed_slot(int cactus_slot, int input_dim, int output_dim) {
    if (s_fm_count >= FULLMODE_T1_MAX_SLOTS) return -1;
    FmSlotInfo& fi = s_fm[s_fm_count];
    fi.active      = true;
    fi.cactus_slot = cactus_slot;
    fi.input_dim   = input_dim;
    fi.output_dim  = output_dim;
    fi.infer_count = 0;
    fi.total_us    = 0;
    snprintf(fi.path, sizeof(fi.path), "/models/test_slot_%d.cactus", s_fm_count);
    return s_fm_count++;
}
#endif
