/* prefix_cache_engine_test: M2.6.3 gate -- in-engine cache correctness.
 *
 * Gate (ROADMAP M2.6 / MVP done-criteria #5): a cached prefill must produce
 * bit-identical greedy output to a fresh full prefill.  This is the
 * corruption guard for the device<->host DMA path: if the restored KV/SSM
 * bytes don't match what a fresh prefill would have written, the generated
 * tokens diverge.
 *
 * Test:
 *   1. Prefill prompt P (>= 1 block = 512 tok) on slot A, snapshotting to a
 *      shared cache.  Decode a few tokens to confirm the engine works post-
 *      snapshot.
 *   2. Reset slot B, prefill_cached(P + delta) -> should restore P's prefix
 *      (restored_to > 0), prefill only delta.  Greedy-decode N tokens.
 *   3. Fresh slot C: prefill(P + delta) in full, greedy-decode N tokens.
 *   4. Slots B and C must produce IDENTICAL decoded text.
 *
 *   ./prefix_cache_engine_test model.gguf prompt.txt delta.txt [-n N]
 *   (prompt.txt should be >= 512 tokens; delta.txt any length.)
 */
#include "qxmx_gpu.h"
#include "qxmx_model.h"
#include "qxmx_prefix_cache.h"
#include "qxmx_tokenizer.h"

#include <cstdio>
#include <fstream>
#include <cstring>
#include <sstream>
#include <string>
#include <vector>

namespace qx { extern qx_model* g_model; }

static std::string read_file(const char* path) {
    std::ifstream f(path);
    if (!f) { std::fprintf(stderr, "cannot open %s\n", path); return ""; }
    std::stringstream ss; ss << f.rdbuf(); return ss.str();
}

/* Greedy decode N tokens after the slot is pre-filled (logits already set for
 * the last prompt token).  Returns decoded text. */
static std::string gen_decode(qx::engine& eng, int slot, int n_gen,
                              const qx::tokenizer& tok, uint32_t eos_id) {
    qx::incremental_decoder dec(tok);
    std::string out;
    int next = eng.sample(slot, qx::sample_params{});
    out += dec.push(next);
    for (int step = 1; step < n_gen; step++) {
        if (next == (int)eos_id) break;
        eng.forward(slot, next);
        next = eng.sample(slot, qx::sample_params{});
        out += dec.push(next);
    }
    out += dec.finish();
    return out;
}

int main(int argc, char** argv) {
    std::vector<const char*> pos;
    int n_gen = 32;
    for (int i = 1; i < argc; i++) {
        std::string a = argv[i];
        if (a == "-n" && i + 1 < argc) n_gen = std::atoi(argv[++i]);
        else pos.push_back(argv[i]);
    }
    if (pos.size() < 3) {
        std::fprintf(stderr, "usage: %s model.gguf prompt.txt delta.txt [-n N]\n", argv[0]);
        return 1;
    }
    const char* model_path = pos[0];
    std::string prompt_str = read_file(pos[1]);
    std::string delta_str  = read_file(pos[2]);
    if (prompt_str.empty() || delta_str.empty()) return 1;

    qx_model m; char err[256];
    if (qx_model_open(&m, model_path, err, sizeof err)) {
        std::fprintf(stderr, "load: %s\n", err); return 1;
    }
    qx::tokenizer tok;
    if (!tok.load(&m.gguf, err, sizeof err)) { std::fprintf(stderr, "tok: %s\n", err); return 1; }

    std::vector<int32_t> P, D, PplusD;
    tok.encode(prompt_str, &P);
    tok.encode(delta_str,  &D);
    PplusD = P;
    PplusD.insert(PplusD.end(), D.begin(), D.end());

    if ((int)P.size() < qx::PC_BLOCK) {
        std::fprintf(stderr, "prompt too short: %zu tokens (need >= %d for one block)\n",
                     P.size(), qx::PC_BLOCK);
        return 1;
    }
    std::printf("prompt: %zu tokens (%d full blocks), delta: %zu tokens\n",
                P.size(), (int)(P.size() / qx::PC_BLOCK), D.size());

    int max_ctx = std::max((int)PplusD.size() + n_gen + 16, 1 << 14);
    qx::g_model = &m;
    qx::engine eng;
    eng.init(&m, max_ctx);

    /* Cache geometry: 16 attn, 48 ssm, q8_0 K (1088 B/row), V (640 B/row),
     * fp32 SSM state/conv.  Matches engine defaults. */
    qx::prefix_cache cache;
    size_t s_state = (size_t)QX_SSM_DT_RANK * QX_SSM_STATE * QX_SSM_STATE * sizeof(float);
    size_t s_conv  = (size_t)(QX_SSM_CONV_K - 1) * QX_SSM_CONVD * sizeof(float);
    cache.init(eng.n_attn, eng.n_ssm, 1088, 640, s_state, s_conv);

    /* ---- 1. Populate the cache by prefilling P on slot A with snapshotting ---- */
    int slotA = eng.open_slot();
    int restored = eng.prefill_cached(slotA, P.data(), (int)P.size(), &cache);
    std::printf("[A] prefill_cached(P) restored_to=%d (expect 0, cold cache)\n", restored);
    if (restored != 0) {
        std::fprintf(stderr, "FAIL: cold cache should restore 0, got %d\n", restored);
        return 1;
    }
    std::printf("[A] cache after P: kv=%zu ssm=%zu\n", cache.kv_size(), cache.ssm_size());
    if (cache.ssm_size() == 0) {
        std::fprintf(stderr, "FAIL: cache empty after prefilling P (snapshotting broken)\n");
        return 1;
    }
    /* Sanity: with stride=1 every block is a group of 1, so kv==ssm==blocks.
     * With stride>1, KV is cached for EVERY block (the group snapshot DMAs
     * all stride blocks' KV) but SSM checkpoints only at group ends -- so
     * kv==full_blocks and ssm==full_blocks/stride. */
    int snap_stride = []{ const char* e = std::getenv("QXMX_SNAP_STRIDE");
                          int v = e ? std::atoi(e) : 1; return v < 1 ? 1 : v; }();
    size_t full_blocks  = P.size() / qx::PC_BLOCK;
    size_t expect_kv   = full_blocks;
    size_t expect_ssm  = full_blocks / snap_stride;
    size_t expect_blocks = expect_ssm * snap_stride;  /* whole groups, in blocks */
    if (cache.kv_size() != expect_kv || cache.ssm_size() != expect_ssm) {
        std::fprintf(stderr, "FAIL: cache counts kv=%zu ssm=%zu, expected kv=%zu ssm=%zu (stride=%d)\n",
                     cache.kv_size(), cache.ssm_size(), expect_kv, expect_ssm, snap_stride);
        return 1;
    }
    eng.close_slot(slotA);

    /* ---- 2. Slot B: prefill_cached(P+D) -- should restore P, prefill only D ---- */
    int slotB = eng.open_slot();
    restored = eng.prefill_cached(slotB, PplusD.data(), (int)PplusD.size(), &cache);
    size_t restored_blocks = (size_t)restored / qx::PC_BLOCK;
    std::printf("[B] prefill_cached(P+D) restored_to=%d (%zu blocks of %zu)\n",
                restored, restored_blocks, expect_blocks);
    /* The cache may hold FEWER than expect_blocks if the budget evicted some;
     * the 16 GB default holds all of P's blocks (P ~10 blocks).  Require at
     * least the majority restored. */
    if (restored_blocks < expect_blocks) {
        std::fprintf(stderr, "FAIL: only %zu/%zu blocks restored\n", restored_blocks, expect_blocks);
        return 1;
    }
    if (restored_blocks == 0) {
        std::fprintf(stderr, "FAIL: cache hit but restored 0 blocks\n");
        return 1;
    }
    std::string outB = gen_decode(eng, slotB, n_gen, tok, tok.eos_id);
    eng.close_slot(slotB);

    /* ---- 3. Slot C: fresh full prefill of P+D, no cache ---- */
    int slotC = eng.open_slot();
    eng.prefill(slotC, PplusD.data(), (int)PplusD.size());
    std::string outC = gen_decode(eng, slotC, n_gen, tok, tok.eos_id);
    eng.close_slot(slotC);

    /* ---- 4. Compare ---- */
    std::printf("[B] cached:  \"%s\"\n", outB.c_str());
    std::printf("[C] fresh:   \"%s\"\n", outC.c_str());
    if (outB != outC) {
        std::fprintf(stderr, "FAIL: cached vs fresh prefill diverge (%zu vs %zu bytes)\n",
                     outB.size(), outC.size());
        return 1;
    }
    std::printf("PASS: cached prefill produces bit-identical output to fresh prefill.\n");
    std::printf("  restored %zu/%zu prefix blocks (%zu tokens), prefilled only the delta.\n",
                restored_blocks, expect_blocks, (size_t)restored);
    qx_model_close(&m);
    return 0;
}