// src/kernels/ngram.cpp - P2.S4: the PLE n-gram hash and the IQ4_NL table read.
//
// See include/strata/kernels/ngram.hpp for the semantics, the rival readings, and the note on MADV_RANDOM.
#include "strata/kernels/ngram.hpp"
#include "strata/artifact/gguf_reader.hpp"
#include "strata/artifact/dequant.hpp"
#include "strata/kernels/f16_bits.hpp"
#include "strata/ngram/ple_reader.hpp"
#include "strata/platform/memory.hpp"

#include <cerrno>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <limits>
#if !defined(_WIN32)
#include <sys/mman.h>
#endif
#include <algorithm>
#include <atomic>
#include <thread>
#include <vector>
#include <stdexcept>

#if defined(_WIN32)
// `PrefetchVirtualMemory` (memoryapi.h, Windows 8+) is the whole point of the change in `gather` below.
#define WIN32_LEAN_AND_MEAN
#define NOMINMAX
#include <windows.h>
#endif

namespace strata::kernels {

namespace {
/// The A/B arm.  Host-token-path only, so it needs no atomics; see the note on `ple_prefetch_enable`.
bool g_ple_prefetch = true;
}  // namespace

void ple_prefetch_enable(bool on) { g_ple_prefetch = on; }
bool ple_prefetch_enabled() { return g_ple_prefetch; }

PleConsts ple_artifact_consts() {
    // docs/gguf-dump-shard1.txt, verbatim.  Written once here rather than derived, because `head_offsets` is
    // ALSO in the metadata as its own array - and deriving one from the other would hide a mismatch between
    // them instead of surfacing it.  The parity test checks the two agree.
    PleConsts c{};
    c.mult[0] = 23703573157769ull;
    c.mult[1] = 20109073645365ull;
    c.mult[2] = 8052911324071ull;
    const uint64_t vocab[PLE_N_HEADS] = {
        20000003, 20000023, 20000033, 20000047, 20000059, 20000063, 20000069, 20000077,
        20000081, 20000093, 20000107, 20000147, 20000153, 20000159, 20000161, 20000171};
    const uint64_t offset[PLE_N_HEADS] = {
        0,        20000003, 40000026, 60000059, 80000106, 100000165, 120000228, 140000297,
        160000374, 180000455, 200000548, 220000655, 240000802, 260000955, 280001114, 300001275};
    for (int i = 0; i < PLE_N_HEADS; ++i) {
        c.vocab[i] = vocab[i];
        c.offset[i] = offset[i];
    }
    return c;
}

uint64_t ngram_mixed(const int64_t* ctx, const uint64_t* mult, int n) {
    // The first term is an ASSIGNMENT and the rest are XORed into it, which is how the source writes it
    // (`uint64_t mixed = ctx[0]*m[0]; for j=1.. mixed ^= ctx[j]*m[j];`).  Every product wraps mod 2^64,
    // which is what the `(uint64_t)` casts in the source make explicit.
    uint64_t mixed = (uint64_t) ctx[0] * mult[0];
    for (int j = 1; j < n; ++j) mixed ^= (uint64_t) ctx[j] * mult[j];
    return mixed;
}

void ngram_rows(const int32_t* tokens, const int32_t* prev, int n_tokens, const PleConsts& c, uint32_t* out) {
    const int n_prev = NGRAM_SIZE - 1;
    for (int i = 0; i < n_tokens; ++i) {
        int64_t ctx[NGRAM_SIZE];
        ctx[0] = tokens[i];
        bool cut = false;
        for (int s = 1; s < NGRAM_SIZE; ++s) {
            // `prev` is OLDEST FIRST, so predecessor `s` positions back is entry (n_prev - s): s=1 reads the
            // NEWEST.  Reading index (s-1) instead walks the window backwards, which still produces indices
            // in range and so cannot be caught by a range check - only by an oracle.
            const int32_t t = cut ? TOKEN_NULL : prev[i * n_prev + (n_prev - s)];
            // The cut is evaluated BEFORE the value is stored, so the position whose predecessor was EOS is
            // itself EOS.  Storing first and then cutting would leave position s holding the real token while
            // position s+1 became EOS - one token of history too much.
            cut = cut || t < 0 || t == PLE_EOS_TOKEN_ID;
            ctx[s] = cut ? PLE_EOS_TOKEN_ID : t;
        }
        for (int n = 2; n <= NGRAM_SIZE; ++n) {
            const uint64_t mixed = ngram_mixed(ctx, c.mult, n);
            const int base = (n - 2) * HEADS_PER_NGRAM;
            for (int g = 0; g < HEADS_PER_NGRAM; ++g) {
                const int h = base + g;
                out[i * PLE_N_HEADS + h] = (uint32_t) (mixed % c.vocab[h] + c.offset[h]);
            }
        }
    }
}

namespace {
const int8_t kIq4Nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113};
}

int iq4nl_code(int code) { return kIq4Nl[code & 15]; }

void iq4nl_dequant_row(const uint8_t* row, float* out160) {
    for (int b = 0; b < PLE_HEAD_DIM / 32; ++b) {
        const uint8_t* blk = row + (size_t) b * 18;
        uint16_t dbits;
        std::memcpy(&dbits, blk, 2);
        const float d = f32_from_f16(dbits);
        const uint8_t* qs = blk + 2;
        // SPLIT HALVES: qs[j] holds elements j and j+16, not 2j and 2j+1.
        for (int j = 0; j < 16; ++j) {
            out160[b * 32 + j] = d * (float) kIq4Nl[qs[j] & 0x0F];
            out160[b * 32 + j + 16] = d * (float) kIq4Nl[qs[j] >> 4];
        }
    }
}

namespace {
struct Fp8Table {
    float v[256];
    Fp8Table() {
        for (int x = 0; x < 256; ++x) {
            const int s = x >> 7, e = (x >> 3) & 15, m = x & 7;
            float f;
            if (e == 15 && m == 7) f = 0.0f;                     // NaN in e4m3fn; never in a real table, read as 0
            else if (e == 0) f = std::ldexp((float) m / 8.0f, -6);
            else f = std::ldexp(1.0f + (float) m / 8.0f, e - 7);
            v[x] = s ? -f : f;
        }
    }
};
const Fp8Table kFp8;
}  // namespace

void fp8_e4m3_dequant_row(const uint8_t* row, float scale, float* out160) {
    for (int j = 0; j < PLE_HEAD_DIM; ++j) out160[j] = kFp8.v[row[j]] * scale;
}

void bf16_dequant_row(const uint8_t* row, float* out160) {
    // bfloat16 IS the top 16 bits of a float32, so widening is exact for every value, normal or not - no rounding,
    // no clamp, no scale. Bytes are assembled little-endian explicitly rather than reinterpreted, so this is the
    // same on a big-endian host.
    for (int j = 0; j < PLE_HEAD_DIM; ++j) {
        const uint32_t bits = (uint32_t) row[2 * j] | ((uint32_t) row[2 * j + 1] << 8);
        const uint32_t wide = bits << 16;
        float v;
        std::memcpy(&v, &wide, sizeof v);
        out160[j] = v;
    }
}

// ---------------------------------------------------------------------------------------------------
// THE FORMAT TABLE (see PleFormatInfo). One entry per table type; the block formats decode 5 blocks of 32.
namespace {
template <void (*Block)(const uint8_t*, float*), int BlockBytes>
void dequant_blocks(const uint8_t* row, float /*scale*/, float* out160) {
    for (int b = 0; b < PLE_HEAD_DIM / 32; ++b) Block(row + (size_t) b * BlockBytes, out160 + b * 32);
}
void dequant_iq4_nl(const uint8_t* row, float, float* out160) { iq4nl_dequant_row(row, out160); }
void dequant_fp8(const uint8_t* row, float scale, float* out160) { fp8_e4m3_dequant_row(row, scale, out160); }
void dequant_bf16(const uint8_t* row, float, float* out160) { bf16_dequant_row(row, out160); }

const PleFormatInfo kPleFormats[] = {
    {PleFormat::IQ4_NL, "IQ4_NL", "IQ4_NL", PLE_ROW_BYTES, false, dequant_iq4_nl},
    {PleFormat::Q5_0, "Q5_0", "Q5_0", (PLE_HEAD_DIM / 32) * 22, false, dequant_blocks<strata::dequantize_q5_0, 22>},
    {PleFormat::F8_E4M3, "F8_E4M3", "I8", PLE_ROW_BYTES_FP8, true, dequant_fp8},
    // Ordinary GGUFs ship the table in other types: Unsloth's UD-Q6_K_XL and Swift-1.5 Q4_K_L as Q8_0, a Q5_K_M finetune as Q5_1
    {PleFormat::Q5_1, "Q5_1", "Q5_1", PLE_ROW_BYTES_Q5_1, false, dequant_blocks<strata::dequantize_q5_1, 24>},
    {PleFormat::Q8_0, "Q8_0", "Q8_0", PLE_ROW_BYTES_Q8_0, false, dequant_blocks<strata::dequantize_q8_0, 34>},
    // Q4_0 (plain llama-quantize Q4_0 files, #599): 90-byte rows like IQ4_NL (18-byte blocks), a linear 4-bit grid
    {PleFormat::Q4_0, "Q4_0", "Q4_0", PLE_ROW_BYTES, false, dequant_blocks<strata::dequantize_q4_0, 18>},
    // BF16 (#586): the checkpoint's own table at full precision, 320-byte rows; self-describing, so no scale and no
    // metadata to trust (tools/ple_fp8_pack.py writes the FP8 form of the same table)
    {PleFormat::BF16, "BF16", "BF16", PLE_ROW_BYTES_BF16, false, dequant_bf16},
};
constexpr int kPleFormatCount = (int) (sizeof kPleFormats / sizeof kPleFormats[0]);
}  // namespace

const PleFormatInfo* ple_formats() { return kPleFormats; }
int ple_format_count() { return kPleFormatCount; }
const PleFormatInfo& ple_format_info(PleFormat f) { return kPleFormats[(int) f]; }
const PleFormatInfo* ple_format_for_type(const char* gguf_type_name) {
    for (const PleFormatInfo& f : kPleFormats)
        if (std::strcmp(f.gguf_type, gguf_type_name) == 0) return &f;
    return nullptr;
}
std::string ple_format_list() {
    std::string out;
    for (int i = 0; i < kPleFormatCount; ++i) {
        const PleFormatInfo& f = kPleFormats[i];
        if (i > 0) out += i + 1 == kPleFormatCount ? " or " : ", ";
        out += std::strcmp(f.gguf_type, "I8") == 0 ? "FP8 (I8)" : f.name;
    }
    return out;
}

static_assert(PLE_ROW_BYTES_MAX >= PLE_ROW_BYTES && PLE_ROW_BYTES_MAX >= PLE_ROW_BYTES_FP8, "row buffers too small");

struct PleTable::Impl {
    GgufFile* file = nullptr;
    const uint8_t* data = nullptr;
    uint64_t n_rows = 0;
    const PleFormatInfo* fmt = &ple_format_info(PleFormat::IQ4_NL);   // the table's format (#296: Q5_0 is one of them)
    mutable uint64_t bytes_read = 0;
    // Direct mode (plan v0.3 P2): the mapping above is released after the header parse and every row comes
    // from an unbuffered SSD read into `raw`.
    PleIo mode = PleIo::Mmap;
    strata::ngram::PleReader reader;
    strata::ngram::PleReader::Ticket ticket;
    bool pending = false;
    bool locked = false;
    uint32_t rows[PLE_N_HEADS] = {};
    uint8_t raw[PLE_N_HEADS * PLE_ROW_BYTES_MAX] = {};
    static constexpr size_t kMaxPrefetch = 16;
    size_t n_prefetch = 0;
    strata::ngram::PleReader::Ticket prefetch_tickets[kMaxPrefetch] = {};
    uint32_t prefetch_keys[kMaxPrefetch][PLE_N_HEADS] = {};
    uint8_t prefetch_raw[kMaxPrefetch][PLE_N_HEADS * PLE_ROW_BYTES_MAX] = {};
    float scale = 1.0f;               // the FP8 table's one scale
    uint32_t rb = PLE_ROW_BYTES;      // bytes per row
    void decode(const uint8_t* row, float* out160) const { fmt->dequant(row, scale, out160); }
};

PleTable::PleTable() : impl_(new Impl) {}
PleTable::~PleTable() { close(); delete impl_; }

bool PleTable::open(const std::string& gguf_path, std::string& err) {
    return open(gguf_path, err, PleIoOptions{});
}

bool PleTable::open(const std::string& gguf_path, std::string& err, const PleIoOptions& io) {
    close();
    try {
        impl_->file = new GgufFile(gguf_path);
    } catch (const std::exception& e) {
        err = e.what();
        return false;
    }
    const TensorInfo* t = impl_->file->find("per_layer_token_embd.weight");
    if (t == nullptr) {
        err = "per_layer_token_embd.weight is not in " + gguf_path;
        close();
        return false;
    }
    // [160, 320001536]: ne0 = 160 is the FAST axis, so the ROW index is shape[1] and a row is contiguous.
    if (t->shape.size() != 2 || t->shape[0] != (uint64_t) PLE_HEAD_DIM) {
        err = "per_layer_token_embd.weight has an unexpected shape";
        close();
        return false;
    }
    // The table's type is looked up in the format list (ple_formats()): IQ4_NL (ISTA-DASLab's shard 2, the original's
    // own GGUF), Q5_0 (#296: OrcaRouter's GGUF), or the FP8 table as shipped: I8 bytes marked strata.ple.format =
    // f8_e4m3 with strata.ple.scale (tools/ple_fp8_pack.py)
    const PleFormatInfo* fmt = ple_format_for_type(t->type_name());
    if (fmt == nullptr) {
        err = std::string("per_layer_token_embd.weight is ") + t->type_name() + ", not " + ple_format_list();
        close();
        return false;
    }
    if (fmt->needs_scale) {
        const MetaValue* f = impl_->file->get("strata.ple.format");
        const MetaValue* s = impl_->file->get("strata.ple.scale");
        if (f == nullptr || f->s != "f8_e4m3" || s == nullptr || !(s->num() > 0.0)) {
            err = "per_layer_token_embd.weight is I8 without strata.ple.format = f8_e4m3 and a positive strata.ple.scale";
            close();
            return false;
        }
        impl_->scale = (float) s->num();
    }
    impl_->fmt = fmt;
    impl_->rb = fmt->row_bytes;
    // PleReader's row_bytes has been a runtime parameter since the FP8 table (160 B rows) needed it; Q5_0's
    // 110 B rows go through the exact same generic path (ple_reader_test --selftest covers both 90 and 110 B
    // rows: straddling, caching, in-flight tickets, keep-alive). This refusal was stale.
    impl_->n_rows = t->shape[1];

    // THE CHECK THAT MAKES THE OFFSET FALSIFIABLE.  The manifest's `shard2_tensor.offset` is 0, but that is
    // the offset within the GGUF's DATA SECTION: the file's first 192 bytes are a header, and reading at 0
    // would decode the header plus 192 bytes of shifted rows - still plausible IQ4_NL, and wrong for every
    // row.  `GgufFile` parses the header, so `tensor_data` is already correct; this asserts the tensor
    // exactly fills the file from there, which is what makes the whole arrangement checkable rather than
    // assumed.  A wrong data offset would leave a different remainder.
    // A shard may hold other tensors too (Swift 1.5's shard 1 holds layers 0-12 and the table): the table must
    // then fit inside the file at its own offset; alone in its shard (the original's shard 2) it fills it exactly.
    // The size arithmetic is checked before it is used (#865): a header that claims enough rows to wrap the 64-bit
    // product, or a data section that starts past the end of the file, must be refused, not compared after wrapping.
    if (impl_->n_rows > std::numeric_limits<uint64_t>::max() / impl_->rb ||
        impl_->file->data_start() > impl_->file->file_size()) {
        err = "PLE table size overflow or invalid data offset in " + gguf_path;
        close();
        return false;
    }
    const uint64_t need = impl_->n_rows * (uint64_t) impl_->rb;
    const uint64_t have = impl_->file->file_size() - impl_->file->data_start();
    const bool alone = impl_->file->tensors().size() == 1;
    if (t->offset > have || need > have - t->offset || (alone && (t->offset != 0 || need != have))) {
        char buf[256];
        std::snprintf(buf, sizeof buf,
                      "PLE table size mismatch: %llu rows x %d B = %llu at offset %llu, but the file holds %llu from "
                      "data_start %llu",
                      (unsigned long long) impl_->n_rows, (int) impl_->rb, (unsigned long long) need,
                      (unsigned long long) t->offset, (unsigned long long) have,
                      (unsigned long long) impl_->file->data_start());
        err = buf;
        if (alone && have > need)   // #657: extra bytes after the table: a damaged or wrong file, not a format question
            err += "; the file is longer than its table - delete it and its .done mark and run setup again";
        close();
        return false;
    }
    impl_->data = impl_->file->tensor_data(*t);        // only a table that fits the file is mapped
    if (io.mode == PleIo::Direct) {
        // The parse above is the validated source of the offset; the mapping itself is not kept, so no page of
        // the table can enter this process's working set or the file cache through it.
        const uint64_t table_offset = impl_->file->data_start() + t->offset;
        const uint64_t n_rows = impl_->n_rows;
        delete impl_->file;
        impl_->file = nullptr;
        impl_->data = nullptr;
        if (!impl_->reader.open(gguf_path, table_offset, n_rows, io.max_inflight, io.cache_rows, err, io.io_thread,
                                impl_->rb)) {
            close();
            return false;
        }
        impl_->reader.set_keepalive(io.keepalive_ms, io.keepalive_window_s);
        impl_->n_rows = n_rows;
    }
    if (io.mode == PleIo::Mmap && io.lock && impl_->data != nullptr) {
#if defined(_WIN32)
        // Windows has no mlock: the platform helper raises the process's minimum working set and
        // VirtualLocks the mapped table instead (ordinary accounts hold the privilege). A failure
        // falls back to touching the pages, i.e. plain mmap behaviour, with a warning - same as POSIX.
        // A table that is more than half of the RAM leaves the rest of the process little (the locked pages cannot be
        // paged out): say so, and go on (recommend, never force).
        MEMORYSTATUSEX ms;
        ms.dwLength = sizeof(ms);
        if (GlobalMemoryStatusEx(&ms) && need > ms.ullTotalPhys / 2)
            std::fprintf(stderr, "strata: warning: --ple-io ram locks %.1f GiB, more than half of this PC's %.1f GiB of RAM\n",
                         (double) need / (1ull << 30), (double) ms.ullTotalPhys / (1ull << 30));
        const strata::platform::LockResult lr =
            strata::platform::lock_resident((void*) impl_->data, need);
        if (lr.ok && lr.locked_bytes >= need) {
            impl_->locked = true;
        } else {
            // a partial lock is not a lock: the unlocked rest would fault on the token path, so give it back
            if (lr.locked_bytes > 0) strata::platform::unlock_resident((void*) impl_->data, lr.locked_bytes);
            std::fprintf(stderr, "strata: PLE table lock failed (%s): touching its pages instead\n",
                         lr.note.c_str());
            volatile uint8_t sink = 0;
            for (uint64_t off = 0; off < need; off += 4096) sink = sink + impl_->data[off];
            (void) sink;
        }
#else
        const uint64_t page = 4096;
        const uintptr_t a0 = (uintptr_t) impl_->data & ~(uintptr_t) (page - 1);
        const uintptr_t a1 = (uintptr_t) impl_->data + (uintptr_t) need;
        madvise((void*) a0, a1 - a0, MADV_WILLNEED);
        // Fault the table in with several threads first: mlock (and a single toucher) brings it in one page at a time
        // from one thread, ~0.5 GB/s from a cold file - 106 s for a 54 GB Q8_0 table.  The pages are then resident
        // and mlock only pins them.
        {
            constexpr uintptr_t kPiece = 64ull << 20;
            const uintptr_t pieces = (a1 - a0 + kPiece - 1) / kPiece;
            const unsigned threads = std::max(1u, std::min(16u, std::thread::hardware_concurrency()));
            std::atomic<uintptr_t> next{0};
            auto touch = [&] {
                volatile uint8_t sink = 0;
                for (uintptr_t i; (i = next.fetch_add(1)) < pieces;)
                    for (uintptr_t p = a0 + i * kPiece, e = std::min(a1, p + kPiece); p < e; p += page)
                        sink = sink + *(const volatile uint8_t*) p;
                (void) sink;
            };
            std::vector<std::thread> pool;
            for (unsigned t = 0; t < threads; ++t) pool.emplace_back(touch);
            for (auto& t : pool) t.join();
        }
        if (mlock((const void*) a0, a1 - a0) == 0) {
            impl_->locked = true;
        } else {
            std::fprintf(stderr, "strata: PLE table mlock failed (%s; raise `ulimit -l`): its pages stay faulted in "
                                 "but may be reclaimed\n", std::strerror(errno));
        }
#endif
    }
    impl_->mode = io.mode;
    return true;
}

void PleTable::wait_prefetches() {
    if (impl_->n_prefetch == 0) return;
    std::string dummy;
    for (size_t i = 0; i < impl_->n_prefetch; ++i) (void) impl_->reader.collect(impl_->prefetch_tickets[i], dummy);
    impl_->n_prefetch = 0;
}

void PleTable::prefetch_rows(const uint32_t* rows16) {
    if (impl_->mode == PleIo::Direct && impl_->reader.is_open()) {
        for (size_t i = 0; i < impl_->n_prefetch; ++i)
            if (std::memcmp(impl_->prefetch_keys[i], rows16, sizeof(uint32_t) * PLE_N_HEADS) == 0) return;
        if (impl_->n_prefetch >= Impl::kMaxPrefetch) wait_prefetches();
        const size_t slot = impl_->n_prefetch++;
        std::memcpy(impl_->prefetch_keys[slot], rows16, sizeof(uint32_t) * PLE_N_HEADS);
        impl_->prefetch_tickets[slot] = impl_->reader.issue(impl_->prefetch_keys[slot], PLE_N_HEADS, impl_->prefetch_raw[slot]);
        return;
    }
#if defined(_WIN32)
    if (impl_->data != nullptr && g_ple_prefetch) {
        WIN32_MEMORY_RANGE_ENTRY ranges[PLE_N_HEADS];
        ULONG_PTR n = 0;
        for (int h = 0; h < PLE_N_HEADS; ++h) {
            if (rows16[h] >= impl_->n_rows) continue;
            ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->rb);
            ranges[n].NumberOfBytes = impl_->rb;
            ++n;
        }
        if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0);
    }
#endif
}

void PleTable::close() {
    wait_prefetches();
    impl_->reader.close();
    impl_->pending = false;
    impl_->locked = false;   // the unmap below releases the lock
    impl_->mode = PleIo::Mmap;
    delete impl_->file;
    impl_->file = nullptr;
    impl_->data = nullptr;
    impl_->n_rows = 0;
    impl_->rb = PLE_ROW_BYTES;
    impl_->fmt = &ple_format_info(PleFormat::IQ4_NL);
    impl_->scale = 1.0f;
    impl_->bytes_read = 0;
}

bool PleTable::is_open() const { return impl_->data != nullptr || impl_->reader.is_open(); }
bool PleTable::locked() const { return impl_->locked; }
const char* PleTable::format() const { return impl_->fmt->name; }
PleIo PleTable::mode() const { return impl_->mode; }
uint64_t PleTable::rows() const { return impl_->n_rows; }
uint64_t PleTable::bytes_read() const { return impl_->bytes_read; }

void PleTable::read_row(uint32_t row, float* out160) const {
    if (impl_->mode == PleIo::Direct && impl_->reader.is_open()) {
        uint8_t raw[PLE_ROW_BYTES_MAX];
        std::string err;
        const auto t = impl_->reader.issue(&row, 1, raw);
        if (!impl_->reader.collect(t, err)) {
            std::memset(out160, 0, (size_t) PLE_HEAD_DIM * sizeof(float));
            return;
        }
        impl_->decode(raw, out160);
        impl_->bytes_read += impl_->rb;
        return;
    }
    if (impl_->data == nullptr || row >= impl_->n_rows) {
        std::memset(out160, 0, (size_t) PLE_HEAD_DIM * sizeof(float));
        return;
    }
    impl_->decode(impl_->data + (size_t) row * impl_->rb, out160);
    impl_->bytes_read += impl_->rb;
}

bool PleTable::issue(const uint32_t* rows16) {
    wait_prefetches();
    std::memcpy(impl_->rows, rows16, sizeof impl_->rows);
    if (impl_->mode == PleIo::Direct) {
        if (impl_->pending) return false;              // one token in flight per table
        impl_->ticket = impl_->reader.issue(impl_->rows, PLE_N_HEADS, impl_->raw);
        impl_->pending = true;
        return true;
    }
#if defined(_WIN32)
    if (impl_->data != nullptr && g_ple_prefetch) {
        WIN32_MEMORY_RANGE_ENTRY ranges[PLE_N_HEADS];
        ULONG_PTR n = 0;
        for (int h = 0; h < PLE_N_HEADS; ++h) {
            if (rows16[h] >= impl_->n_rows) continue;
            ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->rb);
            ranges[n].NumberOfBytes = impl_->rb;
            ++n;
        }
        if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0);
    }
#endif
    impl_->pending = true;
    return true;
}

bool PleTable::collect(float* out2560, std::string& err) {
    if (!impl_->pending) { err = "PleTable::collect without issue"; return false; }
    impl_->pending = false;
    if (impl_->mode == PleIo::Direct) {
        if (!impl_->reader.collect(impl_->ticket, err)) return false;
        for (int h = 0; h < PLE_N_HEADS; ++h)
            impl_->decode(impl_->raw + (size_t) h * impl_->rb, out2560 + (size_t) h * PLE_HEAD_DIM);
        impl_->bytes_read += (uint64_t) PLE_N_HEADS * impl_->rb;
        return true;
    }
    for (int h = 0; h < PLE_N_HEADS; ++h) read_row(impl_->rows[h], out2560 + (size_t) h * PLE_HEAD_DIM);
    return true;
}

bool PleTable::gather_batch(const uint32_t* rows, size_t n_tokens, float* out, std::string& err) {
    if (impl_->pending) { err = "PleTable::gather_batch while a token is in flight"; return false; }
    const size_t n = n_tokens * (size_t) PLE_N_HEADS;
    if (impl_->mode == PleIo::Direct) {
        // Fast path: if all requested tokens match our in-flight/completed prefetch slots, collect them directly
        // without issuing a second PleReader ticket or copying through RowCache.
        if (impl_->n_prefetch >= n_tokens && n_tokens <= Impl::kMaxPrefetch) {
            bool exact_match = true;
            for (size_t t = 0; t < n_tokens; ++t) {
                if (std::memcmp(impl_->prefetch_keys[t], rows + t * PLE_N_HEADS, sizeof(uint32_t) * PLE_N_HEADS) != 0) {
                    exact_match = false;
                    break;
                }
            }
            if (exact_match) {
                for (size_t i = 0; i < impl_->n_prefetch; ++i) {
                    if (!impl_->reader.collect(impl_->prefetch_tickets[i], err)) {
                        impl_->n_prefetch = 0;
                        return false;
                    }
                }
                impl_->n_prefetch = 0;
                for (size_t t = 0; t < n_tokens; ++t) {
                    for (int h = 0; h < PLE_N_HEADS; ++h) {
                        impl_->decode(impl_->prefetch_raw[t] + (size_t) h * impl_->rb,
                                      out + (t * PLE_N_HEADS + (size_t) h) * PLE_HEAD_DIM);
                    }
                }
                impl_->bytes_read += (uint64_t) n * impl_->rb;
                return true;
            }
        }
        wait_prefetches();
        std::vector<uint8_t> raw(n * impl_->rb);
        const auto ticket = impl_->reader.issue(rows, n, raw.data());
        if (!impl_->reader.collect(ticket, err)) return false;
        for (size_t i = 0; i < n; ++i) impl_->decode(raw.data() + i * impl_->rb, out + i * PLE_HEAD_DIM);
        impl_->bytes_read += (uint64_t) n * impl_->rb;
        return true;
    }
    wait_prefetches();
    // #1984 #1971: a mapped table that is not in the file cache served this loop as serial demand faults, one random
    // 4 KiB page at a time (16 rows a token, 16 x the chunk's tokens for a prompt), and the first chunk of a prompt has
    // only layer 0 to hide it behind: a cold 27 GB table at ~0.2-1 ms a fault is minutes with no progress mark, which
    // is what the stall watchdog stops.  Ask the OS for the next window's pages in one go (PrefetchVirtualMemory /
    // MADV_WILLNEED, a hint that cannot change a value: an unfetched page is faulted by the read) while this window is read.
    constexpr size_t kWin = 4096;   // rows a window
    auto hint = [&](size_t a, size_t b) {
        if (impl_->data == nullptr || !g_ple_prefetch) return;
#if defined(_WIN32)
        WIN32_MEMORY_RANGE_ENTRY ranges[kWin];
        ULONG_PTR k = 0;
        for (size_t i = a; i < b; ++i) {
            if (rows[i] >= impl_->n_rows) continue;
            ranges[k].VirtualAddress = (PVOID) (impl_->data + (size_t) rows[i] * impl_->rb);
            ranges[k].NumberOfBytes = impl_->rb;
            ++k;
        }
        if (k > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), k, ranges, 0);
#else
        for (size_t i = a; i < b; ++i) {
            if (rows[i] >= impl_->n_rows) continue;
            const uintptr_t r0 = (uintptr_t) (impl_->data + (size_t) rows[i] * impl_->rb);
            const uintptr_t p0 = r0 & ~(uintptr_t) 4095;
            (void) madvise((void*) p0, (size_t) (r0 + impl_->rb - p0), MADV_WILLNEED);
        }
#endif
    };
    hint(0, std::min(n, kWin));
    for (size_t w = 0; w < n; w += kWin) {
        if (w + kWin < n) hint(w + kWin, std::min(n, w + 2 * kWin));
        const size_t e = std::min(n, w + kWin);
        for (size_t i = w; i < e; ++i) read_row(rows[i], out + i * PLE_HEAD_DIM);
    }
    return true;
}

void PleTable::set_injected_delay_us(double us) { impl_->reader.set_injected_delay_us(us); }

std::string PleTable::io_report() const {
    if (impl_->mode != PleIo::Direct || !impl_->reader.is_open()) return {};
    const strata::ngram::ReaderStats s = impl_->reader.snapshot();
    char buf[400];
    int n = std::snprintf(buf, sizeof buf,
                  "ple io: %llu rows, %.1f%% row-cache hits, %llu SSD reads (%.1f MB), read p50 %.0f us p99 %.0f us, "
                  "blocked %.3f ms total (submit %.3f ms), cache %llu/%llu rows",
                  (unsigned long long) s.requests, s.requests ? 100.0 * (double) s.cache_hits / (double) s.requests : 0.0,
                  (unsigned long long) s.reads, (double) s.bytes / 1e6, s.percentile(0.5), s.percentile(0.99),
                  s.wait_us / 1000.0, s.submit_us / 1000.0, (unsigned long long) impl_->reader.cache_size(),
                  (unsigned long long) impl_->reader.cache_capacity());
    if (s.keepalive_reads > 0 && n > 0 && n < (int) sizeof buf)
        std::snprintf(buf + n, sizeof buf - (size_t) n, ", SSD kept awake by %llu reads (slowest %.1f ms)",
                      (unsigned long long) s.keepalive_reads, s.keepalive_us_max / 1000.0);
    return buf;
}

void PleTable::gather(const uint32_t* rows16, float* out2560) const {
    if (impl_->mode == PleIo::Direct) {
        // `gather` stays const for its existing callers; the reader's state is the table's I/O state.
        PleTable* self = const_cast<PleTable*>(this);
        std::string err;
        if (!self->issue(rows16) || !self->collect(out2560, err)) {
            std::fprintf(stderr, "PleTable::gather: %s\n", err.empty() ? "a token is already in flight" : err.c_str());
            std::memset(out2560, 0, (size_t) NG_N_EMBD * sizeof(float));
        }
        return;
    }
    // ================================ SIXTEEN SERIAL PAGE FAULTS, MEASURED ================================
    //
    // **THIS COST 2.10-2.61 ms PER TOKEN AND HAD NEVER BEEN IN THE PLAN'S BUDGET AT ALL.**  The round-309
    // `token host phases` line put it second behind the layer loop among avoidable terms, and the arithmetic
    // says why: the table is 320,001,536 rows of `PLE_ROW_BYTES` = 90 B in a 26.8 GB mapping, so the sixteen
    // rows a token needs are 1,440 B - **0.5 MB/s**.  That is not bandwidth, it is latency: sixteen reads into
    // sixteen different 4 KB pages scattered across 26.8 GB, taken ONE AT A TIME, and on this machine the PLE
    // shard is 26.8 GB against 63 GB of RAM that the 31.6 GB expert arena is also competing for, so they are
    // not in the OS cache.  Sixteen serial NVMe reads at ~150 us is 2.4 ms, which is the measurement.
    //
    // `PrefetchVirtualMemory` issues all sixteen in ONE call and lets them complete in parallel.  It is a hint
    // and cannot change the answer - a range it does not fetch is simply faulted in by the read that follows -
    // so the only risk is that it does nothing.
#if defined(_WIN32)
    if (impl_->data != nullptr && g_ple_prefetch) {
        WIN32_MEMORY_RANGE_ENTRY ranges[PLE_N_HEADS];
        ULONG_PTR n = 0;
        for (int h = 0; h < PLE_N_HEADS; ++h) {
            // Out-of-range rows are handled by `read_row` as zeros and have no address to prefetch.
            if (rows16[h] >= impl_->n_rows) continue;
            ranges[n].VirtualAddress = (PVOID) (impl_->data + (size_t) rows16[h] * impl_->rb);
            ranges[n].NumberOfBytes = impl_->rb;
            ++n;
        }
        if (n > 0) (void) PrefetchVirtualMemory(GetCurrentProcess(), n, ranges, 0);
    }
#endif
    // HEAD-SLOWEST, which is what `ggml_get_rows` does and what the source's own comment says: head h's 160
    // values occupy [h*160, (h+1)*160).  A head-fastest layout would put element (d, h) at d*16 + h and needs
    // a real transpose - a reshape of the same flat buffer compares equal and would make the check vacuous.
    for (int h = 0; h < PLE_N_HEADS; ++h) read_row(rows16[h], out2560 + (size_t) h * PLE_HEAD_DIM);
}

}  // namespace strata::kernels