/* prefix_cache_test: M2.6.1 gate for the hash-chained KV block cache +
 * SSM checkpoint store.  Pure host, no SYCL.
 *
 * Validates:
 *   - block hashing: chained (parent mixed), distinct ids -> distinct hash.
 *   - KV put/get round-trips bytes byte-exact; ids verify-on-hit.
 *   - Hash collision (forced) falls back to miss via id verify (corruption
 *     guard).
 *   - KV LRU eviction picks the oldest ref==0 block; ref protects.
 *   - SSM put/get round-trips; evict-anywhere.
 *
 * The device<->host DMA path (engine integration) is M2.6.3; this test
 * exercises the cache module with synthetic byte blobs.
 */
#include "qxmx_prefix_cache.h"

#include <cassert>
#include <cstdio>
#include <cstring>
#include <vector>

using namespace qx;

static int failures = 0;
#define CHECK(cond) do { if (!(cond)) { \
    std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); \
    failures++; } } while (0)

int main() {
    /* Geometry matching the engine: 16 attn, 48 ssm, q8_0 K, TurboQuant V.
     * (KROW_BYTES_Q8=1088, VROW_U32*4=640.) */
    const int n_attn = 16, n_ssm = 48;
    const size_t k_row = 1088, v_row = 640;
    const size_t ssm_state = 48 * 128 * 128 * sizeof(float);
    const size_t ssm_conv  = 3 * 10240 * sizeof(float);
    prefix_cache pc;
    pc.init(n_attn, n_ssm, k_row, v_row, ssm_state, ssm_conv);

    /* Per-block byte sizes the cache expects. */
    const size_t k_blk = (size_t)n_attn * PC_BLOCK * k_row;
    const size_t v_blk = (size_t)n_attn * PC_BLOCK * v_row;

    /* ---- block hashing ---- */
    int32_t ids0[PC_BLOCK], ids1[PC_BLOCK];
    for (int i = 0; i < PC_BLOCK; i++) { ids0[i] = i; ids1[i] = i + 1; }
    block_hash h0 = hash_block({0}, ids0, PC_BLOCK);
    block_hash h1 = hash_block({0}, ids1, PC_BLOCK);
    block_hash h0b = hash_block({0}, ids0, PC_BLOCK);
    CHECK(h0 != h1);                /* different ids -> different hash */
    CHECK(h0 == h0b);               /* deterministic */

    /* Chaining: same block at different positions (different parent) differs. */
    block_hash h0_at1 = hash_block(h0, ids0, PC_BLOCK);
    CHECK(h0_at1 != hash_block({0xdeadbeef}, ids0, PC_BLOCK));

    /* ---- KV put/get round-trip ---- */
    std::vector<uint8_t> k_data(k_blk, 0), v_data(v_blk, 0);
    for (size_t i = 0; i < k_blk; i++) k_data[i] = (uint8_t)(i * 7 + 1);
    for (size_t i = 0; i < v_blk; i++) v_data[i] = (uint8_t)(i * 3 + 5);
    pc.kv_put(h0, ids0, PC_BLOCK, k_data.data(), v_data.data());
    CHECK(pc.kv_size() == 1);

    const kv_block* got = pc.kv_get(h0, ids0, PC_BLOCK);
    CHECK(got != nullptr);
    CHECK(std::memcmp(got->k.data(), k_data.data(), k_blk) == 0);
    CHECK(std::memcmp(got->v.data(), v_data.data(), v_blk) == 0);

    /* ---- miss on absent hash ---- */
    CHECK(pc.kv_get(h1, ids1, PC_BLOCK) == nullptr);

    /* ---- collision guard: same hash, different ids -> miss ----
     * Simulate by inserting ids0 under h0, then asking for h0 with ids1. */
    CHECK(pc.kv_get(h0, ids1, PC_BLOCK) == nullptr);
    CHECK(pc.kv_has(h0, ids1, PC_BLOCK) == false);
    CHECK(pc.kv_has(h0, ids0, PC_BLOCK) == true);  /* still present for right ids */

    /* ---- LRU eviction ---- */
    /* Insert 3 blocks; evict the oldest (h0).  h0's lru_seq was bumped by the
     * kv_has above, so insert h1, h2, then touch order is h0 < h1 < h2... */
    std::vector<uint8_t> k1(k_blk, 0xAA), v1(v_blk, 0xBB);
    std::vector<uint8_t> k2(k_blk, 0xCC), v2(v_blk, 0xDD);
    block_hash hh1 = hash_block(h0, ids1, PC_BLOCK);
    block_hash hh2 = hash_block(hh1, ids1, PC_BLOCK);
    pc.kv_put(hh1, ids1, PC_BLOCK, k1.data(), v1.data());
    pc.kv_put(hh2, ids1, PC_BLOCK, k2.data(), v2.data());
    CHECK(pc.kv_size() == 3);
    /* Now touch h0 so it's most-recent; hh1 is oldest. */
    pc.kv_get(h0, ids0, PC_BLOCK);
    CHECK(pc.kv_evict_one());
    CHECK(pc.kv_size() == 2);
    CHECK(pc.kv_get(hh1, ids1, PC_BLOCK) == nullptr);  /* hh1 was evicted */
    CHECK(pc.kv_get(h0, ids0, PC_BLOCK)  != nullptr);  /* h0 kept */
    CHECK(pc.kv_get(hh2, ids1, PC_BLOCK) != nullptr);  /* hh2 kept */

    /* ---- refcount protects from eviction ---- */
    pc.kv_ref(h0);
    CHECK(!pc.kv_evict_one() || pc.kv_get(h0, ids0, PC_BLOCK) != nullptr);
    /* Either no evictable block (hh2 is the only ref==0), or we evicted hh2
     * and h0 survived because ref>0.  Verify h0 survived. */
    CHECK(pc.kv_get(h0, ids0, PC_BLOCK) != nullptr);
    pc.kv_deref(h0);

    /* ---- SSM checkpoint round-trip ---- */
    std::vector<uint8_t> st((size_t)n_ssm * ssm_state, 0xEE);
    std::vector<uint8_t> cv((size_t)n_ssm * ssm_conv,  0x77);
    pc.ssm_put(h0, 1, st.data(), cv.data());
    CHECK(pc.ssm_size() == 1);
    const ssm_checkpoint* cp = pc.ssm_get(h0);
    CHECK(cp != nullptr);
    CHECK(cp->n_blocks == 1);
    CHECK(std::memcmp(cp->state.data(), st.data(), st.size()) == 0);
    CHECK(std::memcmp(cp->conv.data(),  cv.data(),  cv.size())  == 0);

    /* SSM evict-anywhere. */
    CHECK(pc.ssm_evict_one());
    CHECK(pc.ssm_size() == 0);
    CHECK(pc.ssm_get(h0) == nullptr);

    /* ---- budget enforcement ----
     * Set a tiny budget (2 checkpoints) and insert 5; the cache must
     * evict-from-anywhere down to <= budget.  Each checkpoint is
     * n_ssm*(ssm_state+ssm_conv) bytes. */
    pc.set_ssm_budget(2 * ((size_t)n_ssm * (ssm_state + ssm_conv)));
    for (int b = 0; b < 5; b++) {
        block_hash bh = hash_block(b == 0 ? block_hash{0} :
                                   hash_block({(uint64_t)(b*7)}, ids0, PC_BLOCK),
                                   ids0, PC_BLOCK);
        std::vector<uint8_t> s((size_t)n_ssm * ssm_state, (uint8_t)b);
        std::vector<uint8_t> cv((size_t)n_ssm * ssm_conv, (uint8_t)(b+1));
        pc.ssm_put(bh, b + 1, s.data(), cv.data());
    }
    CHECK(pc.ssm_bytes() <= pc.ssm_budget() + (size_t)n_ssm * (ssm_state + ssm_conv));
    CHECK(pc.ssm_size() <= 5);
    /* The most-recently-put checkpoint must survive (freshest lru_seq). */
    CHECK(pc.ssm_size() >= 1);

    if (failures == 0) {
        std::printf("prefix_cache_test: all checks passed (kv=%zu, ssm=%zu)\n",
                    pc.kv_size(), pc.ssm_size());
        return 0;
    }
    std::fprintf(stderr, "prefix_cache_test: %d FAILURE(S)\n", failures);
    return 1;
}