// ============================================================================
// test/test_ota_adapter/test_ota_adapter.cpp
//
// Unity tests for src/ota_adapter.{h,cpp} -- double-buffered PSRAM model
// hot-swap over chunked LoRa transfer. Previously zero test coverage and
// unconditionally `#include <Arduino.h>` / `#include <esp_heap_caps.h>`'d,
// so it could not be compiled natively. Applied the same #ifdef ARDUINO /
// native-stub pattern used elsewhere in this sweep (see health_report.cpp,
// uav_relay.cpp): heap_caps_malloc/free stubbed to plain malloc/free,
// classifier_swap_model() (real impl uses TFLite-Micro, ESP32-only) is
// stubbed here and its calls captured for assertions.
//
// On the live mesh, ota_adapter's chunk handler only ever receives payload
// that has ALREADY passed ota_signed_verify_chunk()'s Ed25519 signature
// check in main.cpp -- an unsigned/corrupt chunk never reaches this code.
// These tests cover the double-buffer/reassembly logic that still runs on
// an already-authenticated chunk stream.
//
// Writing these tests surfaced a real capacity gap (not fixed here --
// flagged, since fixing it means widening chunk_index/total_chunks from
// uint8_t on the wire, a protocol change that needs the T2 gateway's
// chunking side updated in lockstep, which this sweep has not verified):
// chunk_index and total_chunks are both uint8_t (wire bytes), and
// OTA_MAX_CHUNK_PAYLOAD is 220 bytes, so the maximum model size this
// protocol can EVER deliver is 255 * 220 = 56,100 bytes (~54.8 KB) --
// even though OTA_MODEL_SLOT_SIZE allocates a full 512 KB PSRAM buffer
// per slot. A real quantized TFLite model easily exceeds 55 KB, meaning
// today's OTA path silently caps out well below the buffer it was given.
// See test_max_deliverable_size_is_far_below_slot_size below.
//
// Run with:
//   pio test -e native -f test_ota_adapter
// ============================================================================

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

// ota_adapter.cpp #include "rf_classifier.h" for classifier_swap_model()'s
// declaration; the real implementation (rf_classifier.cpp) uses TFLite-Micro
// and is ESP32-only, so we provide a native stub with the exact declared
// signature and capture every call for assertions.
static int    g_swap_call_count = 0;
static size_t g_swap_last_size  = 0;
static bool   g_swap_should_succeed = true;
bool classifier_swap_model(const uint8_t* model_data, size_t model_size) {
    (void)model_data;
    g_swap_call_count++;
    g_swap_last_size = model_size;
    return g_swap_should_succeed;
}

#include "../../src/ota_adapter.cpp"

// ── Test helpers ─────────────────────────────────────────────────────────

static uint16_t test_crc16_ccitt(const uint8_t* data, size_t len) {
    uint16_t crc = 0xFFFF;
    for (size_t i = 0; i < len; i++) {
        crc ^= (uint16_t)data[i] << 8;
        for (int j = 0; j < 8; j++) {
            if (crc & 0x8000) crc = (crc << 1) ^ 0x1021;
            else crc <<= 1;
        }
    }
    return crc;
}

// Builds one wire-format OTA chunk packet:
//   [0..32]  detection header (unused by ota_handle_chunk, just skipped)
//   [33]     chunk_index
//   [34]     total_chunks
//   [35..36] chunk_crc16 (BE)
//   [37..38] model_version (BE)
//   [39]     reserved
//   [40..]   payload
static std::vector<uint8_t> build_chunk(
    uint8_t chunk_index, uint8_t total_chunks, uint16_t model_version,
    const std::vector<uint8_t>& payload, bool corrupt_crc = false
) {
    std::vector<uint8_t> pkt(DETECTION_PACKET_SIZE + OTA_CHUNK_HEADER_SIZE);
    // Detection header content is irrelevant to ota_handle_chunk.
    std::fill(pkt.begin(), pkt.begin() + DETECTION_PACKET_SIZE, 0);

    uint16_t crc = test_crc16_ccitt(payload.data(), payload.size());
    if (corrupt_crc) crc ^= 0xFFFF;

    size_t off = DETECTION_PACKET_SIZE;
    pkt[off + 0] = chunk_index;
    pkt[off + 1] = total_chunks;
    pkt[off + 2] = (uint8_t)((crc >> 8) & 0xFF);
    pkt[off + 3] = (uint8_t)(crc & 0xFF);
    pkt[off + 4] = (uint8_t)((model_version >> 8) & 0xFF);
    pkt[off + 5] = (uint8_t)(model_version & 0xFF);
    pkt[off + 6] = 0;  // reserved

    pkt.insert(pkt.end(), payload.begin(), payload.end());
    return pkt;
}

static std::vector<uint8_t> make_payload(size_t n, uint8_t fill = 0xAB) {
    return std::vector<uint8_t>(n, fill);
}

void setUp(void) {
    g_swap_call_count = 0;
    g_swap_last_size  = 0;
    g_swap_should_succeed = true;
    ota_deinit();
    TEST_ASSERT_TRUE(ota_init());
}

void tearDown(void) {
    ota_deinit();
}

// ── Init / reset state ──────────────────────────────────────────────────

void test_init_starts_with_no_staged_or_active_model(void) {
    TEST_ASSERT_EQUAL_UINT16(0, ota_staged_version());
    TEST_ASSERT_EQUAL_UINT16(0, ota_active_version());
    TEST_ASSERT_FALSE(ota_model_ready());
}

// ── Chunk validation ─────────────────────────────────────────────────────

void test_handle_chunk_too_short_rejected(void) {
    std::vector<uint8_t> tiny(DETECTION_PACKET_SIZE + OTA_CHUNK_HEADER_SIZE);  // no payload byte
    TEST_ASSERT_FALSE(ota_handle_chunk(tiny.data(), tiny.size()));
    TEST_ASSERT_FALSE(ota_model_ready());
}

void test_handle_chunk_bad_crc_rejected(void) {
    auto pkt = build_chunk(0, 1, 7, make_payload(10), /*corrupt_crc=*/true);
    TEST_ASSERT_FALSE(ota_handle_chunk(pkt.data(), pkt.size()));
    TEST_ASSERT_FALSE(ota_model_ready());
}

void test_handle_chunk_bad_crc_does_not_advance_staged_version(void) {
    auto pkt = build_chunk(0, 1, 7, make_payload(10), /*corrupt_crc=*/true);
    ota_handle_chunk(pkt.data(), pkt.size());
    // A rejected chunk must not silently start tracking a new transfer.
    TEST_ASSERT_EQUAL_UINT16(0, ota_staged_version());
}

// ── Single/multi-chunk transfer completion ───────────────────────────────

void test_single_chunk_transfer_completes(void) {
    auto pkt = build_chunk(0, 1, 5, make_payload(50));
    TEST_ASSERT_TRUE(ota_handle_chunk(pkt.data(), pkt.size()));
    TEST_ASSERT_TRUE(ota_model_ready());
    TEST_ASSERT_EQUAL_UINT16(5, ota_staged_version());
}

void test_multi_chunk_transfer_incomplete_until_all_received(void) {
    auto c0 = build_chunk(0, 3, 9, make_payload(OTA_MAX_CHUNK_PAYLOAD));
    auto c1 = build_chunk(1, 3, 9, make_payload(OTA_MAX_CHUNK_PAYLOAD));
    auto c2 = build_chunk(2, 3, 9, make_payload(30));

    TEST_ASSERT_TRUE(ota_handle_chunk(c0.data(), c0.size()));
    TEST_ASSERT_FALSE(ota_model_ready());
    TEST_ASSERT_TRUE(ota_handle_chunk(c1.data(), c1.size()));
    TEST_ASSERT_FALSE(ota_model_ready());
    TEST_ASSERT_TRUE(ota_handle_chunk(c2.data(), c2.size()));
    TEST_ASSERT_TRUE(ota_model_ready());
}

void test_multi_chunk_transfer_completes_out_of_order(void) {
    auto c0 = build_chunk(0, 3, 9, make_payload(OTA_MAX_CHUNK_PAYLOAD));
    auto c1 = build_chunk(1, 3, 9, make_payload(OTA_MAX_CHUNK_PAYLOAD));
    auto c2 = build_chunk(2, 3, 9, make_payload(30));

    // Last chunk arrives first -- dest_offset is a deterministic function
    // of chunk_index, so arrival order must not matter.
    TEST_ASSERT_TRUE(ota_handle_chunk(c2.data(), c2.size()));
    TEST_ASSERT_FALSE(ota_model_ready());
    TEST_ASSERT_TRUE(ota_handle_chunk(c0.data(), c0.size()));
    TEST_ASSERT_TRUE(ota_handle_chunk(c1.data(), c1.size()));
    TEST_ASSERT_TRUE(ota_model_ready());
}

void test_duplicate_chunk_does_not_double_count(void) {
    auto c0 = build_chunk(0, 2, 4, make_payload(20));
    auto c1 = build_chunk(1, 2, 4, make_payload(20));

    TEST_ASSERT_TRUE(ota_handle_chunk(c0.data(), c0.size()));
    TEST_ASSERT_TRUE(ota_handle_chunk(c0.data(), c0.size()));  // duplicate
    TEST_ASSERT_TRUE(ota_handle_chunk(c0.data(), c0.size()));  // duplicate
    TEST_ASSERT_FALSE(ota_model_ready());  // still missing chunk 1

    TEST_ASSERT_TRUE(ota_handle_chunk(c1.data(), c1.size()));
    TEST_ASSERT_TRUE(ota_model_ready());
}

void test_new_version_transfer_resets_partial_previous_transfer(void) {
    // Start a v1 transfer, receive only chunk 0 of 2.
    auto v1c0 = build_chunk(0, 2, 1, make_payload(20));
    TEST_ASSERT_TRUE(ota_handle_chunk(v1c0.data(), v1c0.size()));
    TEST_ASSERT_FALSE(ota_model_ready());

    // A v2 transfer starts before v1 finished -- must fully reset, not
    // merge state with the abandoned v1 partial transfer.
    auto v2c0 = build_chunk(0, 2, 2, make_payload(20));
    auto v2c1 = build_chunk(1, 2, 2, make_payload(20));
    TEST_ASSERT_TRUE(ota_handle_chunk(v2c0.data(), v2c0.size()));
    TEST_ASSERT_TRUE(ota_handle_chunk(v2c1.data(), v2c1.size()));
    TEST_ASSERT_TRUE(ota_model_ready());
    TEST_ASSERT_EQUAL_UINT16(2, ota_staged_version());
}

// ── ota_apply_model() ────────────────────────────────────────────────────

void test_apply_model_fails_when_not_ready(void) {
    TEST_ASSERT_FALSE(ota_apply_model());
    TEST_ASSERT_EQUAL_INT(0, g_swap_call_count);
}

void test_apply_model_calls_classifier_swap_with_received_size(void) {
    auto pkt = build_chunk(0, 1, 3, make_payload(123));
    ota_handle_chunk(pkt.data(), pkt.size());
    TEST_ASSERT_TRUE(ota_apply_model());
    TEST_ASSERT_EQUAL_INT(1, g_swap_call_count);
    TEST_ASSERT_EQUAL_UINT32(123, (uint32_t)g_swap_last_size);
}

void test_apply_model_success_updates_active_version_and_clears_ready(void) {
    auto pkt = build_chunk(0, 1, 8, make_payload(10));
    ota_handle_chunk(pkt.data(), pkt.size());
    TEST_ASSERT_TRUE(ota_apply_model());
    TEST_ASSERT_EQUAL_UINT16(8, ota_active_version());
    TEST_ASSERT_FALSE(ota_model_ready());  // consumed, not re-appliable
}

void test_apply_model_failure_keeps_previous_active_version(void) {
    // Successfully apply v1 first so there's a real "previous" version.
    auto pkt1 = build_chunk(0, 1, 1, make_payload(10));
    ota_handle_chunk(pkt1.data(), pkt1.size());
    TEST_ASSERT_TRUE(ota_apply_model());
    TEST_ASSERT_EQUAL_UINT16(1, ota_active_version());

    // v2 transfer completes, but the swap itself fails (e.g. corrupt
    // model that fails TFLite-Micro validation).
    g_swap_should_succeed = false;
    auto pkt2 = build_chunk(0, 1, 2, make_payload(10));
    ota_handle_chunk(pkt2.data(), pkt2.size());
    TEST_ASSERT_FALSE(ota_apply_model());

    // A failed swap must not silently promote v2 to "active".
    TEST_ASSERT_EQUAL_UINT16(1, ota_active_version());
}

void test_apply_model_failure_leaves_model_ready_for_retry(void) {
    // Intentional/documented behavior: a failed swap does NOT clear
    // s_model_complete, so the caller can retry ota_apply_model() without
    // needing the gateway to resend the whole transfer.
    g_swap_should_succeed = false;
    auto pkt = build_chunk(0, 1, 4, make_payload(10));
    ota_handle_chunk(pkt.data(), pkt.size());
    TEST_ASSERT_FALSE(ota_apply_model());
    TEST_ASSERT_TRUE(ota_model_ready());

    g_swap_should_succeed = true;
    TEST_ASSERT_TRUE(ota_apply_model());
    TEST_ASSERT_FALSE(ota_model_ready());
}

// ── Buffer-bounds safety ─────────────────────────────────────────────────

void test_chunk_overflowing_slot_is_rejected_not_crashed(void) {
    // Craft a chunk_index that would push dest_offset + payload past
    // OTA_MODEL_SLOT_SIZE. With real OTA_MODEL_SLOT_SIZE=512KB and
    // OTA_MAX_CHUNK_PAYLOAD=220, this branch is unreachable via a real
    // uint8_t chunk_index (max deliverable is ~55KB, see below) -- but
    // the bounds check must still hold and not crash if it somehow does.
    size_t big_index = OTA_MODEL_SLOT_SIZE / OTA_MAX_CHUNK_PAYLOAD;  // out of uint8_t range in practice
    if (big_index > 255) big_index = 255;  // clamp to what the wire can actually carry
    auto pkt = build_chunk((uint8_t)big_index, (uint8_t)(big_index + 1), 6,
                            make_payload(OTA_MAX_CHUNK_PAYLOAD));
    // Must return a bool either way without corrupting memory (ASan/UBSan
    // would catch a real overflow if this ever regresses).
    bool result = ota_handle_chunk(pkt.data(), pkt.size());
    (void)result;
    TEST_PASS();
}

void test_max_deliverable_size_is_far_below_slot_size(void) {
    // Documents a real capacity gap found while writing these tests:
    // chunk_index/total_chunks are uint8_t on the wire (max 255), and
    // OTA_MAX_CHUNK_PAYLOAD is 220 bytes/chunk, so the largest model this
    // protocol can ever deliver is 255*220 = 56,100 bytes -- far below
    // the 512KB PSRAM slot ota_init() allocates for it. Not fixed here
    // (widening chunk_index needs a coordinated wire-format change with
    // T2's chunking side, not verified in this sweep) -- flagged.
    const size_t max_deliverable = 255u * OTA_MAX_CHUNK_PAYLOAD;
    TEST_ASSERT_EQUAL_UINT32(56100u, (uint32_t)max_deliverable);
    TEST_ASSERT_TRUE(max_deliverable < OTA_MODEL_SLOT_SIZE);
}

// ── Reset / deinit ────────────────────────────────────────────────────────

void test_reset_discards_partial_transfer(void) {
    auto c0 = build_chunk(0, 2, 11, make_payload(20));
    ota_handle_chunk(c0.data(), c0.size());

    ota_reset();

    // After reset, the same chunk index is treated as fresh (not a dup),
    // proving the received-chunk bitmask was actually cleared.
    auto c1 = build_chunk(1, 2, 11, make_payload(20));
    TEST_ASSERT_TRUE(ota_handle_chunk(c0.data(), c0.size()));
    TEST_ASSERT_TRUE(ota_handle_chunk(c1.data(), c1.size()));
    TEST_ASSERT_TRUE(ota_model_ready());
}

void test_deinit_then_reinit_starts_clean(void) {
    auto pkt = build_chunk(0, 1, 42, make_payload(10));
    ota_handle_chunk(pkt.data(), pkt.size());
    TEST_ASSERT_TRUE(ota_model_ready());

    ota_deinit();
    TEST_ASSERT_TRUE(ota_init());
    TEST_ASSERT_FALSE(ota_model_ready());
    TEST_ASSERT_EQUAL_UINT16(0, ota_staged_version());
    TEST_ASSERT_EQUAL_UINT16(0, ota_active_version());
}

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

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

    RUN_TEST(test_init_starts_with_no_staged_or_active_model);

    RUN_TEST(test_handle_chunk_too_short_rejected);
    RUN_TEST(test_handle_chunk_bad_crc_rejected);
    RUN_TEST(test_handle_chunk_bad_crc_does_not_advance_staged_version);

    RUN_TEST(test_single_chunk_transfer_completes);
    RUN_TEST(test_multi_chunk_transfer_incomplete_until_all_received);
    RUN_TEST(test_multi_chunk_transfer_completes_out_of_order);
    RUN_TEST(test_duplicate_chunk_does_not_double_count);
    RUN_TEST(test_new_version_transfer_resets_partial_previous_transfer);

    RUN_TEST(test_apply_model_fails_when_not_ready);
    RUN_TEST(test_apply_model_calls_classifier_swap_with_received_size);
    RUN_TEST(test_apply_model_success_updates_active_version_and_clears_ready);
    RUN_TEST(test_apply_model_failure_keeps_previous_active_version);
    RUN_TEST(test_apply_model_failure_leaves_model_ready_for_retry);

    RUN_TEST(test_chunk_overflowing_slot_is_rejected_not_crashed);
    RUN_TEST(test_max_deliverable_size_is_far_below_slot_size);

    RUN_TEST(test_reset_discards_partial_transfer);
    RUN_TEST(test_deinit_then_reinit_starts_clean);

    return UNITY_END();
}
