// metal/src/strata_window_cli.cpp - Direct zero-bloat Metal runner for Strata's Window + ExpertStream + PleTable.
// Supports both:
//   1. Prefill (PromptGemm) + decode (Window::run + ExpertStream F_NOCACHE pread + speculative router lookahead)
//   2. Teacher-forced full-vocab (248,320) FP64 KLD, PPL, and top-1 agreement evaluation (--logits-out / --logits-ref)

#include "strata/core/conversation_cache.hpp"
#include "strata/core/layer.hpp"
#include "strata/core/session.hpp"
#include "strata/kernels/ngram.hpp"
#include "strata/spec/suffix_drafter.hpp"
#include "strata/spec/spec_gate.hpp"
#include "engine/scr_state.hpp"
#include "engine/suffix_window.hpp"
#include "strata/core/scr.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <filesystem>
#include <fstream>
#include <sstream>
#include <string>
#include <sys/stat.h>
#include <unistd.h>
#include <vector>

namespace strata::metal_engine {
void set_model_shard(const char* shard1);
void* window_create();
void window_destroy(void* w);
bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::SessionState& ss, int max_t, int64_t max_cells,
                 int64_t* n_vocab, std::string& err);
void window_mrope_host_set(const int32_t* host_table);
bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* const* embd_rows,
                const std::function<bool(const float*, int64_t, int64_t, std::string&)>& on_chunk, std::string& err);
bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0,
                     const std::function<bool(const float*, int64_t, int64_t, std::string&)>& on_chunk, std::string& err);
bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err);
bool window_commit(void* w, int n_keep, std::string& err);
bool window_copy_logits(void* w, int t, float* host);
void window_reset_stats(void* w);
void window_reset();
bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err);
bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err);
bool window_session_save(const char* path, const int32_t* ids, int64_t n_cells, std::string& err, int64_t base_cells = 0);
bool window_session_load(const char* path, std::vector<int32_t>& ids, int64_t* n_cells, std::string& err, const char* parent_path = nullptr);
bool window_scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err);
bool mtp_load(const std::string& rt_dir, int max_t, std::string& err);
bool mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err);
bool mtp_draft(int T, const int32_t* tokens, int64_t p, int a, int max_steps, int32_t* drafts, float* probs, float min_p,
               int* n_drafts, std::string& err);
}  // namespace strata::metal_engine

static double now_ms() {
    return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now().time_since_epoch()).count();
}

// One policy instance per decode: a new prompt does not inherit another workload's timings.
struct SpecControl {
    strata::spec::SpecGate gate;
    bool enabled = true, full = false;
    int mtp_used = 0;

    SpecControl() {
        const char* mode = std::getenv("STRATA_SPEC_GATE");
        enabled = mode && std::strcmp(mode, "0") != 0;
        full = mode && std::strcmp(mode, "1") == 0;
    }

    int propose(strata::spec::SuffixDrafter& suffix, int32_t* input,
                const int32_t* mtp, const float* probs, int n_mtp, int suffix_max,
                int rem, int64_t pos, int64_t cells) {
        int nd = strata::spec::suffix_limit(n_mtp, rem, pos, cells);
        if (full) nd = gate.spec_gate_pick(0, nd, probs);
        mtp_used = nd;
        for (int j = 0; j < nd; ++j) input[1 + j] = mtp[j];
        int limit = strata::spec::suffix_limit(std::max(suffix_max, nd), rem, pos, cells) - nd;
        if (suffix_max <= 0 || limit <= 0) return nd;
        if (enabled) limit = gate.spec_gate_pick(1, limit, nullptr, nd, probs);
        if (limit <= 0) return nd;
        const double begin = now_ms();
        const int sk = nd > 0 ? suffix.propose_after(input + 1, nd, limit, input + 1 + nd)
                              : suffix.propose(limit, input + 1);
        if (sk > 0) gate.spec_gate_draft_cost(1, (now_ms() - begin) / (1000 * sk));
        return nd + sk;
    }

    void result(int nd, int keep) {
        const int accepted = keep - 1;
        gate.spec_gate_result(0, std::min(mtp_used, keep), std::min(mtp_used, accepted));
        if (accepted >= mtp_used)
            gate.spec_gate_result(1, std::min(nd - mtp_used, keep - mtp_used), accepted - mtp_used);
    }

    void describe() const {
        if (enabled) std::fprintf(stderr, "%s\n", gate.spec_gate_describe().c_str());
    }
};

static double logsumexp_f64(const float* lg, int64_t n, int32_t* argmax_out) {
    float mx = lg[0];
    int32_t am = 0;
    for (int64_t i = 1; i < n; ++i) {
        if (lg[i] > mx) {
            mx = lg[i];
            am = (int32_t) i;
        }
    }
    if (argmax_out) *argmax_out = am;
    double sum = 0.0;
    for (int64_t i = 0; i < n; ++i) {
        sum += std::exp((double) lg[i] - (double) mx);
    }
    return (double) mx + std::log(sum);
}

static std::vector<int32_t> read_ids(const std::string& path) {
    std::vector<int32_t> v;
    std::ifstream f(path);
    int64_t x;
    while (f >> x) v.push_back((int32_t) x);
    return v;
}

// teacher-forced next-token distributions, as log-probabilities: inputs[i] runs at pos + i and predicts targets[i]
struct TfRun {
    std::vector<float> lp;   // [n][n_vocab]
    double nll = 0;
    double ms = 0;
};
static bool tf_run(void* w, const std::vector<int32_t>& inputs, const std::vector<int32_t>& targets, int64_t pos, int64_t n_vocab,
                   TfRun& out, std::string& err) {
    namespace me = strata::metal_engine;
    const size_t n = inputs.size();
    out.lp.assign(n * (size_t) n_vocab, 0.0f);
    out.nll = 0;
    const double t0 = now_ms();
    for (size_t i = 0; i < n; ++i) {
        int32_t tok = inputs[i], next = 0;
        if (!me::window_run(w, 1, &tok, pos + (int64_t) i, true, &next, err) || !me::window_commit(w, 1, err)) return false;
        float* lg = out.lp.data() + i * (size_t) n_vocab;
        if (!me::window_copy_logits(w, 0, lg)) { err = "copy_logits failed"; return false; }
        const double lse = logsumexp_f64(lg, n_vocab, nullptr);
        for (int64_t v = 0; v < n_vocab; ++v) lg[v] = (float) ((double) lg[v] - lse);
        out.nll -= lg[targets[i]];
    }
    out.ms = now_ms() - t0;
    return true;
}
// KL(ref || cand) per position, nats
static void tf_compare(const TfRun& ref, const TfRun& c, int64_t n_vocab, double& mean, double& p90, double& mx, double& top1) {
    const size_t n = ref.lp.size() / (size_t) n_vocab;
    std::vector<double> k(n);
    int agree = 0;
    for (size_t i = 0; i < n; ++i) {
        const float* a = ref.lp.data() + i * (size_t) n_vocab;
        const float* b = c.lp.data() + i * (size_t) n_vocab;
        double s = 0;
        int64_t ia = 0, ib = 0;
        for (int64_t v = 0; v < n_vocab; ++v) {
            s += std::exp((double) a[v]) * ((double) a[v] - (double) b[v]);
            if (a[v] > a[ia]) ia = v;
            if (b[v] > b[ib]) ib = v;
        }
        k[i] = std::max(0.0, s);
        agree += ia == ib;
    }
    mean = 0;
    for (double x : k) mean += x;
    mean /= (double) n;
    std::vector<double> s = k;
    std::sort(s.begin(), s.end());
    p90 = s[(size_t) (0.9 * (double) (n - 1))];
    mx = s.back();
    top1 = 100.0 * agree / (double) n;
}

// --scr-test: the SCR splice [A B C] -> [A B' C] against a fresh prefill of [A B' C], scored teacher-forced on D.
// C's last token is D's first input (so every variant's state is "A, the middle span, C minus its last token").
// Variants, each from the saved [A B C] session:
//   orig       no edit: the old context (how much B -> B' matters at all)
//   keep J     C's cells relocated, B' + C[:J] prefilled from the state at |A|, then the running state of [A B C]
//              restored (the CLM rule: the post-edit GDN keeps having read B)
//   gdnex J    as keep, but the running state of the fresh [A B' C] restored: only the KV staleness is left
//   lossy J    as keep, without the restore: the GDN state has read A B' C[:J] only
// J is raised to the junction block's minimum; lossy with J = |C| is the fresh prefill again (a sanity row: KL 0).
static int scr_test(void* w, strata::core::SessionState& ss, int64_t n_vocab, const std::string& pa, const std::string& pb,
                    const std::string& pbp, const std::string& pc, const std::string& pd, const std::string& js, int32_t pad,
                    int n_eval, const std::string& tmp) {
    namespace me = strata::metal_engine;
    std::string err;
    const std::vector<int32_t> A = read_ids(pa), B = read_ids(pb), D = read_ids(pd);
    std::vector<int32_t> Bp = read_ids(pbp), C = read_ids(pc);
    if (A.size() < 4 || B.empty() || Bp.empty() || C.size() < 2 || D.size() < 2) { std::fprintf(stderr, "scr-test: empty span\n"); return 2; }
    while ((Bp.size() - B.size()) % 4 != 0) Bp.push_back(pad);   // the shift must be a multiple of 4 cells
    const int32_t c_last = C.back();
    C.pop_back();
    const int64_t nA = (int64_t) A.size(), nB = (int64_t) B.size(), nBp = (int64_t) Bp.size(), nC = (int64_t) C.size();
    std::vector<int32_t> inputs{c_last}, targets;
    for (size_t i = 0; i + 1 < D.size() && (int) inputs.size() < n_eval; ++i) inputs.push_back(D[i]);
    targets.assign(D.begin(), D.begin() + (ptrdiff_t) inputs.size());
    auto prefill = [&](const std::vector<int32_t>& ids, int64_t pos) {
        std::vector<int64_t> v(ids.begin(), ids.end());
        return v.empty() || me::prompt_run_rows(v.data(), (int64_t) v.size(), pos, nullptr, err);
    };
    auto cat = [](std::vector<int32_t> a, const std::vector<int32_t>& b, size_t nb) { a.insert(a.end(), b.begin(), b.begin() + (ptrdiff_t) nb); return a; };
    auto fresh = [&] { me::window_reset(); ss.ple_prev[0] = -1; ss.ple_prev[1] = -1; };
    const int32_t prev_c[2] = {C.size() >= 2 ? C[C.size() - 2] : Bp.back(), C.back()};
    auto fail = [&](const char* what) { std::fprintf(stderr, "scr-test: %s: %s\n", what, err.c_str()); return 1; };

    // the reference: [A B' C] prefilled fresh
    strata::core::ConversationCheckpoint ckA, ckABC, ckRef;
    int32_t prevA[2];
    fresh();
    if (!prefill(A, 0) || !me::window_checkpoint_save(ckA, err)) return fail("prefill A");
    prevA[0] = ss.ple_prev[0]; prevA[1] = ss.ple_prev[1];
    double t0 = now_ms();
    if (!prefill(cat(Bp, C, C.size()), nA)) return fail("prefill B' C");
    const double ref_ms = now_ms() - t0;
    if (!me::window_checkpoint_save(ckRef, err)) return fail("checkpoint");
    TfRun ref;
    if (!tf_run(w, inputs, targets, nA + nBp + nC, n_vocab, ref, err)) return fail("ref eval");
    // the old context [A B C], saved as a session
    fresh();
    if (!prefill(cat(A, B, B.size()), 0) || !prefill(C, nA + nB) || !me::window_checkpoint_save(ckABC, err)) return fail("prefill A B C");
    const std::vector<int32_t> abc = cat(cat(A, B, B.size()), C, C.size());
    t0 = now_ms();
    if (!me::window_session_save(tmp.c_str(), abc.data(), (int64_t) abc.size(), err)) return fail("session save");
    const double save_ms = now_ms() - t0;
    TfRun orig;
    if (!tf_run(w, inputs, targets, nA + nB + nC, n_vocab, orig, err)) return fail("orig eval");
    const double n = (double) inputs.size();
    std::printf("SCR_SPANS A=%lld B=%lld B'=%lld (padded) C=%lld eval=%zu delta=%lld session_save_ms=%.1f ref_prefill_ms(B'+C)=%.1f ref_ppl=%.4f\n",
                (long long) nA, (long long) nB, (long long) nBp, (long long) nC, inputs.size(), (long long) (nBp - nB), save_ms, ref_ms,
                std::exp(ref.nll / n));
    double km, kp, kx, t1;
    tf_compare(ref, orig, n_vocab, km, kp, kx, t1);
    std::printf("SCR_ROW variant=orig J=- mean_kld=%.6f p90_kld=%.6f max_kld=%.6f top1=%.2f ppl=%.4f\n", km, kp, kx, t1, std::exp(orig.nll / n));

    std::vector<int64_t> Js;
    {
        std::stringstream s(js);
        std::string x;
        while (std::getline(s, x, ',')) Js.push_back(std::atoll(x.c_str()));
        Js.push_back(nC);
    }
    for (const char* mode : {"keep", "gdnex", "lossy"})
        for (int64_t J : Js) {
            if (J == nC && std::string(mode) != "lossy") continue;
            const int64_t Je = std::min(nC, std::max(J, scr::scr_min_junction(nA + nBp)));
            std::vector<int32_t> ids;
            int64_t cells = 0;
            t0 = now_ms();
            if (!me::window_session_load(tmp.c_str(), ids, &cells, err)) return fail("session load");
            const double load_ms = now_ms() - t0;
            t0 = now_ms();
            if (!me::window_scr_relocate(nA + nB, nA + nBp, nC, err)) return fail("relocate");
            const double rel_ms = now_ms() - t0;
            t0 = now_ms();
            if (!me::window_checkpoint_restore(ckA, err)) return fail("restore A");
            ss.ple_prev[0] = prevA[0]; ss.ple_prev[1] = prevA[1];
            if (!prefill(cat(Bp, C, (size_t) Je), nA)) return fail("prefill B' C[:J]");
            if (std::string(mode) == "keep" && !me::window_checkpoint_restore(ckABC, err)) return fail("restore ABC");
            if (std::string(mode) == "gdnex" && !me::window_checkpoint_restore(ckRef, err)) return fail("restore ref");
            ss.ple_prev[0] = prev_c[0]; ss.ple_prev[1] = prev_c[1];   // token ids are never stale
            const double splice_ms = now_ms() - t0;
            TfRun r;
            if (!tf_run(w, inputs, targets, nA + nBp + nC, n_vocab, r, err)) return fail("eval");
            tf_compare(ref, r, n_vocab, km, kp, kx, t1);
            std::printf("SCR_ROW variant=%s J=%lld mean_kld=%.6f p90_kld=%.6f max_kld=%.6f top1=%.2f ppl=%.4f load_ms=%.1f relocate_ms=%.1f "
                        "splice_prefill_ms=%.1f speedup_vs_ref=%.2f\n", mode, (long long) Je, km, kp, kx, t1, std::exp(r.nll / n), load_ms,
                        rel_ms, splice_ms, ref_ms / std::max(1e-3, rel_ms + splice_ms));
            std::fflush(stdout);
            if (std::string(mode) == "keep" && (J == 0 || J == 8)) {
                // Connectome loop check (PLAN.md §2.9): restore spliced state and free-generate 128 tokens greedy
                if (!me::window_session_load(tmp.c_str(), ids, &cells, err)) return fail("loop session load");
                if (!me::window_scr_relocate(nA + nB, nA + nBp, nC, err)) return fail("loop relocate");
                if (!me::window_checkpoint_restore(ckA, err)) return fail("loop restore A");
                ss.ple_prev[0] = prevA[0]; ss.ple_prev[1] = prevA[1];
                if (!prefill(cat(Bp, C, (size_t) Je), nA)) return fail("loop prefill");
                if (!me::window_checkpoint_restore(ckABC, err)) return fail("loop restore ABC");
                ss.ple_prev[0] = prev_c[0]; ss.ple_prev[1] = prev_c[1];
                std::vector<int32_t> gen_toks;
                int32_t cur = inputs[0], nxt = 0;
                int64_t gpos = nA + nBp + nC;
                for (int step = 0; step < 128; ++step) {
                    if (!me::window_run(w, 1, &cur, gpos, true, &nxt, err) || !me::window_commit(w, 1, err)) break;
                    gen_toks.push_back(nxt);
                    cur = nxt;
                    ++gpos;
                }
                int rep4 = 0, tot4 = std::max(0, (int) gen_toks.size() - 3);
                for (int i = 0; i < tot4; ++i) {
                    for (int k = 0; k < i; ++k) {
                        if (gen_toks[i] == gen_toks[k] && gen_toks[i+1] == gen_toks[k+1] &&
                            gen_toks[i+2] == gen_toks[k+2] && gen_toks[i+3] == gen_toks[k+3]) {
                            ++rep4;
                            break;
                        }
                    }
                }
                std::printf("SCR_LOOP_CHECK variant=keep J=%lld gen_tokens=%zu repeated_4grams=%d/%d (%.2f%%)\n",
                            (long long) Je, gen_toks.size(), rep4, tot4, tot4 > 0 ? 100.0 * rep4 / tot4 : 0.0);
                std::fflush(stdout);
            }
        }
    std::remove(tmp.c_str());
    return 0;
}


#include <iostream>
#include <random>

// sample_logits' distribution: the top-k logits, the repetition / presence penalties over the last 128 tokens of
// `recent`, the temperature, min-p and top-p.  tok / p: its tokens and unnormalized probabilities (sum: theirs); empty
// for greedy decoding.
struct Dist {
    std::vector<int32_t> tok;
    std::vector<double> p;
    double sum = 0.0;
};
static void sample_dist(const float* lg, int64_t n_vocab, float temp, int top_k, float top_p, float min_p, float pen_rep,
                        float pen_pres, const std::vector<int64_t>& recent, Dist& d) {
    d.tok.clear();
    d.p.clear();
    d.sum = 0.0;
    if (temp <= 0.0f || top_k == 1) return;
    const int k = top_k > 0 ? std::min<int>((int) n_vocab, top_k) : 64;
    std::vector<std::pair<float, int32_t>> cand;
    cand.reserve(256);
    float cutoff = -INFINITY;
    for (int64_t i = 0; i < n_vocab; ++i) {
        float v = lg[i];
        if (std::isnan(v) || v <= cutoff) continue;
        cand.push_back({v, (int32_t) i});
        if ((int) cand.size() >= k * 4) {
            std::nth_element(cand.begin(), cand.begin() + k - 1, cand.end(),
                             [](const auto& a, const auto& b) { return a.first > b.first; });
            cand.resize((size_t) k);
            cutoff = cand.back().first;
        }
    }
    if (cand.empty()) return;
    if (pen_rep != 1.0f || pen_pres != 0.0f) {
        const size_t r0 = recent.size() > 128 ? recent.size() - 128 : 0;
        for (auto& kv : cand) {
            bool seen = false;
            for (size_t ri = r0; ri < recent.size(); ++ri) if (recent[ri] == kv.second) { seen = true; break; }
            if (seen) {
                if (pen_rep > 1.0f) kv.first = kv.first > 0.0f ? kv.first / pen_rep : kv.first * pen_rep;
                kv.first -= pen_pres;
            }
        }
    }
    std::partial_sort(cand.begin(), cand.begin() + std::min<size_t>(cand.size(), (size_t) k), cand.end(),
                      [](const auto& a, const auto& b) { return a.first > b.first; });
    if ((int) cand.size() > k) cand.resize((size_t) k);
    const float mx = cand[0].first;
    double sum = 0.0;
    std::vector<double> probs(cand.size());
    for (size_t i = 0; i < cand.size(); ++i) {
        probs[i] = std::exp((double) (cand[i].first - mx) / (double) temp);
        sum += probs[i];
    }
    if (sum <= 0.0) {
        d.tok = {cand[0].second};
        d.p = {1.0};
        d.sum = 1.0;
        return;
    }
    if (min_p > 0.0f && cand.size() > 1) {
        const double p_best = probs[0] / sum;
        size_t keep = 1;
        while (keep < cand.size() && (probs[keep] / sum) >= p_best * (double) min_p) ++keep;
        cand.resize(keep);
        probs.resize(keep);
        sum = 0.0;
        for (double p : probs) sum += p;
    }
    if (top_p > 0.0f && top_p < 1.0f) {
        double cum = 0.0;
        size_t keep = cand.size();
        for (size_t i = 0; i < cand.size(); ++i) {
            cum += probs[i] / sum;
            if (cum >= (double) top_p) {
                keep = i + 1;
                break;
            }
        }
        cand.resize(keep);
        probs.resize(keep);
        sum = 0.0;
        for (double p : probs) sum += p;
    }
    d.tok.resize(cand.size());
    for (size_t i = 0; i < cand.size(); ++i) d.tok[i] = cand[i].second;
    d.p = std::move(probs);
    d.sum = sum;
}

// a draw from d, token `skip` left out (-1: none)
static int32_t dist_draw(const Dist& d, std::mt19937& rng, int32_t skip = -1) {
    double tot = d.sum;
    for (size_t i = 0; i < d.tok.size(); ++i) if (d.tok[i] == skip) tot -= d.p[i];
    if (tot <= 0.0) return skip;
    std::uniform_real_distribution<double> dist(0.0, tot);
    double r = dist(rng);
    int32_t last = d.tok[0];
    for (size_t i = 0; i < d.tok.size(); ++i) {
        if (d.tok[i] == skip) continue;
        last = d.tok[i];
        r -= d.p[i];
        if (r <= 0.0) return d.tok[i];
    }
    return last;
}

static int32_t sample_logits(const float* lg, int64_t n_vocab, int32_t greedy_tok, float temp, int top_k, float top_p, float min_p,
                             float pen_rep, float pen_pres, const std::vector<int64_t>& recent, std::mt19937& rng) {
    Dist d;
    sample_dist(lg, n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, recent, d);
    return d.tok.empty() ? greedy_tok : dist_draw(d, rng);
}

// Speculative sampling (Leviathan et al. 2023, sec. 2.3) of a verified window: input[0] the last token, input[1..nd] the
// drafts, output[i] the window's greedy pick after input[0..i].  The drafts are deterministic (q = one-hot), so draft
// x after row i is accepted with probability p_i(x) (p_i: sample_dist of row i's logits, history `hist` + input[0..i]);
// a rejected row samples p_i without x (the residual max(p - q, 0) normalized), and a fully accepted window samples its
// last row.  The tokens emitted are distributed as sample_logits one at a time.  Returns keep (output[0..keep-1] the
// tokens, input[0..keep-1] appended to hist).
static int spec_sample(void* w, const int32_t* input, int32_t* output, int nd, int64_t n_vocab, float temp, int top_k, float top_p,
                       float min_p, float pen_rep, float pen_pres, std::vector<int64_t>& hist, std::vector<float>& lg,
                       std::mt19937& rng) {
    namespace me = strata::metal_engine;
    Dist d;
    std::uniform_real_distribution<double> uni(0.0, 1.0);
    for (int i = 0;; ++i) {
        hist.push_back(input[i]);
        me::window_copy_logits(w, i, lg.data());
        sample_dist(lg.data(), n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, hist, d);
        if (d.tok.empty()) {   // greedy (top_k 1): the drafts must be the picks
            if (i == nd || output[i] != input[i + 1] || strata::spec::suffix_stop(output[i])) return i + 1;
            continue;
        }
        if (i == nd) {
            output[i] = dist_draw(d, rng);
            return i + 1;
        }
        const int32_t x = input[i + 1];
        static const bool typical = std::getenv("STRATA_SPEC_SAMPLE") && std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "typical";
        double px = 0.0;
        for (size_t k = 0; k < d.tok.size(); ++k) if (d.tok[k] == x) px = d.p[k];
        const bool accept = typical ? (!d.tok.empty() && (d.tok[0] == x || px >= std::max(0.05 * d.sum, 0.25 * d.p[0])))
                                    : (uni(rng) * d.sum < px);
        if (accept) {
            output[i] = x;
            if (strata::spec::suffix_stop(x)) return i + 1;
            continue;
        }
        output[i] = typical ? dist_draw(d, rng) : dist_draw(d, rng, x);
        return i + 1;
    }
}

struct DiskCache {
    struct Entry {
        std::string path;
        std::string parent_path;
        std::vector<int64_t> ids;
        std::vector<int64_t> pads;
        int64_t base_cells = 0;
        size_t bytes = 0;
        bool pinned = false;
    };

    static constexpr size_t kLearnMinTokens = 32;
    static constexpr size_t kLearnMaxTokens = 65536;
    static constexpr size_t kMaxPins = 16;
    static constexpr size_t kSeenMax = 32;

    std::string dir;
    size_t limit_bytes = 0;
    std::vector<Entry> entries;
    std::deque<std::vector<int64_t>> seen_prompts;

    bool enabled() const { return !dir.empty(); }

    static uint64_t fnv1a(const int64_t* data, size_t n) {
        uint64_t h = 1469598103934665603ull;
        const auto* p = reinterpret_cast<const uint8_t*>(data);
        const size_t nb = n * sizeof(int64_t);
        for (size_t i = 0; i < nb; ++i) h = (h ^ p[i]) * 1099511628211ull;
        return h;
    }

    static bool is_prefix(const std::vector<int64_t>& a, const std::vector<int64_t>& b) {
        return a.size() <= b.size() && std::equal(a.begin(), a.end(), b.begin());
    }

    static size_t common_prefix(const std::vector<int64_t>& a, const std::vector<int64_t>& b, size_t cap) {
        const size_t lim = std::min({a.size(), b.size(), cap});
        size_t i = 0;
        while (i < lim && a[i] == b[i]) ++i;
        return i;
    }

    static void read_pads(const std::string& ppath, size_t n_ids, int32_t pad_tok, const std::vector<int64_t>& ids,
                          std::vector<int64_t>& out_pads) {
        out_pads.clear();
        FILE* pf = std::fopen(ppath.c_str(), "rb");
        if (!pf) return;
        uint64_t np = 0;
        if (std::fread(&np, sizeof(np), 1, pf) == 1 && np <= n_ids) {
            std::vector<int64_t> p((size_t) np);
            if (std::fread(p.data(), sizeof(int64_t), (size_t) np, pf) == (size_t) np) {
                bool ok = true;
                for (int64_t at : p) {
                    if (at < 0 || (size_t) at >= n_ids || ids[(size_t) at] != pad_tok) ok = false;
                }
                if (ok) out_pads = std::move(p);
            }
        }
        std::fclose(pf);
    }

    static void write_pads(const std::string& ppath, const std::vector<int64_t>& pads) {
        if (pads.empty()) {
            std::remove(ppath.c_str());
            return;
        }
        if (FILE* pf = std::fopen(ppath.c_str(), "wb")) {
            const uint64_t np = (uint64_t) pads.size();
            std::fwrite(&np, sizeof(np), 1, pf);
            std::fwrite(pads.data(), sizeof(int64_t), pads.size(), pf);
            std::fclose(pf);
        }
    }

    static std::string read_parent(const std::string& ppath) {
        std::ifstream in(ppath, std::ios::binary);
        if (!in) return "";
        std::string s;
        std::getline(in, s);
        while (!s.empty() && (s.back() == '\r' || s.back() == '\n')) s.pop_back();
        return s;
    }

    static void write_parent(const std::string& ppath, const std::string& parent_path) {
        if (parent_path.empty()) {
            std::remove(ppath.c_str());
            return;
        }
        std::ofstream out(ppath, std::ios::binary | std::ios::trunc);
        if (out) out << parent_path << "\n";
    }

    static void remove_files(const std::string& p) {
        namespace fs = std::filesystem;
        std::error_code ec;
        fs::remove(p, ec);
        fs::remove(p + ".pads", ec);
        fs::remove(p + ".pin", ec);
        fs::remove(p + ".parent", ec);
    }

    void load_seen() {
        seen_prompts.clear();
        if (dir.empty()) return;
        const std::string spath = dir + "/prompts-seen.bin";
        FILE* f = std::fopen(spath.c_str(), "rb");
        if (!f) return;
        uint32_t magic = 0, count = 0;
        if (std::fread(&magic, 4, 1, f) == 1 && magic == 0x5345454Eu &&
            std::fread(&count, 4, 1, f) == 1 && count <= kSeenMax) {
            for (uint32_t i = 0; i < count; ++i) {
                uint32_t len = 0;
                if (std::fread(&len, 4, 1, f) != 1 || len > kLearnMaxTokens) break;
                std::vector<int32_t> i32(len);
                if (len > 0 && std::fread(i32.data(), 4, len, f) != len) break;
                seen_prompts.emplace_back(i32.begin(), i32.end());
            }
        }
        std::fclose(f);
    }

    void save_seen() const {
        if (dir.empty()) return;
        const std::string spath = dir + "/prompts-seen.bin";
        const std::string tmp = spath + ".tmp";
        FILE* f = std::fopen(tmp.c_str(), "wb");
        if (!f) return;
        const uint32_t magic = 0x5345454Eu;
        const uint32_t count = (uint32_t) seen_prompts.size();
        bool ok = std::fwrite(&magic, 4, 1, f) == 1 && std::fwrite(&count, 4, 1, f) == 1;
        for (const auto& p : seen_prompts) {
            const uint32_t len = (uint32_t) std::min(p.size(), kLearnMaxTokens);
            std::vector<int32_t> i32(p.begin(), p.begin() + len);
            ok = ok && std::fwrite(&len, 4, 1, f) == 1 &&
                 (len == 0 || std::fwrite(i32.data(), 4, len, f) == len);
        }
        std::fclose(f);
        if (ok) std::rename(tmp.c_str(), spath.c_str());
        else std::remove(tmp.c_str());
    }

    void init(const std::string& cache_dir, int64_t limit_mib, int64_t max_cells, int32_t pad_tok) {
        namespace fs = std::filesystem;
        if (cache_dir.empty()) return;
        std::error_code ec;
        fs::create_directories(cache_dir, ec);
        if (ec) {
            std::fprintf(stderr, "strata metal cache: cannot create %s: %s\n", cache_dir.c_str(), ec.message().c_str());
            return;
        }
        dir = cache_dir;
        limit_bytes = limit_mib > 0 ? (size_t) limit_mib * 1024ull * 1024ull : 0;
        std::vector<std::pair<fs::file_time_type, Entry>> found;
        for (const auto& de : fs::directory_iterator(dir, ec)) {
            const auto p = de.path();
            const auto ext = p.extension().string();
            if (ext == ".partial" || ext == ".tmp") {
                fs::remove(p, ec);
                continue;
            }
            if (ext != ".kvc") continue;
            FILE* f = std::fopen(p.string().c_str(), "rb");
            if (!f) continue;
            scr::SessionHeader h, want;
            const bool hdr_ok = std::fread(&h, sizeof(h), 1, f) == 1;
            const int64_t n_ids = hdr_ok ? (int64_t) (uint32_t) h.n_ids : 0;
            const int64_t base_cells = (hdr_ok && (h.flags & (1u << 23))) ? (int64_t) ((uint64_t) h.n_ids >> 32) : 0;
            bool ok = hdr_ok &&
                      std::memcmp(h.magic, want.magic, 8) == 0 &&
                      h.version == want.version &&
                      h.n_cells > 0 && h.n_cells < max_cells &&
                      n_ids == h.n_cells &&
                      base_cells >= 0 && base_cells < h.n_cells && (base_cells & 3) == 0;
            std::vector<int32_t> i32;
            if (ok) {
                i32.resize((size_t) n_ids);
                ok = std::fread(i32.data(), 4, (size_t) n_ids, f) == (size_t) n_ids;
            }
            std::fclose(f);
            std::string parent_path = base_cells > 0 ? read_parent(p.string() + ".parent") : "";
            if (base_cells > 0 && parent_path.empty()) ok = false;
            if (!ok) {
                remove_files(p.string());
                continue;
            }
            Entry e;
            e.path = p.string();
            e.parent_path = std::move(parent_path);
            e.ids.assign(i32.begin(), i32.end());
            read_pads(e.path + ".pads", e.ids.size(), pad_tok, e.ids, e.pads);
            e.base_cells = base_cells;
            e.bytes = (size_t) de.file_size(ec);
            e.pinned = base_cells == 0 && fs::exists(e.path + ".pin", ec);
            found.push_back({de.last_write_time(ec), std::move(e)});
        }
        std::sort(found.begin(), found.end(), [](const auto& a, const auto& b) { return a.first < b.first; });
        for (auto& kv : found) entries.push_back(std::move(kv.second));
        prune_orphans();
        enforce_limits();
        load_seen();
        size_t total = 0, n_pin = 0, n_shared = 0;
        for (const auto& e : entries) {
            total += e.bytes;
            if (e.pinned) ++n_pin;
            if (e.base_cells > 0) ++n_shared;
        }
        std::fprintf(stderr, "strata metal cache: disk tier %s, %zu sessions (%zu pinned, %zu shared), %.2f GiB (limit %.1f GiB)\n",
                     dir.c_str(), entries.size(), n_pin, n_shared, total / 1073741824.0, limit_bytes / 1073741824.0);
    }

    void prune_orphans() {
        for (auto it = entries.begin(); it != entries.end();) {
            if (it->base_cells > 0) {
                bool parent_ok = false;
                for (const auto& p : entries) {
                    if (p.path == it->parent_path && p.base_cells == 0 &&
                        (int64_t) common_prefix(p.ids, it->ids, p.ids.size()) >= it->base_cells) {
                        parent_ok = true;
                        break;
                    }
                }
                if (!parent_ok) {
                    remove_files(it->path);
                    it = entries.erase(it);
                    continue;
                }
            }
            ++it;
        }
    }

    void enforce_limits() {
        namespace fs = std::filesystem;
        std::error_code ec;
        size_t n_pin = 0;
        for (const auto& e : entries) if (e.pinned) ++n_pin;
        for (size_t i = 0; n_pin > kMaxPins && i < entries.size(); ++i) {
            if (entries[i].pinned) {
                entries[i].pinned = false;
                fs::remove(entries[i].path + ".pin", ec);
                --n_pin;
            }
        }
        if (limit_bytes == 0) return;
        size_t total = 0;
        for (const auto& e : entries) total += e.bytes;
        for (auto it = entries.begin(); total > limit_bytes && entries.size() > 1 && it != entries.end();) {
            if (it->pinned) {
                ++it;
                continue;
            }
            const std::string gone = it->path;
            total -= it->bytes;
            remove_files(gone);
            it = entries.erase(it);
            for (auto cit = entries.begin(); cit != entries.end();) {
                if (cit->parent_path == gone) {
                    if (cit < it) --it;
                    total -= cit->bytes;
                    remove_files(cit->path);
                    cit = entries.erase(cit);
                } else {
                    ++cit;
                }
            }
        }
    }

    int64_t observe_prefix(const std::vector<int64_t>& ids, int64_t n_pre) {
        if (!enabled() || n_pre < (int64_t) kLearnMinTokens) return -1;
        const size_t cap = std::min<size_t>((size_t) n_pre, kLearnMaxTokens);
        std::vector<int64_t> cur(ids.begin(), ids.begin() + cap);
        size_t best_lcp = 0;
        bool added = false;
        for (auto& s : seen_prompts) {
            if (is_prefix(s, cur)) {
                s = cur;
                added = true;
                break;
            }
            if (is_prefix(cur, s)) {
                added = true;
                break;
            }
            const size_t lcp = common_prefix(cur, s, kLearnMaxTokens);
            const auto spans = strata::core::scr_detail::find_spans(
                s, cur, (int64_t) lcp, strata::core::ScrParams{32, 64, 16, 4});
            if (!spans.empty()) {
                // Same CLM conversation after a mid-prompt edit/compaction: update in place, do not pin the fold point
                s = cur;
                added = true;
                break;
            }
            if (lcp >= kLearnMinTokens && lcp < cur.size() && lcp < s.size()) {
                best_lcp = std::max(best_lcp, lcp);
            }
        }
        if (!added) {
            if (seen_prompts.size() >= kSeenMax) seen_prompts.pop_front();
            seen_prompts.push_back(std::move(cur));
        }
        save_seen();
        if (best_lcp >= kLearnMinTokens) {
            for (const auto& e : entries) {
                if (e.ids.size() == best_lcp && is_prefix(e.ids, ids)) return -1;
            }
            return (int64_t) best_lcp;
        }
        return -1;
    }

    void touch(size_t idx) {
        if (idx >= entries.size()) return;
        namespace fs = std::filesystem;
        std::error_code ec;
        fs::last_write_time(entries[idx].path, fs::file_time_type::clock::now(), ec);
        if (idx + 1 < entries.size()) {
            Entry e = std::move(entries[idx]);
            entries.erase(entries.begin() + (long) idx);
            entries.push_back(std::move(e));
        }
    }

    void save_current(const std::vector<int64_t>& ids, const std::vector<int64_t>& pads, bool pinned, std::string& err,
                      std::string* active_path = nullptr) {
        if (!enabled() || ids.size() < kLearnMinTokens) return;
        namespace fs = std::filesystem;
        std::error_code ec;
        for (size_t i = 0; i < entries.size(); ++i) {
            if (is_prefix(ids, entries[i].ids) && entries[i].pads == pads) {
                if (!pinned || (entries[i].ids.size() == ids.size() && entries[i].base_cells == 0)) {
                    if (pinned && !entries[i].pinned) {
                        entries[i].pinned = true;
                        if (FILE* pf = std::fopen((entries[i].path + ".pin").c_str(), "wb")) std::fclose(pf);
                    }
                    if (!pinned && active_path != nullptr) *active_path = entries[i].path;
                    touch(i);
                    return;
                }
            }
        }
        std::string parent_path;
        int64_t base_cells = 0;
        if (!pinned) {
            size_t best_len = 0;
            for (const auto& e : entries) {
                if (e.pinned && e.base_cells == 0 && e.pads.empty() &&
                    e.ids.size() >= kLearnMinTokens && ids.size() >= kLearnMinTokens + 4) {
                    const size_t lcp = common_prefix(e.ids, ids, std::min(e.ids.size(), ids.size() - 4));
                    const int64_t c0 = ((int64_t) lcp / 4) * 4;
                    if (c0 >= (int64_t) kLearnMinTokens && (size_t) c0 > best_len &&
                        (pads.empty() || pads.front() >= c0)) {
                        best_len = (size_t) c0;
                        base_cells = c0;
                        parent_path = e.path;
                    }
                }
            }
        }
        const uint64_t h = fnv1a(ids.data(), ids.size());
        char fname[96];
        std::snprintf(fname, sizeof(fname), "/%016llx-%zu.kvc", (unsigned long long) h, ids.size());
        const std::string path = dir + fname;
        std::vector<int32_t> i32(ids.begin(), ids.end());
        if (!strata::metal_engine::window_session_save(path.c_str(), i32.data(), (int64_t) i32.size(), err, base_cells)) return;
        write_pads(path + ".pads", pads);
        write_parent(path + ".parent", parent_path);
        if (pinned) {
            if (FILE* pf = std::fopen((path + ".pin").c_str(), "wb")) std::fclose(pf);
        } else {
            fs::remove(path + ".pin", ec);
        }
        const size_t sz = (size_t) fs::file_size(path, ec);
        const std::string prev_active = (!pinned && active_path != nullptr) ? *active_path : "";
        for (auto it = entries.begin(); it != entries.end();) {
            if (it->path == path ||
                (!it->pinned && (is_prefix(it->ids, ids) || (!prev_active.empty() && it->path == prev_active)))) {
                if (it->path != path) {
                    remove_files(it->path);
                }
                it = entries.erase(it);
            } else {
                ++it;
            }
        }
        if (!pinned && active_path != nullptr) *active_path = path;
        entries.push_back({path, parent_path, ids, pads, base_cells, sz, pinned});
        prune_orphans();
        enforce_limits();
    }

    int find_best(const std::vector<int64_t>& raw_ids, int64_t max_new, int64_t max_cells, int32_t pad_tok,
                  int64_t& out_reused) const {
        out_reused = 0;
        if (!enabled()) return -1;
        int best_idx = -1;
        int64_t best_score = 0;
        for (size_t i = 0; i < entries.size(); ++i) {
            const auto& e = entries[i];
            if (e.ids.empty()) continue;
            std::vector<int64_t> q_ids = raw_ids;
            if (!e.pads.empty()) {
                std::vector<int64_t> carried;
                strata::core::scr_carry(raw_ids, e.ids, e.pads, pad_tok, carried);
                if ((int64_t) carried.size() + max_new <= max_cells) q_ids = std::move(carried);
            }
            const int64_t q_n_pre = std::max<int64_t>(0, (int64_t) q_ids.size() - 1);
            int64_t cand = 0;
            if ((int64_t) e.ids.size() <= q_n_pre && is_prefix(e.ids, q_ids)) {
                cand = (int64_t) e.ids.size();
            } else if (e.ids.size() >= 32 && q_n_pre >= 32) {
                const strata::core::ScrPlan plan = strata::core::scr_plan(
                    e.ids, q_ids, {(int64_t) e.ids.size()}, strata::core::ScrParams{32, 64, 16, 4});
                if (!plan.empty() && (int64_t) q_ids.size() + plan.pads() + max_new <= max_cells) {
                    cand = plan.n_prefix + plan.reused();
                }
            }
            if (cand >= 32 && cand >= best_score) {
                best_score = cand;
                best_idx = (int) i;
            }
        }
        out_reused = best_score;
        return best_idx;
    }
};

static int serve_loop(void* w, strata::core::SessionState& ss, int64_t n_vocab, int64_t max_cells, int suffix_draft,
                      bool use_mtp, int spec, float spec_min_p, const std::string& cache_disk_dir, int64_t cache_disk_mib) {
    namespace me = strata::metal_engine;
    constexpr int64_t kScrPadToken = 198;   // '\n' pad for 4-cell QSA block alignment (ported from 01a11408 / scr.hpp)
    DiskCache disk_cache;
    disk_cache.init(cache_disk_dir, cache_disk_mib, max_cells, kScrPadToken);

    std::printf("INFO kv=%lld\n", (long long) max_cells);
    std::printf("READY %lld\n", (long long) max_cells);
    std::fflush(stdout);

    std::string line, err;
    struct CkSlot {
        std::vector<int64_t> ids;
        strata::core::ConversationCheckpoint ck;
        int32_t prev0 = -1, prev1 = -1;
    };
    std::vector<CkSlot> ck_ring;
    auto save_ck_at = [&](const std::vector<int64_t>& pfx) {
        if (pfx.empty()) return;
        for (size_t ci = 0; ci < ck_ring.size(); ++ci) {
            if (ck_ring[ci].ids == pfx) {
                CkSlot slot = std::move(ck_ring[ci]);
                ck_ring.erase(ck_ring.begin() + (long) ci);
                if (me::window_checkpoint_save(slot.ck, err)) {
                    slot.prev0 = ss.ple_prev[0];
                    slot.prev1 = ss.ple_prev[1];
                    ck_ring.push_back(std::move(slot));
                }
                return;
            }
        }
        if (ck_ring.size() >= 2) ck_ring.erase(ck_ring.begin());
        CkSlot slot;
        slot.ids = pfx;
        if (me::window_checkpoint_save(slot.ck, err)) {
            slot.prev0 = ss.ple_prev[0];
            slot.prev1 = ss.ple_prev[1];
            ck_ring.push_back(std::move(slot));
        }
    };
    std::vector<int64_t> cached_ids;
    std::vector<int64_t> cached_pads;
    std::vector<std::pair<int64_t, uint64_t>> cached_imgs;
    std::vector<int32_t> mrope_host((size_t) (max_cells + 16) * 3, 0);
    std::vector<float> img_rows;
    std::vector<const float*> row_ptr;
    std::string active_path;
    bool ram_dirty = false;
    std::vector<float> lg_buf((size_t) n_vocab);
    std::mt19937 rng(42);
    // temp > 0 with drafts: speculative sampling (spec_sample), else STRATA_SPEC_SAMPLE=greedy: the drafts kept while
    // they are the greedy picks and only the last token sampled (fast, but those tokens are greedy, unpenalized)
    static const bool spec_exact = std::getenv("STRATA_SPEC_SAMPLE") &&
                                   (std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "exact" ||
                                    std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "typical");

    while (std::getline(std::cin, line)) {
        if (line.empty()) continue;
        if (line == "QUIT" || line.rfind("QUIT ", 0) == 0) {
            if (ram_dirty && cached_ids.size() >= 32) disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path);
            break;
        }
        if (line == "STOP") continue;
        if (line == "VRAM" || line.rfind("VRAM ", 0) == 0) {
            std::printf("VRAM expert_slots=4129 expert_cache_mib=6656 expert_cache_full_mib=6656 vram_free_mib=13312\n");
            std::fflush(stdout);
            continue;
        }
        if (line.rfind("SAVE ", 0) == 0) {
            std::string path = line.substr(5);
            while (!path.empty() && (path.back() == '\r' || path.back() == ' ')) path.pop_back();
            const double t0 = now_ms();
            std::vector<int32_t> i32(cached_ids.begin(), cached_ids.end());
            if (!me::window_session_save(path.c_str(), i32.data(), (int64_t) i32.size(), err)) {
                std::printf("SERR io 0 %s\n", err.c_str());
            } else {
                DiskCache::write_pads(path + ".pads", cached_pads);
                struct stat st{};
                const size_t fbytes = ::stat(path.c_str(), &st) == 0 ? (size_t) st.st_size
                                                                     : scr::fixed_bytes() + scr::cell_bytes(i32.size());
                std::printf("SAVED %zu %zu %.1f\n", i32.size(), fbytes, now_ms() - t0);
            }
            std::fflush(stdout);
            continue;
        }
        if (line.rfind("RESTORE ", 0) == 0) {
            std::string path = line.substr(8);
            while (!path.empty() && (path.back() == '\r' || path.back() == ' ')) path.pop_back();
            const double t0 = now_ms();
            std::vector<int32_t> i32;
            int64_t cells = 0;
            if (!me::window_session_load(path.c_str(), i32, &cells, err)) {
                std::printf("SERR io 0 %s\n", err.c_str());
            } else {
                cached_ids.assign(i32.begin(), i32.end());
                DiskCache::read_pads(path + ".pads", cached_ids.size(), kScrPadToken, cached_ids, cached_pads);
                ram_dirty = false;
                struct stat st{};
                const size_t fbytes = ::stat(path.c_str(), &st) == 0 ? (size_t) st.st_size
                                                                     : scr::fixed_bytes() + scr::cell_bytes((size_t) cells);
                std::printf("RESTORED %lld %zu %.1f\n", (long long) cells, fbytes, now_ms() - t0);
            }
            std::fflush(stdout);
            continue;
        }
        const bool geni = line.rfind("GENI ", 0) == 0;
        if (!geni && line.rfind("GEN ", 0) != 0) continue;

        std::istringstream iss(line.substr(geni ? 5 : 4));
        int max_new = 512;
        if (!(iss >> max_new)) continue;
        float temp = 0.0f, top_p = 1.0f, min_p = 0.05f, pen_rep = 1.05f, pen_pres = 0.2f;
        int top_k = 40;
        std::vector<uint64_t> req_img_keys;
        std::string tok_str;
        while (iss >> tok_str) {
            auto eq = tok_str.find('=');
            if (eq != std::string::npos) {
                std::string k = tok_str.substr(0, eq), v = tok_str.substr(eq + 1);
                if (k == "temperature") temp = std::stof(v);
                else if (k == "top_p") top_p = std::stof(v);
                else if (k == "top_k") top_k = std::stoi(v);
                else if (k == "min_p") min_p = std::stof(v);
                else if (k == "penalty_repeat") pen_rep = std::stof(v);
                else if (k == "penalty_present") pen_pres = std::stof(v);
                else if (k == "seed") rng.seed((uint32_t) std::stoul(v));
                else if (k == "img_keys") {
                    req_img_keys.clear();
                    const char* q = v.c_str();
                    while (*q != '\0') {
                        char* qe = nullptr;
                        const unsigned long long hk = std::strtoull(q, &qe, 16);
                        if (qe == q) { req_img_keys.clear(); break; }
                        req_img_keys.push_back((uint64_t) hk);
                        q = *qe == ',' ? qe + 1 : qe;
                        if (*qe != ',' && *qe != '\0') { req_img_keys.clear(); break; }
                    }
                }
            } else {
                break;
            }
        }
        std::string emb_path;
        if (geni) {
            emb_path = tok_str;
            tok_str.clear();
            iss >> tok_str;
        }
        std::vector<int64_t> raw_ids;
        std::stringstream tss(tok_str);
        std::string item;
        while (std::getline(tss, item, ',')) {
            if (!item.empty()) raw_ids.push_back(std::stoll(item));
        }
        if (raw_ids.empty()) {
            std::printf("DONE 0 0 0.0 0.0 stop 0 0 0\n");
            std::fflush(stdout);
            continue;
        }
        if ((int64_t) raw_ids.size() + max_new > max_cells) {
            max_new = (int) std::max<int64_t>(1, max_cells - (int64_t) raw_ids.size() - 1);
        }

        const double t_pre0 = now_ms();
        std::vector<std::pair<int64_t, uint64_t>> req_imgs;
        row_ptr.clear();
        img_rows.clear();
        if (geni) {
            constexpr int64_t kImagePad = 248056;
            const int64_t n = (int64_t) raw_ids.size();
            row_ptr.assign((size_t) n, nullptr);
            const int64_t cells = (int64_t) mrope_host.size() / 3;
            auto put = [&](int64_t c, int64_t t, int64_t h, int64_t w) {
                mrope_host[(size_t) c * 3] = (int32_t) t;
                mrope_host[(size_t) c * 3 + 1] = (int32_t) h;
                mrope_host[(size_t) c * 3 + 2] = (int32_t) w;
            };
            struct Img { int64_t n, nx, ny; size_t off; };
            std::vector<Img> imgs;
            std::string ve;
            std::FILE* f = std::fopen(emb_path.c_str(), "rb");
            if (!f) ve = "cannot open " + emb_path;
            while (f && ve.empty()) {
                int32_t hdr[5];
                const size_t got = std::fread(hdr, sizeof(int32_t), 5, f);
                if (got == 0) break;
                if (got != 5 || hdr[0] != 0x31455653 || hdr[1] < 1 || hdr[2] < 1 || hdr[3] < 1 ||
                    (int64_t) hdr[2] * hdr[3] != hdr[1] || hdr[4] != 2560) {
                    ve = "bad embeddings file (expected strata-vision records of width 2560)";
                    break;
                }
                const size_t off = img_rows.size(), cnt = (size_t) hdr[1] * (size_t) hdr[4];
                img_rows.resize(off + cnt);
                if (std::fread(img_rows.data() + off, sizeof(float), cnt, f) != cnt) { ve = "short embeddings file"; break; }
                imgs.push_back({hdr[1], hdr[2], hdr[3], off});
            }
            if (f) std::fclose(f);
            int64_t p = 0, i = 0;
            size_t k = 0;
            while (ve.empty() && i < n) {
                if (raw_ids[(size_t) i] != kImagePad) { put(i, p, p, p); ++p; ++i; continue; }
                if (k >= imgs.size()) { ve = "the prompt has more images than the embeddings file"; break; }
                const Img& im = imgs[k++];
                uint64_t h = DiskCache::fnv1a(&im.n, 3);
                if (req_img_keys.size() == imgs.size()) {
                    h ^= req_img_keys[k - 1] * 1099511628211ull;
                } else {
                    h ^= DiskCache::fnv1a(reinterpret_cast<const int64_t*>(img_rows.data() + im.off),
                                          (size_t) im.n * (2560 * sizeof(float) / sizeof(int64_t)));
                }
                req_imgs.push_back({i, h});
                for (int64_t j = 0; j < im.n && ve.empty(); ++j) {
                    if (i + j >= n || raw_ids[(size_t) (i + j)] != kImagePad) {
                        ve = "image has fewer <|image_pad|> tokens than rows";
                    }
                }
                for (int64_t j = 0; j < im.n && ve.empty(); ++j) {
                    const int64_t y = j / im.nx, x = j % im.nx;
                    put(i + j, p, p + y, p + x);
                    row_ptr[(size_t) (i + j)] = img_rows.data() + im.off + (size_t) j * 2560;
                }
                i += im.n;
                p += std::max(im.nx, im.ny);
            }
            if (ve.empty() && k != imgs.size()) ve = "the embeddings file has more images than the prompt";
            if (ve.empty() && n > 0 && raw_ids[(size_t) (n - 1)] == kImagePad) ve = "the prompt cannot end in an image";
            for (int64_t c = n; ve.empty() && c < cells; ++c) put(c, p + (c - n), p + (c - n), p + (c - n));
            if (!ve.empty()) {
                me::window_mrope_host_set(nullptr);
                std::printf("ERR %s\n", ve.c_str());
                std::fflush(stdout);
                continue;
            }
            me::window_mrope_host_set(mrope_host.data());
        } else {
            me::window_mrope_host_set(nullptr);
        }

        std::vector<int64_t> ids = raw_ids;
        // Carry forward any 0..3 SCR alignment pads from prior turns that still lie in the shared prefix
        if (!geni && cached_imgs.empty() && !cached_pads.empty() && !cached_ids.empty()) {
            std::vector<int64_t> carried_ids;
            cached_pads = strata::core::scr_carry(raw_ids, cached_ids, cached_pads, kScrPadToken, carried_ids);
            if ((int64_t) carried_ids.size() + max_new <= max_cells) {
                ids = std::move(carried_ids);
            } else {
                cached_pads.clear();
            }
        }
        int64_t n_pre = std::max<int64_t>(0, (int64_t) ids.size() - 1);
        size_t lcp = 0;
        const size_t lim0 = std::min(cached_ids.size(), (size_t) n_pre);
        while (lcp < lim0 && cached_ids[lcp] == ids[lcp]) ++lcp;
        if (geni != !cached_imgs.empty()) {
            // If switching between text and vision (or if image hashes differ inside 0..lcp), clamp lcp before the first differing image
            size_t img_cap = lim0;
            const size_t n_im = std::max(cached_imgs.size(), req_imgs.size());
            for (size_t ii = 0; ii < n_im; ++ii) {
                if (ii >= cached_imgs.size() || ii >= req_imgs.size() || cached_imgs[ii] != req_imgs[ii]) {
                    const int64_t p_cut = std::min(ii < cached_imgs.size() ? cached_imgs[ii].first : (int64_t) lim0,
                                                   ii < req_imgs.size() ? req_imgs[ii].first : (int64_t) lim0);
                    img_cap = (size_t) std::max<int64_t>(0, p_cut);
                    break;
                }
            }
            lcp = std::min(lcp, img_cap);
        } else if (geni) {
            for (size_t ii = 0; ii < req_imgs.size(); ++ii) {
                if (ii >= cached_imgs.size() || cached_imgs[ii] != req_imgs[ii]) {
                    const int64_t p_cut = std::min(ii < cached_imgs.size() ? cached_imgs[ii].first : (int64_t) lcp,
                                                   req_imgs[ii].first);
                    lcp = std::min(lcp, (size_t) std::max<int64_t>(0, p_cut));
                    break;
                }
            }
        }

        // Estimate how many tokens the live RAM state + ck_ring can reuse before checking the NVMe tier
        int64_t ram_reused = 0;
        if (lcp == cached_ids.size() && lcp > 0) {
            ram_reused = (int64_t) lcp;
        } else if (!geni && cached_imgs.empty() && cached_ids.size() >= 32 && n_pre >= 32) {
            std::vector<int64_t> ck_lens = {(int64_t) cached_ids.size()};
            for (const auto& s : ck_ring) {
                if (!s.ids.empty() && s.ids.size() <= cached_ids.size() &&
                    std::equal(s.ids.begin(), s.ids.end(), cached_ids.begin())) {
                    ck_lens.push_back((int64_t) s.ids.size());
                }
            }
            const strata::core::ScrPlan plan = strata::core::scr_plan(
                cached_ids, ids, ck_lens, strata::core::ScrParams{32, 64, 16, 4});
            if (!plan.empty() && (int64_t) ids.size() + plan.pads() + max_new <= max_cells) {
                ram_reused = plan.n_prefix + plan.reused();
            }
        }
        if (ram_reused == 0) {
            for (const auto& s : ck_ring) {
                if (!s.ids.empty() && (int64_t) s.ids.size() <= n_pre && (int64_t) s.ids.size() > ram_reused &&
                    std::equal(s.ids.begin(), s.ids.end(), ids.begin())) {
                    ram_reused = (int64_t) s.ids.size();
                }
            }
        }

        // Check NVMe disk tier if it can beat the live in-RAM reuse by >= 32 tokens (text conversations only)
        if (!geni && disk_cache.enabled()) {
            int64_t disk_reused = 0;
            const int disk_idx = disk_cache.find_best(raw_ids, max_new, max_cells, kScrPadToken, disk_reused);
            if (disk_idx >= 0 && disk_reused >= ram_reused + 32) {
                if (ram_dirty && cached_ids.size() >= 32) {
                    disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path);
                    ram_dirty = false;
                }
                const int disk_idx2 = disk_cache.find_best(raw_ids, max_new, max_cells, kScrPadToken, disk_reused);
                if (disk_idx2 >= 0) {
                    const std::string load_path = disk_cache.entries[(size_t) disk_idx2].path;
                    const std::string load_parent = disk_cache.entries[(size_t) disk_idx2].parent_path;
                    const bool load_pinned = disk_cache.entries[(size_t) disk_idx2].pinned;
                    const std::vector<int64_t> load_pads = disk_cache.entries[(size_t) disk_idx2].pads;
                    std::vector<int32_t> i32;
                    int64_t cells = 0;
                    if (me::window_session_load(load_path.c_str(), i32, &cells, err, load_parent.c_str())) {
                        cached_ids.assign(i32.begin(), i32.end());
                        cached_pads = load_pads;
                        active_path = load_pinned ? "" : load_path;
                        ram_dirty = false;
                        ck_ring.clear();
                        save_ck_at(cached_ids);
                        disk_cache.touch((size_t) disk_idx2);
                        ids = raw_ids;
                        if (!cached_pads.empty() && !cached_ids.empty()) {
                            std::vector<int64_t> carried_ids;
                            cached_pads = strata::core::scr_carry(raw_ids, cached_ids, cached_pads, kScrPadToken, carried_ids);
                            if ((int64_t) carried_ids.size() + max_new <= max_cells) ids = std::move(carried_ids);
                            else cached_pads.clear();
                        }
                        n_pre = std::max<int64_t>(0, (int64_t) ids.size() - 1);
                        lcp = 0;
                        const size_t lim1 = std::min(cached_ids.size(), (size_t) n_pre);
                        while (lcp < lim1 && cached_ids[lcp] == ids[lcp]) ++lcp;
                    }
                }
            } else if (ram_reused == 0 && ram_dirty && cached_ids.size() >= 32) {
                disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path);
                ram_dirty = false;
            }
        }

        const int64_t learn_at = (!geni && disk_cache.enabled() && cached_pads.empty())
                                     ? disk_cache.observe_prefix(ids, n_pre)
                                     : -1;
        int64_t reused = 0;

        auto on_chunk = [&](const float*, int64_t k, int64_t p0, std::string& e) -> bool {
            if (!use_mtp) return true;
            std::vector<int32_t> nxt((size_t) k);
            for (int64_t t = 0; t < k; ++t) nxt[(size_t) t] = (int32_t) ids[(size_t) (p0 + t + 1)];
            return me::mtp_prefill(nxt.data(), k, p0, e);
        };
        const float* const* embd_ptr = geni ? row_ptr.data() : nullptr;
        auto run_prefill_span = [&](int64_t start_p, int64_t end_p) -> bool {
            if (end_p <= start_p) return true;
            if (learn_at > start_p && learn_at <= end_p && cached_pads.empty()) {
                if (!me::prompt_run(ids.data() + start_p, learn_at - start_p, start_p, embd_ptr, on_chunk, err)) return false;
                std::vector<int64_t> pfx(ids.begin(), ids.begin() + learn_at);
                save_ck_at(pfx);
                disk_cache.save_current(pfx, {}, /*pinned=*/true, err);
                if (end_p > learn_at) {
                    if (!me::prompt_run(ids.data() + learn_at, end_p - learn_at, learn_at, embd_ptr, on_chunk, err)) return false;
                }
                return true;
            }
            return me::prompt_run(ids.data() + start_p, end_p - start_p, start_p, embd_ptr, on_chunk, err);
        };

        if (lcp == cached_ids.size() && lcp > 0) {
            reused = (int64_t) lcp;
            if (reused > 0) {
                std::printf("RESUME %lld\n", (long long) reused);
                std::fflush(stdout);
            }
            if (n_pre > reused) {
                if (!run_prefill_span(reused, n_pre)) {
                    std::printf("ERR %s\n", err.c_str());
                    std::fflush(stdout);
                    cached_ids.clear();
                    ram_dirty = false;
                    continue;
                }
            }
            cached_ids.assign(ids.begin(), ids.begin() + n_pre);
            ram_dirty = true;
        } else {
            // 2. Check SCR mid-context fold [A B C] -> [A B' C D] (including stripped <think> blocks) via scr_plan
            if (!geni && cached_imgs.empty() && cached_ids.size() >= 32 && n_pre >= 32) {
                std::vector<int64_t> ck_lens = {(int64_t) cached_ids.size()};
                for (const auto& s : ck_ring) {
                    if (!s.ids.empty() && s.ids.size() <= cached_ids.size() &&
                        std::equal(s.ids.begin(), s.ids.end(), cached_ids.begin())) {
                        ck_lens.push_back((int64_t) s.ids.size());
                    }
                }
                const strata::core::ScrPlan plan = strata::core::scr_plan(
                    cached_ids, ids, ck_lens, strata::core::ScrParams{32, 64, 16, 4});
                if (!plan.empty() && (int64_t) ids.size() + plan.pads() + max_new <= max_cells) {
                    ids = strata::core::scr_padded(ids, plan, kScrPadToken);
                    n_pre = std::max<int64_t>(0, (int64_t) ids.size() - 1);
                    const auto new_pads = strata::core::scr_pad_positions(plan);
                    cached_pads.insert(cached_pads.end(), new_pads.begin(), new_pads.end());
                    std::sort(cached_pads.begin(), cached_pads.end());
                    cached_pads.erase(std::unique(cached_pads.begin(), cached_pads.end()), cached_pads.end());
                    save_ck_at(cached_ids);
                    bool scr_ok = true;
                    // Pass 1: relocate left-shifting runs (shift < 0) in forward order 0 .. m-1
                    for (size_t ri = 0; scr_ok && ri < plan.runs.size(); ++ri) {
                        const auto& r = plan.runs[ri];
                        if (r.shift() < 0 && !me::window_scr_relocate(r.old_start, r.at, r.n, err)) scr_ok = false;
                    }
                    // Pass 2: relocate right-shifting runs (shift > 0) in reverse order m-1 .. 0
                    for (size_t ri = plan.runs.size(); scr_ok && ri-- > 0;) {
                        const auto& r = plan.runs[ri];
                        if (r.shift() > 0 && !me::window_scr_relocate(r.old_start, r.at, r.n, err)) scr_ok = false;
                    }
                    int64_t cur_p = plan.restore > 0 ? plan.restore : (int64_t) lcp;
                    if (scr_ok && plan.restore > 0) {
                        for (const auto& s : ck_ring) {
                            if ((int64_t) s.ids.size() == plan.restore) {
                                me::window_checkpoint_restore(s.ck, err);
                                ss.ple_prev[0] = s.prev0;
                                ss.ple_prev[1] = s.prev1;
                                break;
                            }
                        }
                    }
                    // Pass 3: prefill edited gaps [cur_p, r.at) now that all runs sit at [r.at, r.at + r.n)
                    for (const auto& r : plan.runs) {
                        if (!scr_ok) break;
                        if (r.at > cur_p) {
                            if (!me::prompt_run_rows(ids.data() + cur_p, r.at - cur_p, cur_p, on_chunk, err)) {
                                scr_ok = false;
                                break;
                            }
                        }
                        cur_p = r.at + r.n;
                    }
                    if (scr_ok && plan.adopt > 0) {
                        for (const auto& s : ck_ring) {
                            if ((int64_t) s.ids.size() == plan.adopt) {
                                me::window_checkpoint_restore(s.ck, err);
                                ss.ple_prev[0] = s.prev0;
                                ss.ple_prev[1] = s.prev1;
                                break;
                            }
                        }
                    }
                    if (scr_ok && n_pre > cur_p) {
                        if (!me::prompt_run_rows(ids.data() + cur_p, n_pre - cur_p, cur_p, on_chunk, err)) {
                            scr_ok = false;
                        }
                    }
                    if (scr_ok) {
                        reused = plan.n_prefix + plan.reused();
                        std::printf("RESUME %lld\n", (long long) reused);
                        std::fflush(stdout);
                        cached_ids.assign(ids.begin(), ids.begin() + n_pre);
                        ram_dirty = true;
                    } else {
                        cached_ids.clear();
                        cached_pads.clear();
                        ram_dirty = false;
                    }
                }
            }
            if (reused == 0) {
                int best_idx = -1;
                size_t best_len = 0;
                for (size_t ci = 0; ci < ck_ring.size(); ++ci) {
                    const auto& cids = ck_ring[ci].ids;
                    if (!cids.empty() && cids.size() <= (size_t) n_pre && cids.size() > best_len &&
                        std::equal(cids.begin(), cids.end(), ids.begin())) {
                        best_len = cids.size();
                        best_idx = (int) ci;
                    }
                }
                if (best_idx >= 0) {
                    me::window_checkpoint_restore(ck_ring[(size_t) best_idx].ck, err);
                    ss.ple_prev[0] = ck_ring[(size_t) best_idx].prev0;
                    ss.ple_prev[1] = ck_ring[(size_t) best_idx].prev1;
                    reused = (int64_t) best_len;
                    std::printf("RESUME %lld\n", (long long) reused);
                    std::fflush(stdout);
                    if (n_pre > reused) {
                        if (!run_prefill_span(reused, n_pre)) {
                            std::printf("ERR %s\n", err.c_str());
                            std::fflush(stdout);
                            cached_ids.clear();
                            ram_dirty = false;
                            continue;
                        }
                    }
                    cached_ids.assign(ids.begin(), ids.begin() + n_pre);
                    ram_dirty = true;
                } else {
                    me::window_reset();
                    ss.ple_prev[0] = ss.ple_prev[1] = -1;
                    cached_ids.clear();
                    cached_pads.clear();
                    active_path.clear();
                    if (n_pre > 0) {
                        if (!run_prefill_span(0, n_pre)) {
                            std::printf("ERR %s\n", err.c_str());
                            std::fflush(stdout);
                            ram_dirty = false;
                            continue;
                        }
                    }
                    cached_ids.assign(ids.begin(), ids.begin() + n_pre);
                    ram_dirty = true;
                }
            }
        }

        const double prefill_ms = now_ms() - t_pre0;
        const double t_dec0 = now_ms();
        int32_t cur_tok = (int32_t) ids.back();
        int64_t pos = n_pre;
        int gen = 0;
        const char* finish = "length";

        if ((use_mtp || (suffix_draft > 0 && temp == 0.0f))) {
            strata::spec::SuffixDrafter suffix(4, 32, (size_t) max_cells);
            for (int64_t id : ids) suffix.append((int32_t) id);
            int32_t mtp_drafts[8] = {};
            float mtp_probs[8] = {};
            int n_mtp_drafts = 0;
            bool first_win = true;
            SpecControl spec_control;
            while (gen < max_new) {
                int32_t input[8] = {cur_tok}, output[8] = {};
                int nd = 0;
                const int rem = max_new - gen;
                nd = spec_control.propose(suffix, input, mtp_drafts, mtp_probs,
                                          !first_win && use_mtp ? n_mtp_drafts : 0,
                                          !first_win ? suffix_draft : 0, rem, pos, max_cells);
                first_win = false;
                const double verify_begin = now_ms();
                if (!me::window_run(w, 1 + nd, input, pos, true, output, err)) {
                    std::printf("ERR %s\n", err.c_str());
                    std::fflush(stdout);
                    cached_ids.clear();
                    finish = "error";
                    break;
                }
                spec_control.gate.spec_gate_forward(1 + nd, (now_ms() - verify_begin) / 1000);
                // the penalties see the window's tokens (cached_ids + input[0..keep-1]), as the one-token path's
                int keep;
                if (temp > 0.0f && spec_exact) {
                    keep = spec_sample(w, input, output, nd, n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, lg_buf, rng);
                } else {
                    keep = strata::spec::suffix_keep(input, output, nd);
                    cached_ids.insert(cached_ids.end(), input, input + keep);
                    if (temp > 0.0f && !strata::spec::suffix_stop(output[keep - 1])) {
                        me::window_copy_logits(w, keep - 1, lg_buf.data());
                        output[keep - 1] = sample_logits(lg_buf.data(), n_vocab, output[keep - 1], temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, rng);
                    }
                }
                spec_control.result(nd, keep);
                if (!me::window_commit(w, keep, err)) {
                    std::printf("ERR %s\n", err.c_str());
                    std::fflush(stdout);
                    cached_ids.clear();
                    finish = "error";
                    break;
                }
                for (int i = 0; i < keep; ++i) {
                    std::printf("T %d\n", (int) output[i]);
                    std::fflush(stdout);
                    suffix.append(output[i]);
                    cur_tok = output[i];
                    ++gen;
                    if (strata::spec::suffix_stop(cur_tok)) break;
                }
                if (!strata::spec::suffix_stop(cur_tok) && gen < max_new && use_mtp) {
                    const double draft_begin = now_ms();
                    if (!me::mtp_draft(1 + nd, output, pos, keep - 1, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) {
                        std::printf("ERR %s\n", err.c_str());
                        std::fflush(stdout);
                        cached_ids.clear();
                        finish = "error";
                        break;
                    }
                    spec_control.gate.spec_gate_draft_cost(0, (now_ms() - draft_begin) / (1000 * std::max(1, spec - 1)));
                } else if (use_mtp) {
                    me::mtp_prefill(output, keep, pos, err);
                }
                pos += keep;
                if (strata::spec::suffix_stop(cur_tok)) {
                    finish = "stop";
                    break;
                }
            }
            spec_control.describe();
        } else {
            for (int step = 0; step < max_new; ++step) {
                int32_t greedy_tok = 0;
                if (!me::window_run(w, 1, &cur_tok, pos, true, &greedy_tok, err) ||
                    !me::window_commit(w, 1, err)) {
                    std::printf("ERR %s\n", err.c_str());
                    std::fflush(stdout);
                    cached_ids.clear();
                    finish = "error";
                    break;
                }
                cached_ids.push_back(cur_tok);
                ++pos;

                int32_t next_tok = greedy_tok;
                if (temp > 0.0f) {
                    me::window_copy_logits(w, 0, lg_buf.data());
                    next_tok = sample_logits(lg_buf.data(), n_vocab, greedy_tok, temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, rng);
                }
                std::printf("T %d\n", (int) next_tok);
                std::fflush(stdout);
                ++gen;

                if (next_tok == 248046 || next_tok == 248044) {
                    finish = "stop";
                    break;
                }
                cur_tok = next_tok;
            }
        }
        cached_imgs = req_imgs;
        if (!geni) save_ck_at(cached_ids);
        const double decode_ms = now_ms() - t_dec0;
        if (std::strcmp(finish, "error") != 0) {
            std::printf("DONE %d %zu %.1f %.1f %s 0 0 %lld\n",
                        gen, ids.size(), prefill_ms, decode_ms, finish, (long long) reused);
            std::fflush(stdout);
            if (!geni && disk_cache.enabled() && cached_ids.size() >= 32) {
                disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path);
                ram_dirty = false;
            }
        }
    }
    return 0;
}

int main(int argc, char** argv) {
    bool serve_mode = false;
    std::string metal_gguf, ple_gguf, mtp_gguf, tokens_file, logits_out, logits_ref, logits_ref2;
    int n_expert = 256;
    int64_t max_cells = 4096;
    int n_prompt = 22;
    int n_gen = 64;
    int suffix_draft = 0;
    int spec = 1;
    float spec_min_p = 0.0f;
    std::string cache_disk_dir = std::getenv("STRATA_METAL_CACHE_DISK") ? std::getenv("STRATA_METAL_CACHE_DISK") : "";
    int64_t cache_disk_mib = std::getenv("STRATA_METAL_CACHE_DISK_MIB") ? std::atoll(std::getenv("STRATA_METAL_CACHE_DISK_MIB")) : 16384;
    std::string save_session, load_session, scr_a, scr_b, scr_bp, scr_c, scr_d, scr_j = "0,8,16", scr_tmp = "/tmp/strata_scr_test.session";
    int32_t scr_pad = 198;

    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        if (a == "--metal-gguf" && i + 1 < argc) metal_gguf = argv[++i];
        else if (a == "--ple-gguf" && i + 1 < argc) ple_gguf = argv[++i];
        else if (a == "--tokens-file" && i + 1 < argc) tokens_file = argv[++i];
        else if (a == "--n-expert" && i + 1 < argc) n_expert = std::atoi(argv[++i]);
        else if ((a == "--ctx" || a == "--max-context") && i + 1 < argc) max_cells = std::atoll(argv[++i]);
        else if (a == "--serve") serve_mode = true;
        else if (a == "--vision") { /* enabled automatically via GENI */ }
        else if (a == "--n-prompt" && i + 1 < argc) n_prompt = std::atoi(argv[++i]);
        else if (a == "--n-gen" && i + 1 < argc) n_gen = std::atoi(argv[++i]);
        else if (a == "--suffix-draft" && i + 1 < argc) suffix_draft = std::max(0, std::min(7, std::atoi(argv[++i])));
        else if (a == "--mtp-gguf" && i + 1 < argc) mtp_gguf = argv[++i];
        else if (a == "--spec" && i + 1 < argc) spec = std::max(1, std::min(4, std::atoi(argv[++i])));
        else if (a == "--spec-min-p" && i + 1 < argc) spec_min_p = (float) std::atof(argv[++i]);
        else if (a == "--cache-disk" && i + 1 < argc) cache_disk_dir = argv[++i];
        else if (a == "--cache-disk-size" && i + 1 < argc) cache_disk_mib = std::atoll(argv[++i]);
        else if (a == "--logits-out" && i + 1 < argc) logits_out = argv[++i];
        else if (a == "--logits-ref" && i + 1 < argc) logits_ref = argv[++i];
        else if (a == "--logits-ref2" && i + 1 < argc) logits_ref2 = argv[++i];
        else if (a == "--save-session" && i + 1 < argc) save_session = argv[++i];
        else if (a == "--load-session" && i + 1 < argc) load_session = argv[++i];
        else if (a == "--scr-a" && i + 1 < argc) scr_a = argv[++i];
        else if (a == "--scr-b" && i + 1 < argc) scr_b = argv[++i];
        else if (a == "--scr-bp" && i + 1 < argc) scr_bp = argv[++i];
        else if (a == "--scr-c" && i + 1 < argc) scr_c = argv[++i];
        else if (a == "--scr-d" && i + 1 < argc) scr_d = argv[++i];
        else if (a == "--scr-j" && i + 1 < argc) scr_j = argv[++i];
        else if (a == "--scr-pad" && i + 1 < argc) scr_pad = std::atoi(argv[++i]);
        else if (a == "--scr-tmp" && i + 1 < argc) scr_tmp = argv[++i];
    }
    if (metal_gguf.empty()) {
        std::fprintf(stderr, "usage: strata_window_cli --metal-gguf <K256-metal.gguf> [--ple-gguf <shard1.gguf>] "
                             "[--n-expert 256] [--ctx 4096] [--n-prompt 22] [--n-gen 64] [--tokens-file <ids.txt>] "
                             "[--logits-out <out.bin>] [--logits-ref <ref.bin>] [--save-session <f> | --load-session <f>] "
                             "[--scr-a A --scr-b B --scr-bp B' --scr-c C --scr-d D [--scr-j 0,8,16] [--scr-pad 198]]\n");
        return 2;
    }

    const bool tf_mode = !logits_out.empty() || !logits_ref.empty();
    std::vector<int64_t> all_ids;
    if (!tokens_file.empty()) {
        std::ifstream f(tokens_file);
        int64_t tid;
        while (f >> tid) all_ids.push_back(tid);
    }
    if (all_ids.empty()) {
        const int64_t def_ids[] = {
            248045, 8678, 198, 2523, 513, 264, 10631, 17313, 13, 248046, 198, 248045, 846, 198, 814, 20139, 1204, 27413, 71284, 3728, 13081, 32907, 12361, 321, 264, 3074, 6700, 11, 321, 948, 264, 28336, 628, 6561, 571, 13, 248046, 198, 248045, 74455, 198, 248068, 198
        };
        const size_t nd = sizeof(def_ids) / sizeof(def_ids[0]);
        for (int i = 0; i < n_prompt + (tf_mode ? n_gen : 0); ++i) all_ids.push_back(def_ids[i % nd]);
    }
    if (tf_mode && (int) all_ids.size() < n_prompt + n_gen) {
        std::fprintf(stderr, "strata_window_cli: teacher-forced mode needs at least n_prompt + n_gen (%d) tokens in --tokens-file, got %zu\n",
                     n_prompt + n_gen, all_ids.size());
        return 1;
    }
    std::vector<int64_t> prompt_ids(all_ids.begin(), all_ids.begin() + std::min<size_t>(all_ids.size(), (size_t) n_prompt));

    const double t_init0 = now_ms();
    std::string err;
    strata::kernels::PleTable ple_table;
    strata::core::SessionState ss{};
    ss.ple_prev[0] = -1;
    ss.ple_prev[1] = -1;
    ss.ple.consts = strata::kernels::ple_artifact_consts();
    if (!ple_gguf.empty()) {
        strata::kernels::PleIoOptions io{};
        io.mode = strata::kernels::PleIo::Direct;
        io.max_inflight = 256;
        io.cache_rows = 1u << 18;
        if (!ple_table.open(ple_gguf, err, io)) {
            std::fprintf(stderr, "strata_window_cli: PleTable::open failed: %s\n", err.c_str());
            return 1;
        }
        ss.ple.table = &ple_table;
    }

    strata::core::ModelGeometry g{};
    g.n_embd = 2560;
    g.hc = 4;
    g.hc_lr = 320;
    g.n_ff = 640;
    g.n_expert = n_expert;
    g.n_layers = 48;

    strata::metal_engine::set_model_shard(metal_gguf.c_str());
    const bool use_mtp = !mtp_gguf.empty() && spec > 1;
    const int max_win_t = std::max({1, 1 + suffix_draft, spec});
    if (use_mtp) {
        setenv("STRATA_METAL_MTP_GGUF", mtp_gguf.c_str(), 1);
        if (!strata::metal_engine::mtp_load(mtp_gguf, max_win_t, err)) {
            std::fprintf(stderr, "strata_window_cli: mtp_load failed: %s\n", err.c_str());
            return 1;
        }
    }
    void* w = strata::metal_engine::window_create();
    int64_t n_vocab = 0;
    if (!strata::metal_engine::window_init(w, g, ss, max_win_t, max_cells, &n_vocab, err)) {
        std::fprintf(stderr, "strata_window_cli: window_init failed: %s\n", err.c_str());
        return 1;
    }
    const double init_ms = now_ms() - t_init0;
    if (serve_mode) return serve_loop(w, ss, n_vocab, max_cells, suffix_draft, use_mtp, spec, spec_min_p, cache_disk_dir, cache_disk_mib);
    if (!scr_a.empty()) return scr_test(w, ss, n_vocab, scr_a, scr_b, scr_bp, scr_c, scr_d, scr_j, scr_pad, n_gen, scr_tmp);

    FILE* fout_lg = nullptr;
    if (!logits_out.empty()) {
        fout_lg = std::fopen(logits_out.c_str(), "wb");
        if (!fout_lg) {
            std::perror("fopen logits_out");
            return 1;
        }
        fcntl(fileno(fout_lg), F_NOCACHE, 1);
    }
    FILE* fref_lg = nullptr;
    if (!logits_ref.empty()) {
        fref_lg = std::fopen(logits_ref.c_str(), "rb");
        if (!fref_lg) {
            std::perror("fopen logits_ref");
            return 1;
        }
        fcntl(fileno(fref_lg), F_NOCACHE, 1);
    }
    FILE* fref2_lg = nullptr;
    if (!logits_ref2.empty()) {
        fref2_lg = std::fopen(logits_ref2.c_str(), "rb");
        if (!fref2_lg) {
            std::perror("fopen logits_ref2");
            return 1;
        }
        fcntl(fileno(fref2_lg), F_NOCACHE, 1);
    }

    strata::spec::SuffixDrafter suffix(4, 32, (size_t) max_cells);
    for (int64_t tid : prompt_ids) suffix.append((int32_t) tid);

    // Prefill all but the last prompt token via prompt_run_rows
    const int64_t n_pre = std::max<int64_t>(0, (int64_t) prompt_ids.size() - 1);
    const double t_pre0 = now_ms();
    if (!load_session.empty()) {
        // the prefix from disk instead of a prefill: it must be this prompt's first n_pre tokens
        std::vector<int32_t> ids;
        int64_t cells = 0;
        if (!strata::metal_engine::window_session_load(load_session.c_str(), ids, &cells, err)) {
            std::fprintf(stderr, "strata_window_cli: %s\n", err.c_str());
            return 1;
        }
        bool same = (std::getenv("STRATA_BENCH_EXPAND_CELLS") != nullptr) || (cells == n_pre && (int64_t) ids.size() == n_pre);
        if (!std::getenv("STRATA_BENCH_EXPAND_CELLS")) {
            for (int64_t i = 0; same && i < n_pre; ++i) same = ids[(size_t) i] == (int32_t) prompt_ids[(size_t) i];
        }
        if (!same) {
            std::fprintf(stderr, "strata_window_cli: the session holds %lld cells of another prefix (want this prompt's first %lld)\n",
                         (long long) cells, (long long) n_pre);
            return 1;
        }
        std::printf("STRATA_SESSION load_ms=%.2f cells=%lld bytes=%zu\n", now_ms() - t_pre0, (long long) cells,
                    scr::fixed_bytes() + scr::cell_bytes((size_t) cells));
    } else if (n_pre > 0) {
        auto on_chunk = [&](const float*, int64_t k, int64_t p0, std::string& e) -> bool {
            if (!use_mtp) return true;
            std::vector<int32_t> nxt((size_t) k);
            for (int64_t t = 0; t < k; ++t) nxt[(size_t) t] = (int32_t) prompt_ids[(size_t) (p0 + t + 1)];
            return strata::metal_engine::mtp_prefill(nxt.data(), k, p0, e);
        };
        if (!strata::metal_engine::prompt_run_rows(prompt_ids.data(), n_pre, 0, on_chunk, err)) {
            std::fprintf(stderr, "strata_window_cli: prompt_run_rows failed: %s\n", err.c_str());
            return 1;
        }
    }
    if (!save_session.empty()) {
        std::vector<int32_t> ids(prompt_ids.begin(), prompt_ids.begin() + n_pre);
        const double t_s = now_ms();
        if (!strata::metal_engine::window_session_save(save_session.c_str(), ids.data(), n_pre, err)) {
            std::fprintf(stderr, "strata_window_cli: %s\n", err.c_str());
            return 1;
        }
        std::printf("STRATA_SESSION save_ms=%.2f cells=%lld\n", now_ms() - t_s, (long long) n_pre);
    }
    // Run the last prompt token with head=true to get the first output distribution
    int32_t cur_tok = (int32_t) prompt_ids.back();
    int32_t next_tok = 0;
    int64_t pos = std::getenv("STRATA_BENCH_EXPAND_CELLS") ? std::atoll(std::getenv("STRATA_BENCH_EXPAND_CELLS")) : n_pre;
    if (!strata::metal_engine::window_run(w, 1, &cur_tok, pos, true, &next_tok, err) ||
        !strata::metal_engine::window_commit(w, 1, err)) {
        std::fprintf(stderr, "strata_window_cli: first window_run failed: %s\n", err.c_str());
        return 1;
    }
    const double prefill_ms = now_ms() - t_pre0;
    int32_t mtp_drafts[8] = {};
    float mtp_probs[8] = {};
    int n_mtp_drafts = 0;
    double mtp_ms = 0.0;
    if (use_mtp && !tf_mode && n_gen > 1) {
        const double tm0 = now_ms();
        if (!strata::metal_engine::mtp_draft(1, &next_tok, pos, 0, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) {
            std::fprintf(stderr, "strata_window_cli: mtp_draft(0) failed: %s\n", err.c_str());
            return 1;
        }
        mtp_ms += now_ms() - tm0;
    }
    ++pos;

    std::vector<float> cur_lg(tf_mode ? (size_t) n_vocab : 0);
    std::vector<float> ref_lg((fref_lg || fref2_lg) ? (size_t) n_vocab : 0);
    double sum_nll_cand = 0.0, sum_nll_ref = 0.0, sum_nll_ref2 = 0.0;
    std::vector<double> klds, klds2;
    int top1_matches = 0, top1_matches2 = 0;

    auto eval_tf_step = [&](int step_idx) -> bool {
        if (!strata::metal_engine::window_copy_logits(w, 0, cur_lg.data())) return false;
        if (fout_lg) {
            if (std::fwrite(cur_lg.data(), sizeof(float), (size_t) n_vocab, fout_lg) != (size_t) n_vocab) return false;
        }
        const int32_t tgt = (int32_t) all_ids[(size_t) n_prompt + (size_t) step_idx];
        int32_t cand_top1 = 0;
        const double lse_cand = logsumexp_f64(cur_lg.data(), n_vocab, &cand_top1);
        sum_nll_cand += -((double) cur_lg[(size_t) tgt] - lse_cand);

        if (fref_lg) {
            if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref_lg) != (size_t) n_vocab) return false;
            int32_t ref_top1 = 0;
            const double lse_ref = logsumexp_f64(ref_lg.data(), n_vocab, &ref_top1);
            sum_nll_ref += -((double) ref_lg[(size_t) tgt] - lse_ref);
            if (cand_top1 == ref_top1) ++top1_matches;
            double kld = 0.0;
            for (int64_t v = 0; v < n_vocab; ++v) {
                const double log_p = (double) ref_lg[(size_t) v] - lse_ref;
                const double log_q = (double) cur_lg[(size_t) v] - lse_cand;
                const double p = std::exp(log_p);
                if (p > 0.0) kld += p * (log_p - log_q);
            }
            klds.push_back(std::max(0.0, kld));
        }
        if (fref2_lg) {
            if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref2_lg) != (size_t) n_vocab) return false;
            int32_t ref2_top1 = 0;
            const double lse_ref2 = logsumexp_f64(ref_lg.data(), n_vocab, &ref2_top1);
            sum_nll_ref2 += -((double) ref_lg[(size_t) tgt] - lse_ref2);
            if (cand_top1 == ref2_top1) ++top1_matches2;
            double kld2 = 0.0;
            for (int64_t v = 0; v < n_vocab; ++v) {
                const double log_p = (double) ref_lg[(size_t) v] - lse_ref2;
                const double log_q = (double) cur_lg[(size_t) v] - lse_cand;
                const double p = std::exp(log_p);
                if (p > 0.0) kld2 += p * (log_p - log_q);
            }
            klds2.push_back(std::max(0.0, kld2));
        }
        return true;
    };

    if (tf_mode && n_gen > 0) {
        if (!eval_tf_step(0)) {
            std::fprintf(stderr, "strata_window_cli: eval_tf_step(0) failed\n");
            return 1;
        }
    }

    std::vector<int32_t> gen_tokens;
    gen_tokens.push_back(next_tok);
    suffix.append(next_tok);
    cur_tok = tf_mode ? (int32_t) all_ids[(size_t) n_prompt] : next_tok;

    strata::metal_engine::window_reset_stats(w);
    const int decode_steps = std::max(0, n_gen - 1);
    const double t_dec0 = now_ms();
    int produced = 0, spec_rounds = 0, spec_drafted = 0, spec_accepted = 0;
    SpecControl spec_control;
    while (produced < decode_steps) {
        int32_t win_in[8] = {cur_tok};
        int32_t win_out[8] = {};
        int nd = 0;
        const int rem = decode_steps - produced;
        nd = spec_control.propose(suffix, win_in, mtp_drafts, mtp_probs,
                                  !tf_mode && use_mtp ? n_mtp_drafts : 0,
                                  !tf_mode ? suffix_draft : 0, rem, pos, max_cells);
        const int T = 1 + nd;
        const double verify_begin = now_ms();
        if (!strata::metal_engine::window_run(w, T, win_in, pos, true, win_out, err)) {
            std::fprintf(stderr, "strata_window_cli: decode step failed: %s\n", err.c_str());
            return 1;
        }
        spec_control.gate.spec_gate_forward(T, (now_ms() - verify_begin) / 1000);
        const int n_keep = strata::spec::suffix_keep(win_in, win_out, nd, false);
        spec_control.result(nd, n_keep);
        if (!strata::metal_engine::window_commit(w, n_keep, err)) {
            std::fprintf(stderr, "strata_window_cli: window_commit failed: %s\n", err.c_str());
            return 1;
        }
        ++spec_rounds;
        spec_drafted += nd;
        spec_accepted += (n_keep - 1);
        if (tf_mode) {
            if (!eval_tf_step(produced + 1)) {
                std::fprintf(stderr, "strata_window_cli: eval_tf_step(%d) failed\n", produced + 1);
                return 1;
            }
            gen_tokens.push_back(win_out[0]);
            ++produced;
            if (produced < decode_steps) {
                cur_tok = (int32_t) all_ids[(size_t) n_prompt + (size_t) produced];
            }
            pos += 1;
        } else {
            for (int i = 0; i < n_keep && produced < decode_steps; ++i) {
                gen_tokens.push_back(win_out[i]);
                suffix.append(win_out[i]);
                cur_tok = win_out[i];
                ++produced;
            }
            if (use_mtp && produced < decode_steps) {
                const double tm0 = now_ms();
                if (!strata::metal_engine::mtp_draft(T, win_out, pos, n_keep - 1, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) {
                    std::fprintf(stderr, "strata_window_cli: mtp_draft failed: %s\n", err.c_str());
                    return 1;
                }
                const double draft_ms = now_ms() - tm0;
                mtp_ms += draft_ms;
                spec_control.gate.spec_gate_draft_cost(0, draft_ms / (1000 * std::max(1, spec - 1)));
            }
            pos += n_keep;
        }
    }
    if (!tf_mode && (use_mtp || suffix_draft > 0)) spec_control.describe();
    const double decode_ms = now_ms() - t_dec0;
    const double decode_tok_s = decode_steps > 0 ? (decode_steps * 1000.0 / decode_ms) : 0.0;
    const double decode_ms_tok = decode_steps > 0 ? (decode_ms / decode_steps) : 0.0;
    const double prefill_tok_s = prompt_ids.size() * 1000.0 / prefill_ms;

    if (fout_lg) std::fclose(fout_lg);
    if (fref_lg) std::fclose(fref_lg);
    if (fref2_lg) std::fclose(fref2_lg);

    std::printf("STRATA_RESULT init_ms=%.2f prompt_tokens=%zu prefill_ms=%.2f prefill_tok_s=%.2f "
                "decode_tokens=%d decode_ms=%.2f decode_tok_s=%.2f decode_ms_per_tok=%.2f\n",
                init_ms, prompt_ids.size(), prefill_ms, prefill_tok_s,
                decode_steps, decode_ms, decode_tok_s, decode_ms_tok);
    if (spec_rounds > 0 && (use_mtp || suffix_draft > 0)) {
        std::printf("STRATA_SPEC rounds=%d drafted=%d accepted=%d accept_pct=%.1f tok_per_round=%.2f mtp_ms=%.2f\n",
                    spec_rounds, spec_drafted, spec_accepted,
                    spec_drafted > 0 ? 100.0 * (double) spec_accepted / (double) spec_drafted : 0.0,
                    (double) decode_steps / (double) spec_rounds, mtp_ms);
    }
    if (tf_mode && n_gen > 0) {
        const double cand_ppl = std::exp(sum_nll_cand / n_gen);
        if (!klds.empty()) {
            const double ref_ppl = std::exp(sum_nll_ref / n_gen);
            double mean_kld = 0.0;
            for (double x : klds) mean_kld += x;
            mean_kld /= klds.size();
            std::vector<double> sk = klds;
            std::sort(sk.begin(), sk.end());
            const double med_kld = sk[sk.size() / 2];
            const double p90_kld = sk[(size_t) (0.90 * (sk.size() - 1))];
            const double top1_pct = 100.0 * (double) top1_matches / (double) n_gen;
            std::printf("STRATA_KLD eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n",
                        n_gen, ref_ppl, cand_ppl, cand_ppl - ref_ppl, mean_kld, med_kld, p90_kld, top1_pct);
            if (!klds2.empty()) {
                const double ref2_ppl = std::exp(sum_nll_ref2 / n_gen);
                double mean_kld2 = 0.0;
                for (double x : klds2) mean_kld2 += x;
                mean_kld2 /= klds2.size();
                std::vector<double> sk2 = klds2;
                std::sort(sk2.begin(), sk2.end());
                const double med_kld2 = sk2[sk2.size() / 2];
                const double p90_kld2 = sk2[(size_t) (0.90 * (sk2.size() - 1))];
                const double top1_pct2 = 100.0 * (double) top1_matches2 / (double) n_gen;
                std::printf("STRATA_KLD_REF2 eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n",
                            n_gen, ref2_ppl, cand_ppl, cand_ppl - ref2_ppl, mean_kld2, med_kld2, p90_kld2, top1_pct2);
            }
        } else {
            std::printf("STRATA_PPL eval_tokens=%d cand_ppl=%.4f\n", n_gen, cand_ppl);
        }
    }
    std::printf("GEN_IDS:");
    for (size_t i = 0; i < gen_tokens.size(); ++i) {
        std::printf(" %d", gen_tokens[i]);
    }
    std::printf("\n");
    std::fflush(stdout);

    strata::metal_engine::window_destroy(w);
    return 0;
}