/* scheduler_test: M2.5 multi-slot isolation gate.
 *
 * Gate (ROADMAP M2.5): two interleaved greedy requests produce bit-identical
 * output to the same two requests run sequentially.
 *
 * Also covers the no-tail prefill branch (M2.1 latent bug): a mult-of-16
 * prompt prefills via prefill_step, and slot_n_tokens must equal exactly
 * the prompt length after prefill (no spurious ++).
 *
 * The scheduler streams decoded TEXT (M2.4 incremental_decoder), not raw
 * token ids, so the isolation gate compares decoded text, not ids.  This is
 * semantically what the user sees and avoids a lossy id->text->id round-trip.
 *
 *   ./scheduler_test model.gguf "prompt A" "prompt B" [-n N]
 */
#include "qxmx_gpu.h"
#include "qxmx_model.h"
#include "qxmx_scheduler.h"
#include "qxmx_tokenizer.h"

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

namespace qx { extern qx_model* g_model; }

/* Sequential reference: prefill + greedy decode on a fresh slot, streaming
 * text through the same incremental_decoder the scheduler uses.  Uses the
 * prefill_step() loop (what the scheduler uses) so the reference matches the
 * scheduler's prefill path exactly. */
static std::string gen_sequential(qx::engine& eng, const std::vector<int32_t>& ids,
                                  int n_gen, const qx::tokenizer& tok) {
    int slot = eng.open_slot();
    int i = 0, lcm = 0; bool done = false;
    while (!done) eng.prefill_step(slot, ids.data(), (int)ids.size(), i, lcm, done);
    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)tok.eos_id) break;
        eng.forward(slot, next);
        next = eng.sample(slot, qx::sample_params{});
        out += dec.push(next);
    }
    out += dec.finish();
    eng.close_slot(slot);
    return out;
}

int main(int argc, char** argv) {
    std::vector<const char*> pos;
    int n_gen = 16;
    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 A\" \"prompt B\" [-n N]\n", argv[0]);
        return 1;
    }
    qx_model m; char err[256];
    if (qx_model_open(&m, pos[0], 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; }

    qx::g_model = &m;
    qx::engine eng;
    eng.k_ctk = qx::engine::KCT::q8_0;
    eng.init(&m, 1 << 14);

    std::vector<int32_t> idsA, idsB;
    tok.encode(pos[1], &idsA);
    tok.encode(pos[2], &idsB);
    std::printf("prompt A: %zu tokens\n", idsA.size());
    std::printf("prompt B: %zu tokens\n", idsB.size());

    /* ---- Sequential reference: A then B on fresh slots ---- */
    std::string seqA = gen_sequential(eng, idsA, n_gen, tok);
    std::string seqB = gen_sequential(eng, idsB, n_gen, tok);

    /* ---- Interleaved via scheduler (FIFO, one step per slot per round) ---- */
    qx::stop_policy stoppA, stoppB;
    stoppA.add_eos(tok.eos_id); stoppA.max_tokens = n_gen;
    stoppB.add_eos(tok.eos_id); stoppB.max_tokens = n_gen;
    std::string textA, textB;
    {
        qx::scheduler sched(&eng, &tok);
        sched.submit({idsA.data(), (int)idsA.size(), qx::sample_params{}, stoppA,
                      [&](int, const std::string& t){ textA += t; }});
        sched.submit({idsB.data(), (int)idsB.size(), qx::sample_params{}, stoppB,
                      [&](int, const std::string& t){ textB += t; }});
        sched.run();
    }

    int fail = 0;
    if (textA != seqA) {
        std::printf("FAIL: slot A interleaved != sequential\n");
        std::printf("  inter (%zu): \"%s\"\n", textA.size(), textA.c_str());
        std::printf("  seq   (%zu): \"%s\"\n", seqA.size(), seqA.c_str());
        fail = 1;
    } else std::printf("PASS: slot A interleaved == sequential (%zu bytes)\n", textA.size());
    if (textB != seqB) {
        std::printf("FAIL: slot B interleaved != sequential\n");
        std::printf("  inter (%zu): \"%s\"\n", textB.size(), textB.c_str());
        std::printf("  seq   (%zu): \"%s\"\n", seqB.size(), seqB.c_str());
        fail = 1;
    } else std::printf("PASS: slot B interleaved == sequential (%zu bytes)\n", textB.size());

    /* ---- No-tail prefill branch (M2.1 latent bug coverage) ---- */
    /* Pad prompt A to a multiple of 16 so the no-tail branch fires; after
     * prefill_step completes, slot_n_tokens must equal exactly the prompt
     * length.  The old code did a spurious ++ here. */
    {
        std::vector<int32_t> ids = idsA;
        while (ids.size() % 16 != 0) ids.push_back(tok.eos_id);
        int slot = eng.open_slot();
        int i = 0, last_chunk_m = 0; bool done = false;
        while (!done) eng.prefill_step(slot, ids.data(), (int)ids.size(), i, last_chunk_m, done);
        int64_t after = eng.slot_n_tokens(slot);
        std::printf("no-tail prefill: %zu-token (mult-16) prompt, slot_n_tokens after = %lld\n",
                    ids.size(), (long long)after);
        if (after != (int64_t)ids.size()) {
            std::printf("FAIL: slot_n_tokens == %lld, expected %zu (no-tail ++ bug)\n",
                        (long long)after, ids.size());
            fail = 1;
        } else std::printf("PASS: no-tail prefill slot_n_tokens == prompt_len\n");
        eng.close_slot(slot);
    }

    /* ---- Solo scheduler with id capture (diagnostic) ----
     * Rules out a scheduler bug vs an interleaving bug: if a solo slot through
     * the scheduler matches gen_sequential ids, the failures are interleaving. */
    {
        qx::stop_policy sp; sp.add_eos(tok.eos_id); sp.max_tokens = n_gen;
        auto gen_seq_ids = [&](const std::vector<int32_t>& ids) {
            int slot = eng.open_slot();
            eng.prefill(slot, ids.data(), (int)ids.size());
            std::vector<int32_t> g;
            int next = eng.sample(slot, qx::sample_params{}); g.push_back(next);
            for (int step = 1; step < n_gen; step++) {
                if (next == (int)tok.eos_id) break;
                eng.forward(slot, next); next = eng.sample(slot, qx::sample_params{}); g.push_back(next);
            }
            eng.close_slot(slot);
            return g;
        };
        auto seqIds = gen_seq_ids(idsA);
        std::vector<int32_t> soloIds;
        qx::scheduler sched(&eng, &tok);
        sched.submit({idsA.data(), (int)idsA.size(), qx::sample_params{}, sp,
                      /*on_token=*/{},
                      /*on_token_id=*/[&](int, int32_t t){ soloIds.push_back(t); }});
        sched.run();
        /* Determinism: two sequential runs on fresh slots must agree. */
        auto seqIds2 = gen_seq_ids(idsA);
        if (seqIds != seqIds2) {
            std::printf("FAIL: sequential decode not deterministic across runs\n");
            fail = 1;
        } else std::printf("PASS: sequential decode deterministic\n");
        if (seqIds != soloIds) {
            std::printf("FAIL: solo scheduler ids != sequential ids\n");
            std::printf("  seq : "); for (int t : seqIds) std::printf("%d ", t); std::printf("\n");
            std::printf("  solo: "); for (int t : soloIds) std::printf("%d ", t); std::printf("\n");
            fail = 1;
        } else std::printf("PASS: solo scheduler ids == sequential ids\n");
    }

    /* ---- Single-slot prefill() vs prefill_step() parity ----
     * Rules out a prefill_step bug independent of interleaving. */
    {
        int sa = eng.open_slot();
        eng.prefill(sa, idsA.data(), (int)idsA.size());
        std::vector<int32_t> gWrap;
        int next = eng.sample(sa, qx::sample_params{}); gWrap.push_back(next);
        for (int step = 1; step < n_gen; step++) {
            if (next == (int)tok.eos_id) break;
            eng.forward(sa, next); next = eng.sample(sa, qx::sample_params{}); gWrap.push_back(next);
        }
        eng.close_slot(sa);

        int sb = eng.open_slot();
        int i = 0, lcm = 0; bool done = false;
        while (!done) eng.prefill_step(sb, idsA.data(), (int)idsA.size(), i, lcm, done);
        std::vector<int32_t> gStep;
        next = eng.sample(sb, qx::sample_params{}); gStep.push_back(next);
        for (int step = 1; step < n_gen; step++) {
            if (next == (int)tok.eos_id) break;
            eng.forward(sb, next); next = eng.sample(sb, qx::sample_params{}); gStep.push_back(next);
        }
        eng.close_slot(sb);
        if (gWrap != gStep) {
            std::printf("FAIL: prefill() != prefill_step() loop (single slot)\n");
            std::printf("  wrap: "); for (int t : gWrap) std::printf("%d ", t); std::printf("\n");
            std::printf("  step: "); for (int t : gStep) std::printf("%d ", t); std::printf("\n");
            fail = 1;
        } else std::printf("PASS: prefill() == prefill_step() loop (single slot)\n");
    }

    qx_model_close(&m);
    return fail;
}