#pragma once

#include "kernels.hpp"

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <mutex>
#include <random>
#include <shared_mutex>
#include <string>
#include <thread>
#include <tuple>
#include <utility>
#include <vector>

namespace cism {

constexpr std::size_t max_weight_bytes = std::size_t{64} * 1024 * 1024 * 1024;
constexpr std::size_t max_kv_bytes = std::size_t{8} * 1024 * 1024 * 1024;
std::size_t checked_mul(std::size_t a, std::size_t b, const char* label);
std::size_t checked_add(std::size_t a, std::size_t b, const char* label);

struct Config {
    std::string model_type;
    std::size_t hidden, intermediate, layers, heads, kv_heads, head_dim, vocab, context;
    float eps, rope_theta;
    bool tied;
    void validate() const;
};

struct SurjoConfig {
    std::string model_type = "surjo";
    std::size_t hidden = 0, intermediate = 0, layers = 0;
    std::size_t heads = 0, kv_heads = 0, head_dim = 0, vocab = 0, context = 0;
    std::size_t prelude = 0, recurrent = 0, coda = 0, groups = 0, passes = 0, per_xsa = 0;
    std::size_t gdn_v_heads = 0, gdn_k_dim = 0, gdn_v_dim = 0, conv_kernel = 4;
    bool allow_neg = false, xsa_projection = true, tied = false;
    float eps = 1e-5f, rope_theta = 10000.0f;
    void validate() const;
    std::size_t num_xsa_slots() const { return prelude + groups * passes + coda; }
    bool is_xsa_layer(std::size_t li) const;
};

using WeightMap = std::map<std::string, std::vector<float>>;
using ShapeMap = std::map<std::string, std::vector<std::size_t>>;
ShapeMap weight_shapes(const Config& config, bool has_head);
ShapeMap surjo_weight_shapes(const SurjoConfig& config, bool has_head);

enum class Storage { fp32, fp16, int8, int4, fp4 };

// Persistent worker threads for row-parallel matvecs. The caller thread is
// always worker zero, so a single-threaded model never spawns anything.
// Barriers spin briefly, then sleep, keeping per-matrix handoff cost in the
// hundreds of nanoseconds while the pool idles at zero CPU when unused.
class WorkerPool {
public:
    // threads = total participants including the caller; 1 means no workers.
    explicit WorkerPool(std::size_t threads);
    ~WorkerPool();
    WorkerPool(const WorkerPool&) = delete;
    WorkerPool& operator=(const WorkerPool&) = delete;
    std::size_t thread_count() const { return workers_ + 1; }
    // task(begin, end) executes that participant's half-open row slice;
    // WorkerPool computes identical slices for the caller (worker zero) and
    // every background worker from the shared row count. const because Models
    // own their pool behind shared_ptr<const Model> in live sessions.
    void run(std::function<void(std::size_t, std::size_t)> task, std::size_t rows) const;
private:
    void worker_loop(std::size_t index);
    std::size_t workers_;
    std::vector<std::thread> threads_;
    mutable std::mutex run_mutex_;
    mutable std::function<void(std::size_t, std::size_t)> task_;
    mutable std::atomic<std::size_t> rows_{0};
    mutable std::atomic<std::uint64_t> generation_{0};
    // Cumulative completions; the target for generation G is G * workers_.
    mutable std::atomic<std::size_t> completed_{0};
    std::atomic<bool> stopping_{false};
};

class Matrix {
public:
    Matrix() = default;
    Matrix(std::vector<float> values, std::size_t rows, std::size_t cols, Storage storage);
    // Row-slice executor shared by gemm() and pair fusion: runs rows
    // [begin, end) with the exact default-path arithmetic (same kernels and
    // order as the historical gemm bodies). Thread-safe on disjoint slices
    // (reads X/weights/perm_ptr, writes disjoint out rows).
    // Raw pointer to the first weight byte (any storage): lets callers
    // prefetch the next layer's matrices across layer boundaries (exact-
    // neutral hint; never dereferenced here).
    const void* stream_head() const;
    void run_rows(const float* X, float* out, std::size_t tokens,
                  std::size_t begin, std::size_t end, const float* perm_ptr) const;
    void multiply(const std::vector<float>& input, std::vector<float>& output,
                  const WorkerPool* pool = nullptr, std::vector<float>* scratch = nullptr) const;
    // Multi-token evaluation: X is tokens x cols (token-major), out is
    // tokens x rows. Weight rows stream once for all tokens, which amortizes
    // DRAM traffic for oversized models during speculative verification.
    void gemm(const float* X, float* out, std::size_t tokens, const WorkerPool* pool,
              std::vector<float>* scratch = nullptr) const;
    void row(std::size_t index, std::vector<float>& output) const;
    std::size_t bytes() const;
    using Regions = std::vector<std::pair<void*, std::size_t>>;
    void collect(Regions& regions) const;
    // Quantized-activation dispatch (int8 activations, pmaddwd kernels).
    // Only takes effect on quantized storage; fp32 matrices ignore it.
    void set_act_q8(bool enabled) { act_q8_ = enabled; }
private:
    std::size_t rows_ = 0, cols_ = 0;
    Storage storage_ = Storage::fp32;
    bool act_q8_ = false;
    std::vector<float> floats_, scales_;
    // FP16 (IEEE-754 binary16) weight storage: half the bytes of fp32,
    // widened to fp32 on load in the dot kernels. No scales.
    std::vector<std::uint16_t> fp16_;
    std::vector<std::int8_t> int8_;
    std::vector<std::uint8_t> int4_;
    // FP4 (E2M1) nibble packing plus one E4M3 scale byte per 16 elements.
    std::vector<std::uint8_t> fp4_, fp4_scales_;
};

struct Layer {
    std::vector<float> input_norm, post_norm, q_norm, k_norm;
    Matrix q, k, v, o, gate, up, down;
};

class Session;
class PagedSession;
class BatchSession;

// ---- Native paged KV BlockPool (PagedAttention-style, dense Llama/Qwen3) ----
// Physical KV storage is two dense pools:
//   k_pages / v_pages : [num_blocks][num_layers][block_size][kv_width] fp32.
// free_stack is LIFO (allocate pops, free pushes); table maps req_id ->
// resident phys block ids in allocation order. Block size is fixed at 16
// tokens to match the Python paging simulator (src/cism/paging.py).
//
// Stream S2 (cross-request reuse / COW prefix sharing / eviction):
// - Blocks are refcounted. Plain allocate/ensure/free never share, so the
//   historical LIFO exact behavior is unchanged (refcount stays 1).
// - fork(dst, src, prefix_len) shares src's whole leading prefix blocks
//   without copying (refcount++); a non-aligned tail is copied once into a
//   fresh private block. Writes to shared blocks must go through
//   copy_on_write() first, which privatizes the block (allocate + full-block
//   memcpy). Shared values are never mutated, so paged math stays
//   bitwise-identical to dense.
// - Eviction is recompute-not-swap on CPU: evict_one() frees the
//   newest-lowest-priority request (lowest priority value; ties broken by
//   newest sequence number) and the victim recomputes from scratch on resume.
//   No KV is ever swapped to host.
// - Cross-request reuse lives one level up: Model owns a SharedPagedState
//   (one BlockPool per model shape, grown once, never per session), so steady
//   state sessions allocate zero KV pages on the hot path.
constexpr std::size_t kPagedBlockSize = 16;
// Default cap for the shared per-model pool (env CISM_PAGED_MAX_BLOCKS
// overrides; clamped to the 8 GiB KV safety limit for the model shape).
constexpr std::size_t kPagedSharedBlocksDefault = 256;

class BlockPool {
public:
    BlockPool() = default;
    BlockPool(std::size_t num_blocks, std::size_t num_layers, std::size_t kv_width,
              std::size_t block_size = kPagedBlockSize);
    BlockPool(const BlockPool&) = delete;
    BlockPool& operator=(const BlockPool&) = delete;
    BlockPool(BlockPool&&) noexcept = default;
    BlockPool& operator=(BlockPool&&) noexcept = default;

    std::size_t num_blocks() const { return num_blocks_; }
    std::size_t num_layers() const { return num_layers_; }
    std::size_t kv_width() const { return kv_width_; }
    std::size_t block_size() const { return block_size_; }

    std::size_t blocks_needed(std::size_t seq_len) const;
    bool can_allocate(std::size_t seq_len) const;

    // Allocate ceil(seq_len/block_size) blocks for req_id. Throws
    // std::invalid_argument on duplicate req_id, std::runtime_error when out
    // of free blocks (no partial allocation). Returns a copy of the resident
    // block list.
    std::vector<std::size_t> allocate(std::int64_t req_id, std::size_t seq_len);
    // Grow req_id to cover new_len tokens (no shrink; no-op when covered).
    // Throws std::invalid_argument on unknown req_id, std::runtime_error on
    // OOM (existing prefix is preserved).
    std::vector<std::size_t> ensure(std::int64_t req_id, std::size_t new_len);
    // Release all blocks for req_id. Throws std::invalid_argument on unknown
    // req_id (covers double-free). Returns the freed block list.
    // Refcount-aware: shared phys blocks are only returned to free_stack
    // once their last referencing request frees them.
    std::vector<std::size_t> free(std::int64_t req_id);

    // Share src's leading prefix_len tokens with a new req dst (no data copy
    // for whole blocks; the partial tail block, if any, is copied once into
    // a fresh private block). Requires prompt[:prefix_len] to be exactly
    // equal upstream (checked by the session fork path, not here). Throws
    // std::invalid_argument on duplicate dst / unknown src / prefix beyond
    // src's allocation, std::runtime_error on OOM (state unchanged).
    // Returns a copy of dst's resident block list.
    std::vector<std::size_t> fork(std::int64_t dst_req, std::int64_t src_req,
                                  std::size_t prefix_len);
    // Make the block holding position pos of req_id private: when its
    // refcount is 1 this is a no-op, otherwise a fresh block is allocated,
    // the full block contents (all layers/slots, K and V) are copied, and the
    // table entry is swapped. Must be called before writing position pos.
    // Throws std::invalid_argument on unknown req_id / out-of-range pos,
    // std::runtime_error on OOM (existing mapping preserved).
    void copy_on_write(std::int64_t req_id, std::size_t pos);
    // Evict one request under pressure: the newest request among those with
    // the lowest priority value (ties -> largest sequence number), excluding
    // except_req. Returns the victim id, or -1 when nothing is evictable.
    // Recompute-not-swap: the victim's blocks are released; it must restart
    // from scratch (no KV is preserved, nothing is swapped to host).
    std::int64_t evict_one(std::int64_t except_req = -1);
    void set_priority(std::int64_t req_id, int priority);
    int priority(std::int64_t req_id) const;
    // Grow the pool by extra_blocks physical blocks (LIFO-appended; new pages
    // zero-filled). Never shrinks. Throws on 8 GiB cap overflow.
    void grow(std::size_t extra_blocks);

    bool contains(std::int64_t req_id) const;
    const std::vector<std::size_t>& get_blocks(std::int64_t req_id) const;
    const std::vector<std::size_t>& free_stack() const { return free_stack_; }
    const std::map<std::int64_t, std::vector<std::size_t>>& table() const { return table_; }
    // Refcount of a physical block (0 == free). Throws on out-of-range phys.
    std::size_t refcount(std::size_t phys) const;

    std::size_t used() const;
    // Physically occupied blocks (accounts for sharing; used() is logical and
    // double-counts shared blocks by design to match the Python simulator).
    std::size_t phys_used() const { return num_blocks_ - free_stack_.size(); }
    std::size_t free_count() const { return free_stack_.size(); }
    std::size_t total() const { return num_blocks_; }
    std::tuple<std::size_t, std::size_t, std::size_t> usage() const;
    std::size_t kv_bytes() const;
    std::size_t kv_elements() const;

    const std::vector<float>& k_pages() const { return k_pages_; }
    const std::vector<float>& v_pages() const { return v_pages_; }

    // Raw gather helpers: pointer to [kv_width] floats for
    // (phys block, layer, slot-in-block). Caller guarantees bounds.
    float* k_ptr(std::size_t phys, std::size_t layer, std::size_t slot);
    const float* k_ptr(std::size_t phys, std::size_t layer, std::size_t slot) const;
    float* v_ptr(std::size_t phys, std::size_t layer, std::size_t slot);
    const float* v_ptr(std::size_t phys, std::size_t layer, std::size_t slot) const;
    // Fast unchecked row access for the decode hot loop (same indices as
    // k_ptr/v_ptr; bounds guaranteed by allocate/ensure + validated shapes).
    // Plain arithmetic: bitwise-identical addressing without per-access
    // overflow checks.
    const float* k_row_fast(std::size_t phys, std::size_t layer, std::size_t slot) const {
        return k_pages_.data() + ((phys * num_layers_ + layer) * block_size_ + slot) * kv_width_;
    }
    float* k_row_fast(std::size_t phys, std::size_t layer, std::size_t slot) {
        return k_pages_.data() + ((phys * num_layers_ + layer) * block_size_ + slot) * kv_width_;
    }
    const float* v_row_fast(std::size_t phys, std::size_t layer, std::size_t slot) const {
        return v_pages_.data() + ((phys * num_layers_ + layer) * block_size_ + slot) * kv_width_;
    }
    float* v_row_fast(std::size_t phys, std::size_t layer, std::size_t slot) {
        return v_pages_.data() + ((phys * num_layers_ + layer) * block_size_ + slot) * kv_width_;
    }
    // Base pointer to slot 0 of (phys, layer): the hot loop advances by
    // slot * kv_width and reloads the base on block boundaries, avoiding a
    // division/modulo per inner-loop step.
    const float* k_base_fast(std::size_t phys, std::size_t layer) const {
        return k_pages_.data() + ((phys * num_layers_ + layer) * block_size_) * kv_width_;
    }
    const float* v_base_fast(std::size_t phys, std::size_t layer) const {
        return v_pages_.data() + ((phys * num_layers_ + layer) * block_size_) * kv_width_;
    }
    float* k_base_fast(std::size_t phys, std::size_t layer) {
        return k_pages_.data() + ((phys * num_layers_ + layer) * block_size_) * kv_width_;
    }
    float* v_base_fast(std::size_t phys, std::size_t layer) {
        return v_pages_.data() + ((phys * num_layers_ + layer) * block_size_) * kv_width_;
    }

private:
    std::size_t offset_of(std::size_t phys, std::size_t layer, std::size_t slot) const;
    std::size_t pop_free();
    void release_table_entry(std::int64_t req_id);
    std::size_t num_blocks_ = 0, num_layers_ = 0, kv_width_ = 0;
    std::size_t block_size_ = kPagedBlockSize;
    std::vector<float> k_pages_, v_pages_;
    std::vector<std::size_t> free_stack_;
    std::map<std::int64_t, std::vector<std::size_t>> table_;
    // S2: per-phys refcount (0 == free), per-req priority (default 0), and a
    // monotonic admission sequence for newest-lowest-priority eviction.
    std::vector<std::size_t> refcounts_;
    std::map<std::int64_t, int> prio_;
    std::map<std::int64_t, std::uint64_t> seq_;
    std::uint64_t seq_ctr_ = 0;
};

// ---- Shared cross-request paged KV store (Stream S2) ----
// One BlockPool per model shape, owned by the Model and borrowed by every
// PagedSession of that model. Pages are allocated once (grown never shrunk);
// steady-state sessions only insert/erase small table entries, so there is
// zero KV-page allocation on the hot path. All BlockPool calls must hold mu
// (PagedSession holds it exclusively for the whole forward, matching the
// WorkerPool serialization the engine already imposes).
struct SharedPagedState {
    mutable std::shared_mutex mu;
    BlockPool pool;
    bool init = false;
    std::uint64_t req_ctr = 0;
    std::size_t max_blocks = 0;
};

class Model : public std::enable_shared_from_this<Model> {
public:
    Model(Config config, WeightMap weights, std::string precision, std::size_t threads = 1,
          const std::string& act_precision = "fp32");
    // "fp32" keeps exact FP32 activations; "int8" quantizes activations per
    // 32-element block (32767/absmax) for the pmaddwd kernel path.
    void set_act_precision(const std::string& act_precision);
    std::shared_ptr<Session> create_session(std::vector<std::int64_t> prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, std::vector<std::int64_t> eos);
    std::shared_ptr<PagedSession> create_paged_session(std::vector<std::int64_t> prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, std::vector<std::int64_t> eos);
    // Fork a paged session sharing src's leading prefix_len tokens (COW, no
    // copy of whole blocks). Requires prompt[:prefix_len] ==
    // src.history[:prefix_len] exactly; throws std::invalid_argument
    // otherwise (correctness first: shared KV is only valid for identical
    // prefix tokens at identical positions). The child prefills only the
    // suffix; shared blocks stay read-only for the child (COW-guarded).
    std::shared_ptr<PagedSession> create_paged_fork_session(
        std::vector<std::int64_t> prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, std::vector<std::int64_t> eos,
        std::shared_ptr<PagedSession> src, std::size_t prefix_len,
        int priority = 0);
    std::shared_ptr<BatchSession> create_batch_session(
        std::vector<std::vector<std::int64_t>> prompts,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, std::vector<std::int64_t> eos,
        // Optional per-sequence RNG seeds (serving: each job keeps its own
        // serial seed, so batched sampled decode is token-identical to
        // serial). Empty = legacy seed+i. Non-empty must match prompts.
        const std::vector<std::uint64_t>& seq_seeds = {});
    std::vector<float> logits(const std::vector<std::int64_t>& prompt);
    std::vector<float> paged_logits(const std::vector<std::int64_t>& prompt);
    // Teacher-forced mean NLL over tokens[1..] (truncated-window scoring:
    // each window is scored with positions starting at zero).
    double nll(const std::vector<std::int64_t>& tokens);
    // Shared paged-KV pool observability (logical used, free, total) plus
    // physically occupied blocks. Idle-safe test/ops hook: paged_pool_reset
    // clears all table entries (live sessions transparently recompute from
    // history on next use); max_blocks==0 keeps the current cap.
    std::tuple<std::size_t, std::size_t, std::size_t> paged_pool_usage() const;
    std::size_t paged_pool_phys_used() const;
    std::size_t paged_pool_max_blocks() const;
    void paged_pool_reset(std::size_t max_blocks = 0) const;
    std::size_t thread_count() const { return pool_.thread_count(); }
    // Weight memory control. x86 cannot pin CPU cache lines; lock_pages pins
    // DRAM pages against OS paging, and touch() refreshes cache residency.
    std::size_t lock_pages();
    std::size_t unlock_pages();
    std::size_t touch();
    // Fast non-volatile full-weight scan for bandwidth probes; touch() is a
    // deliberately serialized residency refresher, not a bandwidth benchmark.
    std::size_t scan();
    bool pages_locked() const { return pages_locked_.load(std::memory_order_relaxed); }
    const Config& config() const { return config_; }
    const std::string& precision() const { return precision_; }
    const std::string& compile_key() const { return compile_key_; }
    const std::string& compile_canonical() const { return compile_canonical_; }
    const char* act_precision() const { return act_q8_ ? "int8" : "fp32"; }
    std::size_t weight_bytes() const { return weight_bytes_; }
private:
    friend class Session;
    friend class PagedSession;
    friend class BatchSession;
    Config config_;
    std::string precision_;
    std::string compile_key_;
    std::string compile_canonical_;
    std::string act_precision_ = "fp32";
    bool act_q8_ = false;
    Matrix embeddings_, head_;
    bool tied_head_;
    std::vector<Layer> layers_;
    std::vector<float> norm_, inv_freq_;
    // Rotary sin/cos tables [context][head_dim/2], built once at load
    // (read-only afterwards; sessions share them freely).
    std::vector<float> rope_cos_, rope_sin_;
    WorkerPool pool_;
    std::size_t weight_bytes_ = 0;
    std::atomic<bool> pages_locked_{false};
    mutable std::uint64_t warm_sink_ = 0;
    mutable std::uint64_t scan_sink_ = 0;
    // S2 shared paged-KV store (lazy init on first paged session).
    mutable std::shared_ptr<SharedPagedState> paged_state_;
    mutable std::mutex paged_init_mu_;
    std::shared_ptr<SharedPagedState> paged_state() const;
    void collect_regions(Matrix::Regions& regions) const;
    void refresh_compile_key();
};

class Session {
public:
    Session(std::shared_ptr<const Model> model, const std::vector<std::int64_t>& prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, const std::vector<std::int64_t>& eos);
    std::vector<std::int64_t> next_tokens(std::int64_t count, std::int64_t spec_k = 0);
    // Greedy speculative verification of explicit candidates: accepts the
    // argmax-matching prefix, appends the correction/bonus token, forwards it.
    // Exact for temperature=0 (enforced by the SDK layer).
    std::vector<std::int64_t> verify(const std::vector<std::int64_t>& candidates);
    void cancel() noexcept { cancelled_.store(true, std::memory_order_relaxed); }
    std::string finish_reason();
    std::size_t generated_tokens();
    // Lifetime speculative-decoding tally: drafted candidates proposed by
    // draft() versus accepted by verify(). Zero when spec is off.
    std::pair<std::size_t, std::size_t> spec_stats() const;
    // Prompt-lookup drafter over the full token history (for tests/debugging).
    std::vector<std::int64_t> draft(std::size_t k) const;
private:
    friend class Model;
    bool prefill();
    bool forward(std::int64_t token);
    bool forward_tokens(const std::vector<std::int64_t>& tokens);
    std::int64_t sample();
    std::shared_ptr<const Model> model_;
    std::atomic<bool> cancelled_{false};
    mutable std::mutex mutex_;
    std::string finish_;
    std::size_t capacity_, kv_elements_, position_ = 0, generated_ = 0, max_new_;
    double temperature_, top_p_;
    std::size_t top_k_;
    std::mt19937_64 rng_;
    std::vector<std::int64_t> eos_, prompt_, history_;
    std::vector<float> keys_, values_, logits_, block_logits_;
    // One caller at a time; all scratch and KV storage belongs to the session.
    std::vector<float> x_, normalized_, q_, k_, v_, attention_, projected_, gate_, up_, scores_;
    // Disjoint per-(token, head) score slices for pooled attention (one row
    // per (t,h); see Session::forward_tokens). Persistent to avoid allocation.
    std::vector<float> attn_scores_;
    // Session-owned persistent scratch for int4/fp4 permuted activations (one
    // session = one generation at a time under the Session mutex, so reuse is
    // race-free; pool workers only read it after caller-thread sizing).
    std::vector<float> perm_scratch_;
    // Session-owned persistent staging for the embedding gather (sequential per
    // token under the Session mutex, so reuse is race-free).
    std::vector<float> row_buf_;
    // Guarded by mutex_ like all other session state.
    std::size_t spec_proposed_ = 0, spec_accepted_ = 0;
};

// ---- Paged decode (blocked SDPA gather over the shared BlockPool) ----
// Same arithmetic as Session::forward_tokens in the same kernel order, but KV
// lives in the model's shared BlockPool (block_size=16, S2) instead of
// contiguous [layers][capacity][kv_width] buffers. Bitwise-identical logits
// to the dense path for fp32. The dense Session API is unchanged; this is an
// additive path.
//
// Sharing model: every session holds one req_id in the shared pool. The pool
// lock is held exclusively for the whole forward (sessions already serialize
// on the engine generation lock and the WorkerPool barrier), so page
// pointers can never dangle across a concurrent grow. Evicted sessions
// (or a pool reset) drop their table entry; the next forward transparently
// recomputes from history_ (recompute-not-swap: no KV is preserved, nothing
// is swapped to host).
class PagedSession : public std::enable_shared_from_this<PagedSession> {
public:
    PagedSession(std::shared_ptr<const Model> model, const std::vector<std::int64_t>& prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, const std::vector<std::int64_t>& eos);
    ~PagedSession();
    std::vector<std::int64_t> next_tokens(std::int64_t count);
    void cancel() noexcept { cancelled_.store(true, std::memory_order_relaxed); }
    std::string finish_reason();
    std::size_t generated_tokens();
    std::size_t position() const;
    std::size_t capacity() const;
    const std::vector<float>& last_logits() const;
    // Eviction priority (default 0; higher survives). Used only by
    // BlockPool::evict_one under pool pressure.
    void set_priority(int priority);
    int priority() const;
    // Lifetime recompute count (evictions survived via history recompute).
    std::size_t recomputes() const;
    // Run the prompt prefill without generating (pre-warms history_/logits_,
    // e.g. to prepare a fork source for an empty-suffix child). Safe to call
    // before next_tokens; a second call is a no-op.
    bool prefill_fork_source();
private:
    friend class Model;
    // Fork constructor: prompt_suffix enumerates the non-shared tail,
    // prefix_history the already-materialized shared prefix (position is set
    // to prefix_history.size(); prefill forwards only the suffix). An empty
    // suffix skips prefill entirely, so seed_logits must carry the source's
    // last forward output (the distribution for the next token); it is
    // assigned to logits_ so the first sample is exact, never garbage.
    PagedSession(std::shared_ptr<const Model> model, std::vector<std::int64_t> prompt_suffix,
        std::vector<std::int64_t> prefix_history,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, const std::vector<std::int64_t>& eos,
        std::shared_ptr<PagedSession> src, std::size_t prefix_len, int priority,
        std::vector<float> seed_logits);
    void init_common();
    bool prefill();
    bool forward(std::int64_t token);
    bool forward_tokens(const std::vector<std::int64_t>& tokens);
    std::int64_t sample();
    std::shared_ptr<const Model> model_;
    std::atomic<bool> cancelled_{false};
    mutable std::mutex mutex_;
    std::string finish_;
    std::size_t capacity_, position_ = 0, generated_ = 0, max_new_;
    double temperature_, top_p_;
    std::size_t top_k_;
    std::mt19937_64 rng_;
    std::vector<std::int64_t> eos_, prompt_, history_;
    std::vector<float> logits_, block_logits_;
    // Paged KV: borrowed shared pool (cross-request reuse) + unique req_id.
    std::shared_ptr<SharedPagedState> state_;
    std::int64_t req_id_ = -1;
    int priority_ = 0;
    std::size_t recomputes_ = 0;
    // One caller at a time; all scratch belongs to the session.
    std::vector<float> x_, normalized_, q_, k_, v_, attention_, projected_, gate_, up_, scores_;
    // Disjoint per-(token, head) score slices for pooled attention (one row
    // per (t,h), mirroring Session::forward_tokens). Persistent to avoid
    // allocation; pool pages are pinned by the pool lock during fan-out.
    std::vector<float> attn_scores_;
    std::vector<float> perm_scratch_;
    std::vector<float> row_buf_;
};

// ---- Batch section (dense Llama/Qwen3 only; Surjo later) ----
// True batched decode: one weight stream, B sequences. Per-seq prompt /
// position / generated / rng / eos / finish; concatenated X:[sumT,hidden]
// per-layer projections via a single gemm; per-seq attention/sampling loops.
// For sumT==1 the kernels run in the same order as Session::forward_tokens,
// giving bitwise-identical results to the single-sequence path.
class BatchSession {
public:
    BatchSession(std::shared_ptr<const Model> model,
        std::vector<std::vector<std::int64_t>> prompts,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, const std::vector<std::int64_t>& eos,
        const std::vector<std::uint64_t>& seq_seeds = {});
    // Generate up to `count` tokens per sequence. Returns B vectors; entry i
    // holds the tokens emitted for sequence i (possibly shorter on stop).
    std::vector<std::vector<std::int64_t>> next_tokens(std::int64_t count);
    void cancel() noexcept { cancelled_.store(true, std::memory_order_relaxed); }
    // Mark one sequence finished (takes effect on the next step boundary).
    void cancel_seq(std::size_t index);
    std::vector<std::string> finish_reasons();
    std::vector<std::size_t> generated_tokens_list();
    std::size_t batch_size() const;
private:
    friend class Model;
    struct SeqState {
        std::vector<std::int64_t> prompt_;
        std::vector<std::int64_t> history_;
        std::vector<float> keys_, values_, logits_;
        std::size_t capacity_ = 0, kv_elements_ = 0, position_ = 0, generated_ = 0, max_new_ = 0;
        std::mt19937_64 rng_;
        std::vector<std::int64_t> eos_;
        std::string finish_;
    };
    bool prefill();
    bool forward_batch(const std::vector<std::vector<std::int64_t>>& batch_tokens);
    std::int64_t sample_seq(SeqState& seq);
    std::shared_ptr<const Model> model_;
    std::atomic<bool> cancelled_{false};
    mutable std::mutex mutex_;
    std::vector<SeqState> seqs_;
    double temperature_ = 0, top_p_ = 1;
    std::size_t top_k_ = 0;
    bool prefilled_ = false;
    // Shared concatenated scratch: token-major [sumT, ...].
    std::vector<float> x_, normalized_, q_, k_, v_, attention_, projected_, gate_, up_, scores_;
    std::vector<float> block_logits_;
    std::vector<float> perm_scratch_, row_buf_;
};

struct SurjoXsaWeights {
    Matrix q, k, v, o;
    std::vector<float> q_norm, k_norm;
};

struct SurjoGdnWeights {
    Matrix q, k, v, b, w, f0, f1, g0, g1, o_proj;
    std::vector<float> q_conv, k_conv, v_conv;  // flat (dim*kernel) FP32
    std::vector<float> a_log, dt_bias, g_bias, o_norm;
    std::size_t conv_kernel = 4;
};

struct SurjoLayer {
    bool is_xsa = true;
    std::vector<float> input_norm, post_norm;
    SurjoXsaWeights xsa;
    SurjoGdnWeights gdn;
    Matrix gate, up, down;
};

class SurjoSession;

class SurjoModel : public std::enable_shared_from_this<SurjoModel> {
public:
    SurjoModel(SurjoConfig config, WeightMap weights, std::string precision,
               std::size_t threads = 1, const std::string& act_precision = "fp32");
    void set_act_precision(const std::string& act_precision);
    std::shared_ptr<SurjoSession> create_session(std::vector<std::int64_t> prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, std::vector<std::int64_t> eos);
    std::vector<float> logits(const std::vector<std::int64_t>& prompt);
    double nll(const std::vector<std::int64_t>& tokens);
    std::size_t thread_count() const { return pool_.thread_count(); }
    std::size_t lock_pages();
    std::size_t unlock_pages();
    std::size_t touch();
    std::size_t scan();
    bool pages_locked() const { return pages_locked_.load(std::memory_order_relaxed); }
    const SurjoConfig& config() const { return config_; }
    const std::string& precision() const { return precision_; }
    std::size_t weight_bytes() const { return weight_bytes_; }
private:
    friend class SurjoSession;
    SurjoConfig config_;
    std::string precision_;
    std::string act_precision_ = "fp32";
    bool act_q8_ = false;
    Matrix embeddings_, head_;
    bool tied_head_ = false;
    std::vector<SurjoLayer> layers_;
    std::vector<float> norm_, inv_freq_;
    // (layer_idx, pass, slot): slot>=0 XSA KV slot, -1 GDN recurrent.
    struct PlanStep { std::size_t layer, pass; int slot; };
    std::vector<PlanStep> plan_;
    std::size_t xsa_slots_ = 0, gdn_states_ = 0;
    WorkerPool pool_;
    std::size_t weight_bytes_ = 0;
    std::atomic<bool> pages_locked_{false};
    mutable std::uint64_t warm_sink_ = 0, scan_sink_ = 0;
    void collect_regions(Matrix::Regions& regions) const;
    void build_plan();
};

class SurjoSession {
public:
    SurjoSession(std::shared_ptr<const SurjoModel> model, const std::vector<std::int64_t>& prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, const std::vector<std::int64_t>& eos);
    std::vector<std::int64_t> next_tokens(std::int64_t count, std::int64_t spec_k = 0);
    // Greedy speculative verification of explicit candidates: accepts the
    // argmax-matching prefix, appends the correction/bonus token, forwards it.
    // Exact for temperature=0 (enforced by the SDK layer). GDN recurrence is
    // destructive, so the pre-block state (GDN S + conv FIFOs + XSA tail +
    // position) is checkpointed before the block forward and restored (memcpy)
    // with accepted-prefix replay on partial reject.
    std::vector<std::int64_t> verify(const std::vector<std::int64_t>& candidates);
    void cancel() noexcept { cancelled_.store(true, std::memory_order_relaxed); }
    std::string finish_reason();
    std::size_t generated_tokens();
    // Lifetime speculative-decoding tally: drafted candidates proposed by
    // draft() versus accepted by verify(). Zero when spec is off.
    std::pair<std::size_t, std::size_t> spec_stats() const;
    // Prompt-lookup drafter over the full token history (for tests/debugging).
    std::vector<std::int64_t> draft(std::size_t k) const;
    std::vector<std::int64_t> neural_draft(std::size_t k);
    void set_draft_steps(const std::vector<std::size_t>& skipped);
private:
    friend class SurjoModel;
    std::vector<bool> draft_skipped_;
    bool drafting_ = false;
    bool prefill();
    bool forward(std::int64_t token);
    bool forward_tokens(const std::vector<std::int64_t>& tokens);
    // Block-batched forward for T>1 (speculative verify / NLL scoring).
    // Per-layer projections stream weights once for all T tokens via a single
    // gemm over concatenated X:[T,hidden]; XSA attention and GDN recurrence /
    // FIFO stay sequential per token in token order, so results are
    // bitwise-identical to T sequential forward() calls.
    bool forward_block(const std::vector<std::int64_t>& tokens);
    std::int64_t sample();
    std::shared_ptr<const SurjoModel> model_;
    std::atomic<bool> cancelled_{false};
    mutable std::mutex mutex_;
    std::string finish_;
    std::size_t capacity_, position_ = 0, generated_ = 0, max_new_;
    double temperature_, top_p_;
    std::size_t top_k_;
    std::mt19937_64 rng_;
    std::vector<std::int64_t> eos_, prompt_, history_;
    std::size_t xsa_slots_total_ = 0, gdn_total_ = 0, gdn_v_total_ = 0;
    // XSA KV: [slots, capacity, kv_width]; GDN S: [gdn_states, H*K*V]; conv FIFO.
    std::vector<float> xsa_keys_, xsa_values_, gdn_state_;
    std::vector<float> gdn_q_fifo_, gdn_k_fifo_, gdn_v_fifo_;
    std::vector<float> logits_, block_logits_;
    std::vector<float> x_, normalized_, q_, k_, v_, attention_, projected_, gate_, up_, scores_;
    std::vector<float> g_, b_, w_, gv_, o_normed_;
    std::vector<float> perm_scratch_, row_buf_;
    // Guarded by mutex_ like all other session state.
    std::size_t spec_proposed_ = 0, spec_accepted_ = 0;
};

struct FwkvConfig {
    std::string model_type = "fwkv";
    std::size_t d_model = 0, d_emb = 0, layers = 0, ffn_mult = 4, vocab = 0, context = 1024;
    float wkv_floor = 0.1f;
    bool tied = true;
    void validate() const;
};

ShapeMap fwkv_weight_shapes(const FwkvConfig& config);

struct FwkvLayer {
    Matrix proj_k, proj_v, proj_r, proj_out, ffn0, ffn2;
    std::vector<float> w, norm_wkv_w, norm_wkv_b, norm_ffn_w, norm_ffn_b;
};

class FwkvSession;

class FwkvModel : public std::enable_shared_from_this<FwkvModel> {
public:
    FwkvModel(FwkvConfig config, WeightMap weights, std::string precision,
              std::size_t threads = 1, const std::string& act_precision = "fp32");
    void set_act_precision(const std::string& act_precision);
    std::shared_ptr<FwkvSession> create_session(std::vector<std::int64_t> prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, std::vector<std::int64_t> eos);
    std::vector<float> logits(const std::vector<std::int64_t>& prompt);
    double nll(const std::vector<std::int64_t>& tokens);
    std::size_t thread_count() const { return pool_.thread_count(); }
    std::size_t lock_pages();
    std::size_t unlock_pages();
    std::size_t touch();
    std::size_t scan();
    bool pages_locked() const { return pages_locked_.load(std::memory_order_relaxed); }
    const FwkvConfig& config() const { return config_; }
    const std::string& precision() const { return precision_; }
    std::size_t weight_bytes() const { return weight_bytes_; }
private:
    friend class FwkvSession;
    FwkvConfig config_;
    std::string precision_;
    std::string act_precision_ = "fp32";
    bool act_q8_ = false;
    // Factorized tied head: shared.weight [vocab, d_emb] as a quantized
    // Matrix (embed gather via row(), logits via multiply), shared.proj
    // [d_model, d_emb] kept FP32 (147K mults per token either direction).
    Matrix shared_emb_, shared_head_;
    std::vector<float> shared_proj_;
    std::vector<FwkvLayer> layers_;
    std::vector<float> norm_w_, norm_b_;
    WorkerPool pool_;
    std::size_t weight_bytes_ = 0;
    std::atomic<bool> pages_locked_{false};
    mutable std::uint64_t warm_sink_ = 0, scan_sink_ = 0;
    void collect_regions(Matrix::Regions& regions) const;
};

class FwkvSession {
public:
    FwkvSession(std::shared_ptr<const FwkvModel> model, const std::vector<std::int64_t>& prompt,
        std::int64_t max_new_tokens, double temperature, double top_p,
        std::int64_t top_k, std::uint64_t seed, const std::vector<std::int64_t>& eos);
    std::vector<std::int64_t> next_tokens(std::int64_t count, std::int64_t spec_k = 0);
    // Greedy speculative verification of explicit candidates: accepts the
    // argmax-matching prefix, appends the correction/bonus token, forwards it.
    // Exact for temperature=0 (enforced by the SDK layer). The WKV recurrent
    // state is destructive, so it is checkpointed (layers*d_model floats +
    // position, ~40KB for Myosotis-1-base) before the block forward and
    // restored (memcpy) with accepted-prefix replay on partial reject.
    std::vector<std::int64_t> verify(const std::vector<std::int64_t>& candidates);
    void cancel() noexcept { cancelled_.store(true, std::memory_order_relaxed); }
    std::string finish_reason();
    std::size_t generated_tokens();
    // Lifetime speculative-decoding tally: drafted candidates proposed by
    // draft() versus accepted by verify(). Zero when spec is off.
    std::pair<std::size_t, std::size_t> spec_stats() const;
    // Prompt-lookup drafter over the full token history (for tests/debugging).
    std::vector<std::int64_t> draft(std::size_t k) const;
private:
    friend class FwkvModel;
    bool prefill();
    bool forward(std::int64_t token);
    bool forward_tokens(const std::vector<std::int64_t>& tokens);
    // Block-batched forward for T>1 (speculative verify / NLL scoring).
    // Per-layer projections stream weights once for all T tokens via a single
    // gemm over concatenated X:[T,d_model]; WKV recurrence, LayerNorms and
    // the factorized head stay sequential per token in token order, so results
    // are bitwise-identical to T sequential forward() calls.
    bool forward_block(const std::vector<std::int64_t>& tokens);
    std::int64_t sample();
    std::shared_ptr<const FwkvModel> model_;
    std::atomic<bool> cancelled_{false};
    mutable std::mutex mutex_;
    std::string finish_;
    std::size_t capacity_, position_ = 0, generated_ = 0, max_new_;
    double temperature_, top_p_;
    std::size_t top_k_;
    std::mt19937_64 rng_;
    std::vector<std::int64_t> eos_, prompt_, history_;
    // WKV recurrent state: [layers, d_model], zeros at start.
    std::vector<float> state_;
    std::vector<float> logits_, block_logits_;
    std::vector<float> x_, h_, k_, v_, r_, g_, o_, ffn_h_, emb_, x_emb_;
    std::vector<float> perm_scratch_, row_buf_;
    // Guarded by mutex_ like all other session state.
    std::size_t spec_proposed_ = 0, spec_accepted_ = 0;
};
}  // namespace cism