/* qxmx_batch: multi-slot batched-decode benchmark (M4.1 spike).
 *
 * Loads the real model, opens N slots, prefills each with the same prompt
 * (so all reach a common depth), then drives N steps of batched decode via
 * engine::forward_batch (FFN run as one M=N gemm_gpu_v4 call).  Reports
 * per-step wall-clock and aggregate tok/s, A/B vs N separate single-slot
 * forward() calls (the serialized baseline).
 *
 *   ./qxmx_batch model.gguf -f prompt.txt -n 64 -N 4
 *
 * -N: number of slots (default 4).  Padded to a multiple of 16 internally.
 */
#include "qxmx_engine.h"
#include "qxmx_gpu.h"
#include "qxmx_model.h"
#include "qxmx_tokenizer.h"
#include "qxmx_device.h"

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

namespace qx { extern qx_model* g_model; }

int main(int argc, char** argv) {
    int  n_gen = 64;
    int  N_slots = 4;
    const char* prompt_file = nullptr;
    const char* model_path = nullptr;
    qx::engine::KCT k_ctk = qx::engine::KCT::q8_0;

    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 if (a == "-N" && i + 1 < argc) N_slots = std::atoi(argv[++i]);
        else if (a == "-f" && i + 1 < argc) prompt_file = argv[++i];
        else if (a == "-ctk" && i + 1 < argc) {
            std::string v = argv[++i];
            if      (v == "fp16") k_ctk = qx::engine::KCT::fp16;
            else if (v == "q8_0") k_ctk = qx::engine::KCT::q8_0;
            else if (v == "fp8")  k_ctk = qx::engine::KCT::fp8;
            else { std::fprintf(stderr, "-ctk: unknown '%s'\n", v.c_str()); return 1; }
        }
        else if (a[0] != '-') model_path = argv[i];
    }
    if (!model_path || !prompt_file) {
        std::fprintf(stderr, "usage: %s model.gguf -f prompt.txt -n N [-N slots] [-ctk fp16|q8_0|fp8]\n",
                     argv[0]);
        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::ifstream pf(prompt_file);
    std::stringstream ss; ss << pf.rdbuf();
    std::string prompt_str = ss.str();
    std::vector<int32_t> ids;
    tok.encode(prompt_str, &ids);
    std::printf("prompt: %zu tokens, N=%d slots, %d decode steps\n",
                ids.size(), N_slots, n_gen);

    int max_ctx = (int)ids.size() + n_gen + 16;
    if (max_ctx < (1 << 14)) max_ctx = 1 << 14;
    qx::g_model = &m;
    auto eng = std::make_unique<qx::engine>();
    eng->k_ctk = k_ctk;
    if (!eng->init(&m, max_ctx)) { std::fprintf(stderr, "engine init failed\n"); return 1; }
    std::printf("device: %s\n", qx::device_name(eng->q).c_str());

    /* Open N slots, prefill each with the same prompt. */
    std::vector<int> slots(N_slots);
    for (int i = 0; i < N_slots; i++) {
        slots[i] = eng->open_slot();
        eng->prefill(slots[i], ids.data(), (int)ids.size());
    }
    std::printf("prefilled %d slots to depth %zu\n", N_slots, ids.size());

    /* Sample the first generated token per slot (greedy). */
    qx::sample_params sp;   /* greedy */
    std::vector<int32_t> next_tok(N_slots);
    for (int i = 0; i < N_slots; i++) next_tok[i] = eng->sample(slots[i], sp);

    /* forward_batch writes per-slot logits to eng->logits_batch_buf[i*VOCAB]
     * (USM shared).  sample_slot copies row i into eng->logits before sampling,
     * since sample() reads eng->logits. */
    auto sample_slot = [&](int i) -> int32_t {
        std::memcpy(eng->logits, eng->logits_batch_buf + (size_t)i * QX_N_VOCAB,
                    (size_t)QX_N_VOCAB * sizeof(float));
        return eng->sample(slots[i], sp);
    };

    /* Correctness check: one batched step vs one serialized step per slot.
     * Both start from the same post-prefill state, run one forward, and we
     * compare the argmax token per slot.  Reassociation drift may flip a
     * near-tie argmax occasionally; we report mismatches but don't fail. */
    {
        std::vector<int32_t> batch_tok(N_slots), ser_tok(N_slots);
        eng->forward_batch(slots.data(), next_tok.data(), N_slots);
        for (int i = 0; i < N_slots; i++) batch_tok[i] = sample_slot(i);
        /* Reset and run serialized. */
        for (int i = 0; i < N_slots; i++) {
            eng->reset(slots[i]);
            eng->prefill(slots[i], ids.data(), (int)ids.size());
            next_tok[i] = eng->sample(slots[i], sp);
            eng->forward(slots[i], next_tok[i]);
            ser_tok[i] = eng->sample(slots[i], sp);
        }
        int mism = 0;
        for (int i = 0; i < N_slots; i++)
            if (batch_tok[i] != ser_tok[i]) mism++;
        std::printf("correctness: %d/%d slots match serialized argmax (%d mismatch)\n",
                    N_slots - mism, N_slots, mism);
        /* Re-prepare for the timed batched run: reset + prefill all slots. */
        for (int i = 0; i < N_slots; i++) {
            eng->reset(slots[i]);
            eng->prefill(slots[i], ids.data(), (int)ids.size());
            next_tok[i] = eng->sample(slots[i], sp);
        }
    }

    /* Warm up (2 steps). */
    for (int w = 0; w < 2; w++) {
        eng->forward_batch(slots.data(), next_tok.data(), N_slots);
        for (int i = 0; i < N_slots; i++) next_tok[i] = sample_slot(i);
    }

    /* Timed batched run.  forward_batch submits async; we drain via a final
     * reset-prefill-free dummy: simplest is to read a USM value the last
     * kernel wrote.  Easier: forward_batch internally does q.wait() when
     * QXMX_PROFILE is unset (non-profiling mode), so each call already
     * drained.  No explicit wait needed. */
    auto t0 = std::chrono::steady_clock::now();
    for (int s = 0; s < n_gen; s++)
        eng->forward_batch(slots.data(), next_tok.data(), N_slots);
    auto t1 = std::chrono::steady_clock::now();
    double batch_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
    double batch_per_step = batch_ms / n_gen;
    double batch_agg_tps = N_slots / (batch_per_step / 1000.0);

    std::printf("\n=== batched (FFN as one M=%d gemm_gpu_v4) ===\n", N_slots);
    std::printf("  total: %.1f ms over %d steps\n", batch_ms, n_gen);
    std::printf("  per-step: %.1f ms\n", batch_per_step);
    std::printf("  aggregate: %.1f tok/s (per-slot %.1f tok/s)\n",
                batch_agg_tps, batch_agg_tps / N_slots);
    std::fflush(stdout);

    /* === Serialized baseline: N separate forward() calls per step ===
     * Reset slots and re-prefill.  Wrapped in try/catch so a DEVICE_LOST
     * here doesn't lose the batched result above. */
    double ser_ms = 0, ser_per_step = 0, ser_agg_tps = 0;
    try {
        for (int i = 0; i < N_slots; i++) {
            eng->reset(slots[i]);
            eng->prefill(slots[i], ids.data(), (int)ids.size());
            next_tok[i] = eng->sample(slots[i], sp);
        }
        auto t2 = std::chrono::steady_clock::now();
        for (int s = 0; s < n_gen; s++) {
            for (int i = 0; i < N_slots; i++) {
                eng->forward(slots[i], next_tok[i]);
                next_tok[i] = eng->sample(slots[i], sp);
            }
        }
        auto t3 = std::chrono::steady_clock::now();
        ser_ms = std::chrono::duration<double, std::milli>(t3 - t2).count();
        ser_per_step = ser_ms / n_gen;
        ser_agg_tps = N_slots / (ser_per_step / 1000.0);

        std::printf("\n=== serialized (N=%d separate forward() calls per step) ===\n", N_slots);
        std::printf("  total: %.1f ms over %d steps\n", ser_ms, n_gen);
        std::printf("  per-step: %.1f ms\n", ser_per_step);
        std::printf("  aggregate: %.1f tok/s (per-slot %.1f tok/s)\n",
                    ser_agg_tps, ser_agg_tps / N_slots);

        std::printf("\n=== speedup ===\n");
        std::printf("  batched / serialized: %.2fx aggregate tok/s\n", batch_agg_tps / ser_agg_tps);
        std::printf("  per-step: %.1f ms -> %.1f ms (-%.0f%%)\n",
                    ser_per_step, batch_per_step,
                    (ser_per_step - batch_per_step) * 100.0 / ser_per_step);
    } catch (const std::exception& e) {
        std::fprintf(stderr, "\nserialized baseline crashed: %s\n", e.what());
        std::fprintf(stderr, "(batched result above is still valid)\n");
    }
    std::fflush(stdout);

    for (int i = 0; i < N_slots; i++) eng->close_slot(slots[i]);
    return 0;
}