/*  test_fullmode_ondevice.cpp — on-device smoke test for cactus_fullmode
 *
 *  Runs on Heltec V4 only (uses SPIFFS + PSRAM + Arduino runtime).
 *  Build with:  pio test -e heltec_wifi_lora_v4 -f test_fullmode
 *
 *  NOTE: native environment cannot run this test because cactus_fullmode.cpp
 *  depends on SPIFFS, heap_caps_malloc (PSRAM), and Arduino Serial.
 *  A host-side unit test would require stubbing all three — use the real
 *  hardware target instead. See INTEGRATION_NOTES.md for rationale.
 *
 *  The test constructs a minimal valid .cactus blob in PSRAM, writes it to
 *  SPIFFS as /models/smoke_test.cactus, loads it via cactus_fm_load(), runs
 *  inference on a known input, and verifies the output is a valid probability
 *  distribution (all positive, sums to ~1.0).
 */

#include <Arduino.h>
#include <LittleFS.h>
#include <unity.h>
#include <string.h>
#include <stdint.h>
#include <math.h>

#include "cactus_fullmode.h"

// ── Minimal .cactus file builder ──────────────────────────────────────────
// We build a tiny 2->2 linear+softmax network by hand so the test has no
// dependency on Python or the DGX.

static const char* SMOKE_PATH = "/models/smoke_test.cactus";

// JSON header for a 2-in, 2-out FP32 linear->softmax network
static const char SMOKE_HEADER[] =
    "{\"arch\":["
      "{\"type\":\"linear\",\"in\":2,\"out\":2,\"dtype\":\"fp32\"},"
      "{\"type\":\"softmax\"}"
    "],"
    "\"input_shape\":[2],"
    "\"output_shape\":[2],"
    "\"meta\":{\"name\":\"smoke\",\"version\":\"0.0.1\",\"dtype\":\"fp32\"}"
    "}";

// Weight matrix W (2x2 fp32, row-major): identity-ish
// [[1, 0], [0, 1]]   b = [0.1, -0.1]
static const float SMOKE_W[4] = { 1.0f, 0.0f, 0.0f, 1.0f };
static const float SMOKE_B[2] = { 0.1f, -0.1f };

static bool write_smoke_cactus() {
    if (!LittleFS.begin(true)) return false;

    File f = LittleFS.open(SMOKE_PATH, "w");
    if (!f) return false;

    // Magic
    f.write((const uint8_t*)"CACTFM01", 8);

    // Header length (little-endian u32)
    uint32_t hlen = (uint32_t)(sizeof(SMOKE_HEADER) - 1); // exclude NUL
    uint8_t hlen_buf[4];
    hlen_buf[0] = hlen & 0xFF;
    hlen_buf[1] = (hlen >> 8) & 0xFF;
    hlen_buf[2] = (hlen >> 16) & 0xFF;
    hlen_buf[3] = (hlen >> 24) & 0xFF;
    f.write(hlen_buf, 4);

    // Header JSON
    f.write((const uint8_t*)SMOKE_HEADER, hlen);

    // Weight blob: W (fp32), then b (fp32)
    f.write((const uint8_t*)SMOKE_W, sizeof(SMOKE_W));
    f.write((const uint8_t*)SMOKE_B, sizeof(SMOKE_B));

    f.close();
    return true;
}

// ── Tests ─────────────────────────────────────────────────────────────────

static int g_slot = -1;

void setUp(void) {}
void tearDown(void) {}

void test_write_smoke_file(void) {
    TEST_ASSERT_TRUE_MESSAGE(write_smoke_cactus(), "write smoke .cactus to SPIFFS");
}

void test_fm_init(void) {
    TEST_ASSERT_TRUE_MESSAGE(cactus_fm_init(), "cactus_fm_init should return true");
}

void test_fm_load(void) {
    g_slot = cactus_fm_load(SMOKE_PATH);
    TEST_ASSERT_TRUE_MESSAGE(g_slot >= 0, "cactus_fm_load should return valid slot");
}

void test_fm_infer_valid_distribution(void) {
    TEST_ASSERT_TRUE_MESSAGE(g_slot >= 0, "slot must be loaded first");

    // Input: [1.0, 0.5]  -> logits = [1.0*1+0+0.1, 0.5*1-0.1] = [1.1, 0.4]
    // Softmax: e^1.1 / (e^1.1 + e^0.4) should favour class 0
    float in_data[2]  = { 1.0f, 0.5f };
    float out_data[2] = { 0.0f, 0.0f };

    bool ok = cactus_fm_infer(g_slot, in_data, 2, out_data, 2);
    TEST_ASSERT_TRUE_MESSAGE(ok, "infer should succeed");

    // Both outputs should be positive
    TEST_ASSERT_TRUE(out_data[0] > 0.0f);
    TEST_ASSERT_TRUE(out_data[1] > 0.0f);

    // They should sum to approximately 1.0 (softmax property)
    float sum = out_data[0] + out_data[1];
    TEST_ASSERT_FLOAT_WITHIN_MESSAGE(0.01f, 1.0f, sum, "softmax outputs must sum to 1.0");

    // Class 0 should have higher probability (logit 1.1 > 0.4)
    TEST_ASSERT_TRUE_MESSAGE(out_data[0] > out_data[1],
                             "class 0 should have higher prob for input [1.0, 0.5]");
}

void test_fm_wrong_input_size_fails(void) {
    float in_data[3] = { 0.f, 0.f, 0.f };
    float out_data[2];
    bool ok = cactus_fm_infer(g_slot, in_data, 3, out_data, 2);
    TEST_ASSERT_FALSE_MESSAGE(ok, "infer with wrong in_len must return false");
}

void test_fm_psram_used_nonzero(void) {
    size_t used = cactus_fm_psram_used();
    TEST_ASSERT_TRUE_MESSAGE(used > 0, "psram_used should be > 0 after load");
}

void test_fm_unload(void) {
    if (g_slot >= 0) {
        bool ok = cactus_fm_unload(g_slot);
        TEST_ASSERT_TRUE_MESSAGE(ok, "unload should succeed");
        g_slot = -1;
    }
}

void test_fm_psram_used_zero_after_unload(void) {
    size_t used = cactus_fm_psram_used();
    TEST_ASSERT_EQUAL_MESSAGE(0, (int)used, "psram_used should drop to 0 after all unloads");
}

// ── Entry point ───────────────────────────────────────────────────────────

void setup() {
    delay(2000);
    UNITY_BEGIN();

    RUN_TEST(test_write_smoke_file);
    RUN_TEST(test_fm_init);
    RUN_TEST(test_fm_load);
    RUN_TEST(test_fm_infer_valid_distribution);
    RUN_TEST(test_fm_wrong_input_size_fails);
    RUN_TEST(test_fm_psram_used_nonzero);
    RUN_TEST(test_fm_unload);
    RUN_TEST(test_fm_psram_used_zero_after_unload);

    UNITY_END();
}

void loop() {
    // Tests are one-shot; nothing to do here.
}
