// ============================================================================
// test/test_fullmode_t1/test_fullmode_t1.cpp
//
// Unity tests for src/fullmode_t1.{h,cpp} -- firmware glue that dispatches
// inference across loaded .cactus (hand-rolled MLP) models and feeds results
// into the same ModelInferenceResult shape model_manager/correlator use.
//
// Previously zero native coverage: fullmode_t1_init() unconditionally used
// LittleFS/File/Serial/micros() (SPIFFS directory scan), so the whole file
// could only build for the ESP32 target (there IS an on-device smoke test,
// test/test_fullmode/test_fullmode_ondevice.cpp, for cactus_fullmode.cpp
// itself -- a different file -- run via `pio test -e heltec_wifi_lora_v4`).
//
// Applied the same #ifdef ARDUINO / native-stub pattern used throughout this
// sweep. The SPIFFS scan has no native equivalent (no filesystem to scan),
// so under UNIT_TEST fullmode_t1_init() is just a state reset, and a new
// test-only hook (fullmode_t1_test_seed_slot()) registers slots directly --
// what's actually worth testing here is fullmode_t1_infer_all_resident()'s
// dispatch/softmax/argmax/backend-sentinel-encoding math, not the filesystem
// glue. cactus_fm_init/load/infer/psram_used (declared in cactus_fullmode.h,
// defined in the SEPARATE cactus_fullmode.cpp which this test does NOT
// include) are stubbed here as controllable test doubles.
//
// Run with:
//   pio test -e native -f test_fullmode_t1
// ============================================================================

#include <unity.h>
#include <cstring>
#include <vector>
#include <cmath>
#include <algorithm>

// ---------------------------------------------------------------------------
// micros() control -- fullmode_t1_infer_all_resident() calls it twice
// (t0, then elapsed = micros() - t0) to time each slot's inference. A
// call-counted sequence lets tests assert a specific elapsed_us.
// ---------------------------------------------------------------------------
#define STUB_MICROS_PROVIDED
static std::vector<uint32_t> g_micros_sequence;
static size_t g_micros_call_idx = 0;
static inline uint32_t micros() {
    if (g_micros_call_idx < g_micros_sequence.size()) {
        return g_micros_sequence[g_micros_call_idx++];
    }
    return 0;
}

// ---------------------------------------------------------------------------
// cactus_fullmode.h test doubles (real impl is a separate .cpp not included
// here -- these ARE the "cactus_fm_*" symbols for this TU).
// ---------------------------------------------------------------------------
static bool g_cactus_init_ret = true;
extern "C" bool cactus_fm_init(void) { return g_cactus_init_ret; }

static int g_cactus_load_ret = 0;
extern "C" int cactus_fm_load(const char*) { return g_cactus_load_ret; }

extern "C" bool cactus_fm_unload(int) { return true; }

static size_t g_psram_used = 0;
extern "C" size_t cactus_fm_psram_used(void) { return g_psram_used; }

static bool g_infer_ret = true;
static std::vector<float> g_infer_out_values;
static int g_infer_call_count = 0;
static std::vector<float> g_infer_last_in;
static size_t g_infer_last_out_len = 0;
extern "C" bool cactus_fm_infer(int /*slot*/, const float* in, size_t in_len,
                                float* out, size_t out_len) {
    g_infer_call_count++;
    g_infer_last_in.assign(in, in + in_len);
    g_infer_last_out_len = out_len;
    if (!g_infer_ret) return false;
    size_t n = std::min(out_len, g_infer_out_values.size());
    for (size_t i = 0; i < n; i++) out[i] = g_infer_out_values[i];
    return true;
}

#include "../../src/fullmode_t1.cpp"

// ---------------------------------------------------------------------------

void setUp(void) {
    fullmode_t1_test_reset();
    g_cactus_init_ret = true;
    g_cactus_load_ret = 0;
    g_psram_used = 0;
    g_infer_ret = true;
    g_infer_out_values.clear();
    g_infer_call_count = 0;
    g_infer_last_in.clear();
    g_infer_last_out_len = 0;
    g_micros_sequence.clear();
    g_micros_call_idx = 0;
}

void tearDown(void) {}

// ── fullmode_t1_init() (native stub) ─────────────────────────────────────

void test_init_resets_slot_count(void) {
    fullmode_t1_test_seed_slot(0, 2, 2);
    TEST_ASSERT_EQUAL_INT(1, s_fm_count);
    fullmode_t1_init();
    TEST_ASSERT_EQUAL_INT(0, s_fm_count);
}

void test_psram_used_passes_through(void) {
    g_psram_used = 123456;
    TEST_ASSERT_EQUAL_UINT32(123456, fullmode_t1_psram_used());
}

// ── fullmode_t1_infer_all_resident() -- dispatch/shape logic ────────────

void test_no_seeded_slots_returns_zero(void) {
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0.1f, 0.2f};
    int n = fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(0, n);
}

void test_input_dim_mismatch_skips_slot(void) {
    fullmode_t1_test_seed_slot(0, /*input_dim=*/128, /*output_dim=*/4);
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0.1f, 0.2f};  // len=2, slot wants 128
    int n = fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(0, n);
    TEST_ASSERT_EQUAL_INT(0, g_infer_call_count);  // never even dispatched
}

void test_unknown_input_dim_zero_accepts_any_length(void) {
    // probe_slot_dims() falls back to input_dim=0/output_dim=0 when it
    // can't determine dims -- fullmode_t1_infer_all_resident()'s check is
    // `fi.input_dim > 0 && input_len != fi.input_dim`, so dim==0 means
    // "accept anything" by construction, not "reject everything".
    fullmode_t1_test_seed_slot(0, /*input_dim=*/0, /*output_dim=*/2);
    g_infer_out_values = {1.0f, 1.0f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[5] = {0, 0, 0, 0, 0};
    int n = fullmode_t1_infer_all_resident(input, 5, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(1, n);
}

void test_cactus_infer_failure_skips_result(void) {
    fullmode_t1_test_seed_slot(0, 2, 2);
    g_infer_ret = false;
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0.1f, 0.2f};
    int n = fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(0, n);
}

void test_max_results_cap_respected(void) {
    fullmode_t1_test_seed_slot(0, 2, 2);
    fullmode_t1_test_seed_slot(1, 2, 2);
    fullmode_t1_test_seed_slot(2, 2, 2);
    g_infer_out_values = {1.0f, 0.0f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0.1f, 0.2f};
    int n = fullmode_t1_infer_all_resident(input, 2, results, /*max_results=*/2);
    TEST_ASSERT_EQUAL_INT(2, n);
}

void test_inactive_slot_is_skipped(void) {
    fullmode_t1_test_seed_slot(0, 2, 2);
    s_fm[0].active = false;  // same-TU direct poke, no public "unseed" hook
    g_infer_out_values = {1.0f, 0.0f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0.1f, 0.2f};
    int n = fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(0, n);
}

// ── softmax / argmax correctness ─────────────────────────────────────────

void test_equal_logits_produce_uniform_softmax(void) {
    fullmode_t1_test_seed_slot(0, 2, 2);
    g_infer_out_values = {0.0f, 0.0f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0, 0};
    int n = fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_FLOAT_WITHIN(1e-5f, 0.5f, results[0].probabilities[0]);
    TEST_ASSERT_FLOAT_WITHIN(1e-5f, 0.5f, results[0].probabilities[1]);
    TEST_ASSERT_FLOAT_WITHIN(1e-5f, 0.5f, results[0].confidence);
    TEST_ASSERT_EQUAL_UINT8(0, results[0].class_label);  // tie -> first index
}

void test_argmax_picks_highest_logit(void) {
    fullmode_t1_test_seed_slot(0, 3, 3);
    g_infer_out_values = {1.0f, 5.0f, 2.0f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[3] = {0, 0, 0};
    int n = fullmode_t1_infer_all_resident(input, 3, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_EQUAL_UINT8(1, results[0].class_label);
    TEST_ASSERT_TRUE(results[0].confidence > 0.5f);  // dominant logit
}

void test_probabilities_sum_to_one(void) {
    fullmode_t1_test_seed_slot(0, 4, 4);
    g_infer_out_values = {2.1f, -0.3f, 0.7f, 1.4f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[4] = {0, 0, 0, 0};
    int n = fullmode_t1_infer_all_resident(input, 4, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(1, n);
    float sum = 0.0f;
    for (int c = 0; c < results[0].num_classes; c++) sum += results[0].probabilities[c];
    TEST_ASSERT_FLOAT_WITHIN(1e-4f, 1.0f, sum);
}

void test_softmax_is_shift_invariant(void) {
    // softmax(x) == softmax(x + C) for any constant C -- verifies the
    // max-subtraction numerical-stability step doesn't change the result.
    fullmode_t1_test_seed_slot(0, 3, 3);
    g_infer_out_values = {100.0f, 105.0f, 98.0f};  // large logits
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[3] = {0, 0, 0};
    int n = fullmode_t1_infer_all_resident(input, 3, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(1, n);
    TEST_ASSERT_EQUAL_UINT8(1, results[0].class_label);
    TEST_ASSERT_FALSE(std::isnan(results[0].confidence));
    TEST_ASSERT_FALSE(std::isinf(results[0].confidence));
}

// ── backend-sentinel slot_id encoding ────────────────────────────────────

void test_slot_id_encodes_backend_sentinel_and_index(void) {
    fullmode_t1_test_seed_slot(10, 2, 2);
    fullmode_t1_test_seed_slot(20, 2, 2);
    g_infer_out_values = {1.0f, 0.0f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0, 0};
    int n = fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_INT(2, n);
    TEST_ASSERT_EQUAL_UINT8(MODEL_BACKEND_CACTUS | 0, results[0].slot_id);
    TEST_ASSERT_EQUAL_UINT8(MODEL_BACKEND_CACTUS | 1, results[1].slot_id);
}

// ── timing / call bookkeeping ─────────────────────────────────────────────

void test_infer_count_and_elapsed_us_tracked(void) {
    fullmode_t1_test_seed_slot(0, 2, 2);
    g_infer_out_values = {1.0f, 0.0f};
    g_micros_sequence = {1000, 2500};  // t0=1000, t1=2500 -> elapsed=1500
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0, 0};
    fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_UINT32(1500, results[0].latency_us);
    TEST_ASSERT_EQUAL_UINT32(1, s_fm[0].infer_count);
    // 64-bit assertion macros are disabled in this Unity build config.
    TEST_ASSERT_TRUE(s_fm[0].total_us == 1500);
}

void test_infer_count_accumulates_across_calls(void) {
    fullmode_t1_test_seed_slot(0, 2, 2);
    g_infer_out_values = {1.0f, 0.0f};
    ModelInferenceResult results[FULLMODE_T1_MAX_SLOTS];
    float input[2] = {0, 0};
    fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    fullmode_t1_infer_all_resident(input, 2, results, FULLMODE_T1_MAX_SLOTS);
    TEST_ASSERT_EQUAL_UINT32(3, s_fm[0].infer_count);
}

// ── seed hook bounds ──────────────────────────────────────────────────────

void test_seed_slot_returns_minus_one_when_full(void) {
    for (int i = 0; i < FULLMODE_T1_MAX_SLOTS; i++) {
        TEST_ASSERT_TRUE(fullmode_t1_test_seed_slot(i, 2, 2) >= 0);
    }
    TEST_ASSERT_EQUAL_INT(-1, fullmode_t1_test_seed_slot(99, 2, 2));
}

// ── main ──────────────────────────────────────────────────────────────────

int main(int, char**) {
    UNITY_BEGIN();

    RUN_TEST(test_init_resets_slot_count);
    RUN_TEST(test_psram_used_passes_through);

    RUN_TEST(test_no_seeded_slots_returns_zero);
    RUN_TEST(test_input_dim_mismatch_skips_slot);
    RUN_TEST(test_unknown_input_dim_zero_accepts_any_length);
    RUN_TEST(test_cactus_infer_failure_skips_result);
    RUN_TEST(test_max_results_cap_respected);
    RUN_TEST(test_inactive_slot_is_skipped);

    RUN_TEST(test_equal_logits_produce_uniform_softmax);
    RUN_TEST(test_argmax_picks_highest_logit);
    RUN_TEST(test_probabilities_sum_to_one);
    RUN_TEST(test_softmax_is_shift_invariant);

    RUN_TEST(test_slot_id_encodes_backend_sentinel_and_index);

    RUN_TEST(test_infer_count_and_elapsed_us_tracked);
    RUN_TEST(test_infer_count_accumulates_across_calls);

    RUN_TEST(test_seed_slot_returns_minus_one_when_full);

    return UNITY_END();
}
