// The first-letters worker for the AAC reference app (abbreviation
// expansion): the person types only the first letter of each word ("iwtgh"),
// and a small language model taught to do this proposes whole sentences ("I
// want to go home"), fitting what their partner said when that is given.
// Qwen3 0.6B, fine-tuned for it and stored in 4 bits, runs on onnxruntime-
// web's WebGPU execution provider; the sentences come from a beam search that
// keeps only words beginning with the letters typed, in order (the same
// algorithm as the model's reference decoder, ae_decode.py, and the
// prototype's ae-engine.mjs, both in the model library: tools/ae/README.md).
//
// Nothing leaves the device, and nothing is kept: what is typed and what the
// partner said come in one message and go once the sentences are found. The
// worker never fetches: the runtime (onnxruntime-web's bundle and its
// WebAssembly), the tokenizer library (Hugging Face's @huggingface/
// tokenizers, beside this file) and the model's files come from the model
// store (Cache API "aac-models-v1", each file SHA-256 checked when it was
// stored), so it works offline.
//
// Graphics chip only. The session is created with the processor fallback
// switched off (session.disable_cpu_ep_fallback): a device that cannot run
// the whole graph on its graphics chip fails to load it, rather than running
// it on the processor. All 371 nodes run on WebGPU (checked in Chrome and in
// WebKit).
//
// The worker has a GPU device of its own (every limit and feature the
// adapter offers, as the runtime's own would), given to the runtime for the
// session. Unloading releases the session and destroys the device, which
// gives all of its graphics memory back; only then may the page end the
// worker. (In WebKit, ending a worker whose WebGPU device the runtime's JSPI
// build still holds ends the page with it: seen with Playwright's WebKit,
// 2026-10-01.)
//
// A module worker, started by src/bindings/abbrev.js from a Blob URL. It
// imports onnxruntime-web's JSPI build (shared with the listener: it is
// offered only in browsers with JavaScript Promise Integration) and the
// tokenizer library from Blob URLs too, and hands the runtime its
// WebAssembly's bytes, so nothing is looked for by name. JSPI lets the runtime
// read the weights one tensor at a time, straight into graphics memory, so
// the model is never held whole in memory. One thread.
//
// Messages in (each with an `id`, answered in order with {id, ok, ...} or
// {id, ok: false, error}; and, unasked, {event: "lost", text} when the
// graphics chip takes the device back):
//   {type: "init", bundle, wasm, tokenizers}   model-store URLs of the runtime
//                                               and the tokenizer library
//   {type: "load", key, model}                 model: {graph, data: [{name,
//                                               parts: [{url, bytes}]}],
//                                               tokenizer, tokenizer_config,
//                                               config, vocab, vocab_table}
//                                               (URLs); loads it, and finds
//                                               one set of sentences so the
//                                               graphics programs are compiled
//   {type: "check"}                            one set of sentences, which
//                                               must not be empty
//   {type: "expand", letters, partner, near_keys}
//                                            -> {sentences: [text, ...], ms}
//   {type: "cancel"}                           the expansion running now, and
//                                               any asked for before, end
//                                               "superseded"
//   {type: "unload"}                           releases the model and destroys
//                                               the device; the worker can then
//                                               be ended, or load again
// An expansion asked for after another (or an unload) supersedes it: the
// earlier one stops at its next model run and ends "superseded".
//
// Vendored into vendor/hf-tokenizers-<version>-<revision>/ as ae-worker.js by
// vendor.sh, unchanged. Its bytes are pinned in crates/aac-web/src/abbrev.rs
// (RUNTIME_FILES): change this file only together with a new
// TOKENIZERS_VENDOR_REV and its new size and hash there.

const CACHE_NAME = "aac-models-v1";

// The search, as tuned on the development set (tools/ae/README.md): eight
// hypotheses, five sentences offered.
const BEAM = 8;
const OFFERED = 5;
// A slip to a key next to the one meant costs this much (natural log of the
// probability), when the person asks for slips to be forgiven.
const NEAR_KEY_COST = 1.0;

let ort = null;
let Tokenizer = null;
let engine = null; // the loaded model
let key = "";
// The latest expansion (or cancel) asked for: any other is superseded.
let latest = 0;

function describe(err) {
    return String(err && err.message ? err.message : err);
}

async function stored(cache, url) {
    const res = await cache.match(url);
    if (!res) throw new Error(`first letters is not on this device (${url.split("/").pop()})`);
    return res;
}

async function importStored(cache, url) {
    const code = await (await stored(cache, url)).blob();
    const blobUrl = URL.createObjectURL(new Blob([code], { type: "text/javascript" }));
    try {
        return await import(blobUrl);
    } finally {
        URL.revokeObjectURL(blobUrl);
    }
}

async function init(msg) {
    if (ort && Tokenizer) return {};
    if (typeof caches === "undefined") throw new Error("this browser keeps no model store in workers");
    const cache = await caches.open(CACHE_NAME);
    const [runtime, tokenizers, wasm] = await Promise.all([
        importStored(cache, msg.bundle),
        importStored(cache, msg.tokenizers),
        stored(cache, msg.wasm).then((r) => r.arrayBuffer()),
    ]);
    runtime.env.wasm.wasmBinary = wasm; // never fetched by name
    runtime.env.wasm.numThreads = 1;
    runtime.env.wasm.proxy = false;
    runtime.env.logLevel = "error";
    ort = runtime;
    Tokenizer = tokenizers.Tokenizer;
    return {};
}

// ─── Loading ───────────────────────────────────────────────────────────────

// The part of the model the runtime is reading now. It reads the weights one
// tensor after another, in file order, so it is mostly the same part again:
// one part (64 MiB at most) is held at a time, never the whole model.
let held = { url: "", blob: null };

function partBlob(cache, url) {
    if (held.url !== url) held = { url, blob: stored(cache, url).then((r) => r.blob()) };
    return held.blob;
}

// A weights file as the runtime reads it: a Blob it slices one tensor at a
// time, whose bytes stay in the model store until then. A file larger than a
// part (one tensor, the 74 MiB table of word embeddings, cannot be split) is
// stored as several parts, read here as one.
class StoredFile extends Blob {
    constructor(cache, parts) {
        super([]);
        this.cache = cache;
        this.parts = parts;
        this.bytes = parts.reduce((n, p) => n + p.bytes, 0);
    }

    get size() {
        return this.bytes;
    }

    slice(start, end) {
        const { cache, parts } = this;
        return {
            stream: () => new ReadableStream({
                async start(controller) {
                    let at = 0;
                    for (const p of parts) {
                        const from = Math.max(start, at);
                        const to = Math.min(end, at + p.bytes);
                        if (from < to) {
                            const blob = await partBlob(cache, p.url);
                            controller.enqueue(new Uint8Array(await blob.slice(from - at, to - at).arrayBuffer()));
                        }
                        at += p.bytes;
                    }
                    controller.close();
                },
            }),
        };
    }
}

async function load(msg) {
    if (engine && key === msg.key) return {};
    await unload();
    const cache = await caches.open(CACHE_NAME);
    const m = msg.model;
    const json = async (url) => (await stored(cache, url)).json();
    const [graph, tj, tc, cfg, vj, vb] = await Promise.all([
        stored(cache, m.graph).then((r) => r.arrayBuffer()),
        json(m.tokenizer),
        json(m.tokenizer_config),
        json(m.config),
        json(m.vocab),
        stored(cache, m.vocab_table).then((r) => r.arrayBuffer()),
    ]);
    const e = new Engine(new Tokenizer(tj, tc), cfg.model.decoder, vj, vb);
    try {
        e.device = await ownDevice();
        // The graphics chip may take the device back (a driver reset, memory
        // short): nothing more runs on it, and the page is told, so it lets
        // this worker go and starts afresh.
        e.device.lost.then((info) => {
            if (info.reason === "destroyed") return; // let go on purpose
            e.lost = true;
            self.postMessage({ event: "lost", text: "the graphics chip stopped it" });
        });
        await e.open(new Uint8Array(graph), m.data.map((d) => ({ path: d.name, data: new StoredFile(cache, d.parts) })));
    } catch (err) {
        await e.release();
        throw err;
    } finally {
        held = { url: "", blob: null };
    }
    engine = e;
    key = msg.key;
    // The first run compiles the graphics chip's programs for the prompt
    // and for the batched steps (a fraction of a second, once).
    await engine.expand("ih", null, { nearKeys: false, live: () => true });
    return {};
}

// A GPU device for this worker alone: every limit and feature the adapter
// offers (the runtime asks for many of them: storage buffers per shader,
// workgroup sizes, half precision).
async function ownDevice() {
    const adapter = navigator.gpu && (await navigator.gpu.requestAdapter());
    if (!adapter) throw new Error("this browser offers no graphics chip to first letters");
    const requiredLimits = {};
    for (const k in adapter.limits) if (typeof adapter.limits[k] === "number") requiredLimits[k] = adapter.limits[k];
    return adapter.requestDevice({ requiredFeatures: [...adapter.features], requiredLimits });
}

async function unload() {
    latest++;
    const old = engine;
    engine = null;
    key = "";
    if (old) await old.release();
    return {};
}

function loaded() {
    if (!engine) throw new Error("first letters is not loaded");
}

async function check() {
    loaded();
    const r = await engine.expand("iwtgh", "What do you want to do?", { nearKeys: false, live: () => true });
    if (r.sentences.length === 0) throw new Error("it found no sentence");
    return {};
}

async function expand(msg) {
    loaded();
    const mine = msg.id;
    const live = () => latest === mine;
    if (!live()) throw new Error("superseded");
    return engine.expand(String(msg.letters || ""), msg.partner ? String(msg.partner) : null, { nearKeys: !!msg.near_keys, live });
}

// ─── The letters ───────────────────────────────────────────────────────────

// The keyboard the slips are measured on (QWERTY, the paper's 3 x 10 grid).
const ROWS = ["qwertyuiop", "asdfghjkl;", "zxcvbnm,./"];
const POS = new Map();
ROWS.forEach((row, r) => [...row].forEach((ch, c) => POS.set(ch, [c, r])));
const isAlnum = (ch) => /[\p{L}\p{N}]/u.test(ch);

function normText(s) {
    s = s.normalize("NFKC").replace(/[’‘ʼ]/g, "'").replace(/[“”]/g, '"')
        .replace(/—/g, " - ").replace(/–/g, "-");
    return s.split(/\s+/).filter(Boolean).join(" ");
}

function stripFinal(s) {
    s = s.trimEnd();
    let i = s.length;
    while (i > 0 && !isAlnum(s[i - 1])) i--;
    return s.slice(0, i);
}

const normMatch = (s) => stripFinal(normText(s)).toLowerCase().split(" ").filter(Boolean).join(" ");

// The prompt the model was taught with: "Partner: …\nLetters: i w t g h\nReply:".
function makePrompt(letters, partner) {
    const head = partner ? `Partner: ${normText(partner)}\n` : "";
    return `${head}Letters: ${[...letters].join(" ")}\nReply:`;
}

function neighbours(ch) {
    const out = new Set();
    const p = POS.get(ch);
    if (!p) return out;
    for (let dr = -1; dr <= 1; dr++) {
        for (let dc = -1; dc <= 1; dc++) {
            if (!dr && !dc) continue;
            const r = p[1] + dr;
            const c = p[0] + dc;
            if (r >= 0 && r < 3 && c >= 0 && c < 10 && /[a-z]/.test(ROWS[r][c])) out.add(ROWS[r][c]);
        }
    }
    return out;
}

// Can the text `e` follow a hypothesis that has used `p` of the letters `Q`?
// -> [letters used after it, cost] or null.
function advance(e, Q, p, punct, nearCost) {
    let cost = 0;
    for (const ch of e) {
        if (isAlnum(ch)) {
            if (p < Q.length && Q[p] === ch) p++;
            else if (nearCost != null && p < Q.length && neighbours(Q[p]).has(ch)) {
                p++;
                cost += nearCost;
            } else return null;
        } else {
            if (punct === "free") continue;
            if (p < Q.length && Q[p] === ch) p++;
            else if (p >= Q.length) continue;
            else return null;
        }
    }
    return [p, cost];
}

let F16 = null;
function f16table() {
    if (F16) return F16;
    F16 = new Float32Array(65536);
    for (let h = 0; h < 65536; h++) {
        const s = h & 0x8000 ? -1 : 1;
        const e = (h >> 10) & 31;
        const m = h & 1023;
        F16[h] = e === 0 ? s * m * 2 ** -24 : e === 31 ? (m ? NaN : s * Infinity) : s * (1 + m / 1024) * 2 ** (e - 15);
    }
    return F16;
}

// ─── The model ─────────────────────────────────────────────────────────────

// The graph's contract (the model library's onnx_patch.py): inputs input_ids
// int64 [B,S], seqlens_k int32 [B], total_sequence_length int32 [],
// beam_idx int32 [B], past_key_values.N.{key,value} fp16 [B_prev,kvh,P,hd];
// outputs logits fp16 [B,S,V], logsumexp f32 [B,S], present.N.* (kept on the
// graphics chip, and reordered there by beam_idx at the next step).
class Engine {
    constructor(tokenizer, dec, vj, vb) {
        this.tok = tokenizer;
        this.kvh = dec.num_key_value_heads;
        this.hd = dec.head_size;
        this.L = dec.num_hidden_layers;
        this.V = vj.V;
        this.eos = vj.eos[0];
        // Which tokens may come next, by what they start with: one table for
        // after a letter, one for after anything else.
        const gN = new Uint16Array(vb, 0, this.V);
        const gL = new Uint16Array(vb, 2 * this.V, this.V);
        this.last = new Int8Array(vb, 4 * this.V, this.V);
        const build = (g) => {
            const counts = new Map();
            for (let i = 0; i < this.V; i++) if (g[i] !== 0xffff) counts.set(g[i], (counts.get(g[i]) || 0) + 1);
            const arrs = new Map();
            const fill = new Map();
            for (const [k, n] of counts) {
                arrs.set(k, new Int32Array(n));
                fill.set(k, 0);
            }
            for (let i = 0; i < this.V; i++) {
                if (g[i] === 0xffff) continue;
                const k = g[i];
                const f = fill.get(k);
                arrs.get(k)[f] = i;
                fill.set(k, f + 1);
            }
            const byFirst = new Map();
            for (const k of arrs.keys()) {
                const h = vj.groups[k].slice(0, 1);
                if (!byFirst.has(h)) byFirst.set(h, []);
                byFirst.get(h).push(k);
            }
            return { arrs, byFirst };
        };
        this.tables = { N: build(gN), L: build(gL) };
        this.groups = vj.groups;
        const io = dec.outputs || {};
        const kn = io.present_key_names || "present.%d.key";
        const vn = io.present_value_names || "present.%d.value";
        this.pres = [];
        for (let i = 0; i < this.L; i++) this.pres.push(kn.replace("%d", i), vn.replace("%d", i));
        this.pasts = this.pres.map((n) => n.replace("present", "past_key_values"));
        this.sess = null;
        this.device = null;
    }

    async open(graph, externalData) {
        this.sess = await ort.InferenceSession.create(graph, {
            // Options reach the WebGPU provider only when it is named as an
            // object; the device is this worker's own.
            executionProviders: [{ name: "webgpu", device: this.device }],
            externalData,
            graphOptimizationLevel: "all",
            logSeverityLevel: 3,
            // The caches stay on the graphics chip from one step to the next.
            preferredOutputLocation: Object.fromEntries(this.pres.map((n) => [n, "gpu-buffer"])),
            // Graphics chip or nothing: no node may fall back to the processor.
            extra: { session: { disable_cpu_ep_fallback: "1" } },
        });
    }

    // The session, then the device: all of its graphics memory goes back.
    async release() {
        if (this.sess) await this.sess.release().catch(() => {});
        this.sess = null;
        if (this.device) this.device.destroy();
        this.device = null;
    }

    async run(rows, parents, state) {
        const S = rows[0].length;
        const B = rows.length;
        state.T += S;
        const flat = new BigInt64Array(B * S);
        rows.forEach((r, i) => r.forEach((t, j) => { flat[i * S + j] = BigInt(t); }));
        const feeds = {
            input_ids: new ort.Tensor("int64", flat, [B, S]),
            seqlens_k: new ort.Tensor("int32", new Int32Array(B).fill(state.T - 1), [B]),
            total_sequence_length: new ort.Tensor("int32", Int32Array.of(state.T), []),
            beam_idx: new ort.Tensor("int32", Int32Array.from(parents), [B]),
            ...state.past,
        };
        const o = await this.sess.run(feeds);
        for (const t of Object.values(state.past)) t.dispose?.();
        state.past = {};
        this.pres.forEach((n, k) => { state.past[this.pasts[k]] = o[n]; });
        const Vo = o.logits.dims[2];
        const So = o.logits.dims[1];
        const raw = o.logits.data;
        const u16 = raw instanceof Uint16Array ? raw : new Uint16Array(raw.buffer, raw.byteOffset, raw.length);
        return { u16, lse: o.logsumexp.data, Vo, So, B };
    }

    // Up to OFFERED sentences whose words begin with `typed`'s letters, in
    // order (punctuation typed must be where it was typed; none typed, and
    // it may go anywhere). `live()` false stops it ("superseded").
    async expand(typed, partner, { nearKeys, live }) {
        const t0 = performance.now();
        if (this.lost) throw new Error("the graphics chip stopped it");
        // What a keyboard may type for the same letter or mark: compatibility
        // forms folded (NFKC), curly quotes straight.
        typed = typed.normalize("NFKC").replace(/[\u2019\u2018\u02bc]/g, "'").replace(/[\u201c\u201d]/g, '"');
        let Q = typed.toLowerCase().replace(/\s+/g, "");
        const punct = [...Q].some((c) => !isAlnum(c)) ? "strict" : "free";
        if (punct === "free") Q = [...Q].filter(isAlnum).join("");
        if (!Q) return { sentences: [], ms: 0 };
        const nearCost = nearKeys ? NEAR_KEY_COST : null;
        const maxSteps = Math.min(48, 6 + 3 * Q.length);
        const ids = this.tok.encode(makePrompt(typed.toLowerCase().replace(/\s+/g, ""), partner), { add_special_tokens: false }).ids;
        const state = { T: 0, past: {} };
        for (const n of this.pasts) state.past[n] = new ort.Tensor("float16", new Uint16Array(0), [1, this.kvh, 0, this.hd]);
        const f16 = f16table();
        const fin = new Map();
        try {
            let out = await this.run([ids], [0], state);
            let beams = [{ score: 0, ids: [], p: 0, prev: 0 }];
            for (let step = 0; step < maxSteps; step++) {
                if (!live()) throw new Error("superseded");
                const cands = [];
                for (let bi = 0; bi < beams.length; bi++) {
                    const b = beams[bi];
                    const off = (bi * out.So + (out.So - 1)) * out.Vo;
                    const lse = out.lse[bi * out.So + out.So - 1];
                    if (b.p === Q.length && b.ids.length) {
                        const k = b.ids.join(",");
                        if (!fin.has(k)) fin.set(k, b.score + f16[out.u16[off + this.eos]] - lse);
                    }
                    const T = b.prev === 1 ? this.tables.L : this.tables.N;
                    const keys = [...(T.byFirst.get("") || [])];
                    if (b.p < Q.length) {
                        const heads = new Set([Q[b.p]]);
                        if (nearCost != null) for (const n of neighbours(Q[b.p])) heads.add(n);
                        for (const h of heads) keys.push(...(T.byFirst.get(h) || []));
                    }
                    if (punct === "free" || b.p >= Q.length) {
                        const qh = b.p < Q.length ? Q[b.p] : "";
                        for (const [h, ks] of T.byFirst) if (h && !isAlnum(h) && h !== qh) keys.push(...ks);
                    }
                    // The best BEAM continuations of this hypothesis.
                    const top = [];
                    let minI = -1;
                    let minV = Infinity;
                    for (const g of keys) {
                        const r = advance(this.groups[g], Q, b.p, punct, nearCost);
                        if (!r) continue;
                        const arr = T.arrs.get(g);
                        const base = b.score - r[1] - lse;
                        for (let j = 0; j < arr.length; j++) {
                            const t = arr[j];
                            const s = base + f16[out.u16[off + t]];
                            if (top.length < BEAM) {
                                top.push([s, t, r[0]]);
                                if (top.length === BEAM) {
                                    minV = Infinity;
                                    for (let q = 0; q < BEAM; q++) if (top[q][0] < minV) { minV = top[q][0]; minI = q; }
                                }
                            } else if (s > minV) {
                                top[minI] = [s, t, r[0]];
                                minV = Infinity;
                                for (let q = 0; q < BEAM; q++) if (top[q][0] < minV) { minV = top[q][0]; minI = q; }
                            }
                        }
                    }
                    for (const [s, t, np] of top) cands.push([s, bi, t, np]);
                }
                if (!cands.length) break;
                cands.sort((a, b) => b[0] - a[0]);
                cands.length = Math.min(cands.length, BEAM);
                if (fin.size >= OFFERED) {
                    const kth = [...fin.values()].sort((a, b) => b - a)[OFFERED - 1];
                    if (cands[0][0] < kth) break;
                }
                beams = cands.map(([s, bi, t, np]) => ({ score: s, ids: [...beams[bi].ids, t], p: np, prev: this.last[t] }));
                out = await this.run(cands.map((c) => [c[2]]), cands.map((c) => c[1]), state);
            }
        } finally {
            for (const t of Object.values(state.past)) t.dispose?.();
        }
        // Hypotheses that read the same (case, spacing, final punctuation
        // aside) are one sentence, with their probabilities added.
        const merged = new Map();
        for (const [k, s] of fin) {
            const text = this.tok.decode(k.split(",").map(Number), { skip_special_tokens: true }).trim();
            const nk = normMatch(text);
            if (!nk) continue;
            const m = merged.get(nk);
            if (m) {
                const hi = Math.max(m.lse, s);
                const lo = Math.min(m.lse, s);
                m.lse = hi + Math.log1p(Math.exp(lo - hi));
                if (s > m.best) {
                    m.best = s;
                    m.text = text;
                }
            } else merged.set(nk, { lse: s, best: s, text });
        }
        const sentences = [...merged.values()].sort((a, b) => b.lse - a.lse).slice(0, OFFERED).map((m) => m.text);
        return { sentences, ms: Math.round(performance.now() - t0) };
    }
}

// ─── Messages ──────────────────────────────────────────────────────────────

const HANDLERS = { init, load, check, expand, unload, cancel: async () => ({}) };

// One thing at a time, in the order it came; an expansion (or a cancel, or
// an unload) supersedes those before it as soon as it arrives.
let queue = Promise.resolve();

self.onmessage = (event) => {
    const msg = event.data || {};
    if (msg.type === "expand" || msg.type === "cancel" || msg.type === "unload") latest = msg.id;
    const handler = HANDLERS[msg.type];
    queue = queue.then(async () => {
        if (!handler) {
            self.postMessage({ id: msg.id, ok: false, error: `unknown message ${msg.type}` });
            return;
        }
        try {
            const reply = await handler(msg);
            self.postMessage({ id: msg.id, ok: true, ...reply });
        } catch (e) {
            self.postMessage({ id: msg.id, ok: false, error: describe(e) });
        }
    });
};