// The listener's worker for the AAC reference app: on-device speech
// recognition that turns what the person's partner says into text. NVIDIA
// Nemotron 3.5 ASR Streaming 0.6B writes the words down; Silero VAD finds
// where speech starts and ends. Both run on onnxruntime-web's WebGPU
// execution provider. Audio in, text out.
//
// Nothing leaves the device, and nothing is kept: the audio arrives in small
// blocks while the person has asked to listen, is turned into text 320 ms at
// a time on the graphics chip, and is dropped. The worker never fetches: the
// runtime (onnxruntime-web's bundle and its WebAssembly) and the model's
// graphs and weights 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. Every session is created with the processor fallback
// switched off (session.disable_cpu_ep_fallback): a graph a device cannot run
// wholly on its graphics chip fails to load, rather than running on the
// processor.
//
// 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 all
// three sessions. Unloading releases the sessions 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-02.)
//
// A module worker, started by src/bindings/listen.js from a Blob URL. It
// imports onnxruntime-web's JSPI build (JavaScript Promise Integration: the
// listener is offered only in browsers that have it) from a Blob URL too, and
// hands it the WebAssembly's bytes, so the runtime never looks for a file by
// name. JSPI lets the runtime read the model's weights one tensor at a time,
// so the whole model is never held in memory. One thread: no model compute
// on the processor.
//
// The graphs and the streaming contract are the asr-proto prototype's
// (tools/asr/README.md): each 320 ms window (5120 samples at 16 kHz, with
// the 201 before it and 40 after it) is one run of the encoder (log-mel front
// end and cache-aware encoder; its caches stay on the GPU), then greedy RNN-T
// decoding with one fused step graph (prediction net, joint, ArgMax for the
// window's four frames) that reads back only token ids.
//
// Messages in (each with an `id`, answered in order with {id, ok, ...} or
// {id, ok: false, error}):
//   {type: "init", bundle, wasm}        model-store URLs of the runtime
//   {type: "load", key, model}          model: {vad, encoder, step: {url,
//                                       data: [{name, url}]}, vocab: url}
//   {type: "check", slot}               two silent windows through every graph
//   {type: "start", slot}               listen for one utterance (slot: the
//                                       language's prompt slot)
//   {type: "audio", samples}            16 kHz mono Float32Array (unanswered)
//   {type: "finish"}                    end it now -> {text}
//   {type: "stop"}                      drop it, unheard
//   {type: "unload"}                    releases the model and destroys the
//                                       device; the worker can then be
//                                       ended, or load again
// Events out: {event: "speech"} someone started speaking; {event: "words",
// text} what is heard so far; {event: "quiet"} a sound with no words in it
// ended (a cough, a door): still listening, as if nothing had been heard;
// {event: "final", text} they stopped (the words in the end; listening is
// over); {event: "error", text}.
//
// Vendored into vendor/onnxruntime-web-<version>-<revision>/ as asr-worker.js
// by vendor.sh, unchanged. Its bytes are pinned in
// crates/aac-web/src/listen.rs (RUNTIME_FILES): change this file only
// together with a new ORT_WEB_VENDOR_REV and its new size and hash there.

const CACHE_NAME = "aac-models-v1";

// The model (soniqo's export of Nemotron 3.5 ASR, as edited by asr-proto).
const BLANK = 13087;
const MAX_SYMBOLS = 10; // tokens per frame at most (greedy RNN-T)
const WIN = 5120; // 320 ms
const LEFT = 201; // samples before a window the front end reads
const IN = 5361; // samples one encoder run reads
const PREEMPH = 0.97;

// When speech starts and ends (asr-proto, Silero VAD v6.2.3, one call per
// window of 10 frames of 512 samples): it starts at a frame of probability
// 0.5 or more, and the two windows before it go in first (640 ms, more than
// the 560 ms the model looks back); it ends after 1 s of frames below 0.35.
// (500 ms, the prototype's rule, cut people off at a pause between two
// sentences: tools/asr/README.md.)
const ON = 0.5;
const OFF = 0.35;
const SILENCE_MS = 1000;
const FRAME_MS = 32;
const PRE_ROLL = 2;
// One utterance is at most this long: the field it goes in holds 300
// characters.
const MAX_SPEECH_MS = 30000;

const ENC_OUT = ["encoded_output", "encoded_length", "new_pre_cache", "new_cache_last_channel",
    "new_cache_last_time", "new_cache_last_channel_len"];

let ort = null;
let S = null; // {vad, enc, step} sessions
let device = null; // the GPU device they run on, this worker's own
let key = "";
let vocab = [];
let special = new Set();
let listening = null; // the Utterance being heard
let generation = 0; // counts utterances: one that ended says nothing more

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(`the listener is not on this device (${url.split("/").pop()})`);
    return res;
}

async function init(msg) {
    if (ort) return {};
    if (typeof caches === "undefined") throw new Error("this browser keeps no model store in workers");
    const cache = await caches.open(CACHE_NAME);
    const [code, wasm] = await Promise.all([
        stored(cache, msg.bundle).then((r) => r.blob()),
        stored(cache, msg.wasm).then((r) => r.arrayBuffer()),
    ]);
    const url = URL.createObjectURL(new Blob([code], { type: "text/javascript" }));
    let runtime;
    try {
        runtime = await import(url);
    } finally {
        URL.revokeObjectURL(url);
    }
    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;
    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.
class StoredPart extends Blob {
    constructor(cache, url, bytes) {
        super([]);
        this.cache = cache;
        this.url = url;
        this.bytes = bytes;
    }

    get size() {
        return this.bytes;
    }

    slice(start, end) {
        const { cache, url } = this;
        return {
            stream: () => new ReadableStream({
                async start(controller) {
                    const blob = await partBlob(cache, url);
                    controller.enqueue(new Uint8Array(await blob.slice(start, end).arrayBuffer()));
                    controller.close();
                },
            }),
        };
    }
}

async function load(msg) {
    if (S && key === msg.key) return {};
    await unload();
    const cache = await caches.open(CACHE_NAME);
    const m = msg.model;
    const graph = async (g) => new Uint8Array(await (await stored(cache, g.url)).arrayBuffer());
    // Each tensor is read from the model store when the runtime needs it.
    const external = (g) => (g.data || []).map((d) => ({ path: d.name, data: new StoredPart(cache, d.url, d.bytes) }));
    const options = (outputs) => ({
        // Options reach the WebGPU provider only when it is named as an
        // object; the device is this worker's own.
        executionProviders: [{ name: "webgpu", device }],
        graphOptimizationLevel: "all",
        preferredOutputLocation: outputs,
        logSeverityLevel: 3,
        // Graphics chip or nothing: no node may fall back to the processor.
        extra: { session: { disable_cpu_ep_fallback: "1" } },
    });
    const sessions = {};
    try {
        device = await ownDevice();
        // The VAD first: onnxruntime-web 1.30 can break a session created
        // before another one with other output locations (asr-proto, step 7).
        sessions.vad = await ort.InferenceSession.create(await graph(m.vad),
            options({ speech_probs: "cpu", hn: "gpu-buffer", cn: "gpu-buffer" }));
        sessions.enc = await ort.InferenceSession.create(await graph(m.encoder), {
            ...options(Object.fromEntries(ENC_OUT.map((n) => [n, "gpu-buffer"]))),
            externalData: external(m.encoder),
        });
        sessions.step = await ort.InferenceSession.create(await graph(m.step), {
            ...options({ token_out: "cpu", h_out: "gpu-buffer", c_out: "gpu-buffer" }),
            externalData: external(m.step),
        });
        const pieces = await (await stored(cache, m.vocab)).json();
        vocab = [];
        for (const [i, p] of Object.entries(pieces)) vocab[Number(i)] = p;
        // <unk> and the language tags (<en-US> ...) the model writes after a
        // sentence's last punctuation even when told the language.
        special = new Set(vocab.flatMap((p, i) => (p === "<unk>" || /^<[a-z]{2}-[A-Z]{2}>$/.test(p) ? [i] : [])));
    } catch (e) {
        await release(sessions);
        throw e;
    } finally {
        held = { url: "", blob: null };
    }
    S = sessions;
    key = msg.key;
    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 the listener");
    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 });
}

// The sessions, then the device: all of its graphics memory goes back.
async function release(sessions) {
    for (const s of Object.values(sessions || {})) await s.release().catch(() => {});
    if (device) device.destroy();
    device = null;
}

async function unload() {
    generation++;
    if (listening) listening.close();
    listening = null;
    const old = S;
    S = null;
    key = "";
    await release(old);
    return {};
}

function detok(ids) {
    let s = "";
    for (const i of ids) if (!special.has(i) && vocab[i]) s += vocab[i];
    return s.replaceAll("▁", " ").replace(/\s+/g, " ").trim();
}

// ─── One stream of speech ──────────────────────────────────────────────────

// The encoder's caches and the prediction net's state, carried from window to
// window on the GPU (only the first window uploads zeros).
class Stream {
    constructor(slot) {
        this.mask = new Float32Array(128);
        this.mask[slot] = 1;
        this.langMask = new ort.Tensor("float32", this.mask, [1, 128]);
        this.state = {
            pre_cache: new ort.Tensor("float32", new Float32Array(128 * 9), [1, 128, 9]),
            cache_last_channel: new ort.Tensor("float32", new Float32Array(24 * 56 * 1024), [24, 1, 56, 1024]),
            cache_last_time: new ort.Tensor("float32", new Float32Array(24 * 1024 * 8), [24, 1, 1024, 8]),
            cache_last_channel_len: new ort.Tensor("int32", new Int32Array(1), [1]),
        };
        this.stateGpu = false;
        this.h = new ort.Tensor("float16", new Uint16Array(1280), [2, 1, 640]);
        this.c = new ort.Tensor("float16", new Uint16Array(1280), [2, 1, 640]);
        this.hGpu = false;
        this.last = BLANK;
        this.ids = [];
        // Samples [5120k - 201, ...) for the next window k; it starts with
        // the 201 (silent) samples before the stream.
        this.buf = new Float32Array(1 << 16);
        this.len = LEFT;
        this.k = 0;
        this.total = 0;
        this.lastSample = 0;
    }

    append(x) {
        if (this.len + x.length > this.buf.length) {
            const nb = new Float32Array(Math.max(this.buf.length * 2, this.len + x.length));
            nb.set(this.buf.subarray(0, this.len));
            this.buf = nb;
        }
        this.buf.set(x, this.len);
        this.len += x.length;
    }

    // Takes 16 kHz samples; decodes every window whose samples (and the 40
    // after it) are in.
    async push(x) {
        if (x.length) {
            this.append(x);
            this.total += x.length;
            this.lastSample = x[x.length - 1];
        }
        while (this.len >= IN) await this.window();
    }

    // The end: the open window is padded so its pre-emphasis is exactly zero
    // (x[L-1]·0.97^i), decoded, and the text returned.
    async end() {
        const windows = Math.max(1, Math.ceil(this.total / WIN));
        while (this.k < windows) {
            const need = IN - this.len;
            if (need > 0) {
                const pad = new Float32Array(need);
                let v = this.lastSample;
                for (let i = 0; i < need; i++) {
                    v *= PREEMPH;
                    pad[i] = v;
                }
                this.lastSample = v;
                this.append(pad);
            }
            await this.window();
        }
        this.dispose();
        return detok(this.ids);
    }

    async window() {
        const w = this.buf.slice(0, IN);
        this.buf.copyWithin(0, WIN, this.len);
        this.len -= WIN;
        this.k++;
        const eo = await S.enc.run({ audio: new ort.Tensor("float32", w, [1, 1, IN]), language_mask: this.langMask, ...this.state });
        eo.encoded_length.dispose();
        if (this.stateGpu) for (const v of Object.values(this.state)) v.dispose();
        this.state = {
            pre_cache: eo.new_pre_cache,
            cache_last_channel: eo.new_cache_last_channel,
            cache_last_time: eo.new_cache_last_time,
            cache_last_channel_len: eo.new_cache_last_channel_len,
        };
        this.stateGpu = true;
        // Greedy RNN-T over the window's four frames. The step scores all four
        // with the current prediction state: frames before the first one that
        // is not blank are blank under it, so a window costs 1 + tokens runs.
        try {
            let t = 0;
            let sym = 0;
            while (t < 4) {
                const r = await S.step.run({
                    encoder_out: eo.encoded_output,
                    token: new ort.Tensor("int32", Int32Array.of(this.last), [1, 1]),
                    h: this.h,
                    c: this.c,
                });
                const toks = r.token_out.data;
                let f = -1;
                for (let i = t; i < 4; i++) {
                    if (Number(toks[i]) !== BLANK) {
                        f = i;
                        break;
                    }
                }
                if (f < 0) {
                    r.h_out.dispose();
                    r.c_out.dispose();
                    break;
                }
                if (f !== t) {
                    t = f;
                    sym = 0;
                }
                const tok = Number(toks[f]);
                this.ids.push(tok);
                this.last = tok;
                if (this.hGpu) {
                    this.h.dispose();
                    this.c.dispose();
                }
                this.h = r.h_out;
                this.c = r.c_out;
                this.hGpu = true;
                if (++sym >= MAX_SYMBOLS) {
                    t++;
                    sym = 0;
                }
            }
        } finally {
            eo.encoded_output.dispose();
        }
    }

    dispose() {
        if (this.stateGpu) for (const v of Object.values(this.state)) v.dispose();
        this.stateGpu = false;
        if (this.hGpu) {
            this.h.dispose();
            this.c.dispose();
        }
        this.hGpu = false;
        // No sound is kept once the stream ends.
        this.buf = new Float32Array(0);
        this.len = 0;
    }
}

// ─── One utterance: the VAD decides when the stream runs ───────────────────

class Utterance {
    constructor(slot, emit) {
        this.slot = slot;
        this.emit = emit;
        this.h = new ort.Tensor("float32", new Float32Array(128), [1, 1, 128]);
        this.c = new ort.Tensor("float32", new Float32Array(128), [1, 1, 128]);
        this.hGpu = false;
        this.ctx = new Float32Array(64); // the 64 samples before the VAD frame
        this.pending = new Float32Array(0); // samples not yet in a whole window
        this.history = []; // the last PRE_ROLL windows, for the pre-roll
        this.stream = null;
        this.silence = 0;
        this.speech = 0;
        this.words = "";
        this.done = false;
    }

    async vad(w) {
        const inp = new Float32Array(10 * 576);
        for (let f = 0; f < 10; f++) {
            inp.set(f === 0 ? this.ctx : w.subarray(f * 512 - 64, f * 512), f * 576);
            inp.set(w.subarray(f * 512, (f + 1) * 512), f * 576 + 64);
        }
        this.ctx = w.slice(WIN - 64);
        const r = await S.vad.run({ input: new ort.Tensor("float32", inp, [10, 576]), h: this.h, c: this.c });
        if (this.hGpu) {
            this.h.dispose();
            this.c.dispose();
        }
        this.h = r.hn;
        this.c = r.cn;
        this.hGpu = true;
        return Array.from(r.speech_probs.data);
    }

    async push(x) {
        if (this.done) return;
        if (this.stream) await this.stream.push(x); // the stream takes audio as it comes
        const p = new Float32Array(this.pending.length + x.length);
        p.set(this.pending);
        p.set(x, this.pending.length);
        this.pending = p;
        while (this.pending.length >= WIN && !this.done) {
            const w = this.pending.slice(0, WIN);
            this.pending = this.pending.slice(WIN);
            await this.window(w);
        }
    }

    async window(w) {
        const probs = await this.vad(w);
        if (!this.stream) {
            if (probs.some((q) => q >= ON)) {
                // Speech: a stream from the pre-roll, this window, and what
                // came after it already.
                this.stream = new Stream(this.slot);
                this.silence = 0;
                this.speech = 0;
                this.emit({ event: "speech" });
                for (const hw of this.history) await this.stream.push(hw);
                await this.stream.push(w);
                if (this.pending.length) await this.stream.push(this.pending);
            }
        } else {
            const words = detok(this.stream.ids);
            if (words !== this.words) {
                this.words = words;
                this.emit({ event: "words", text: words });
            }
        }
        if (this.stream) {
            for (const q of probs) this.silence = q < OFF ? this.silence + FRAME_MS : 0;
            this.speech += WIN / 16;
            if (this.silence >= SILENCE_MS || this.speech >= MAX_SPEECH_MS) {
                const text = await this.endStream();
                if (text) {
                    this.close();
                    this.emit({ event: "final", text });
                    return;
                }
                // A sound with no words in it (a cough, a door) is not the
                // end: go on listening, as if nothing had been heard.
                this.emit({ event: "quiet" });
            }
        }
        this.history.push(w);
        if (this.history.length > PRE_ROLL) this.history.shift();
    }

    async endStream() {
        const s = this.stream;
        this.stream = null;
        this.words = "";
        return s ? s.end() : "";
    }

    // What was heard, now.
    async finish() {
        const text = this.done ? "" : await this.endStream();
        this.close();
        return text;
    }

    close() {
        this.done = true;
        if (this.stream) {
            this.stream.dispose();
            this.stream = null;
        }
        if (this.hGpu) {
            this.h.dispose();
            this.c.dispose();
        }
        this.hGpu = false;
        // No sound is kept once listening ends.
        this.pending = new Float32Array(0);
        this.history = [];
        this.ctx = new Float32Array(64);
    }
}

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

function emitter(gen) {
    return (event) => {
        if (gen === generation) self.postMessage(event);
    };
}

function loaded() {
    if (!S) throw new Error("the listener is not loaded");
}

async function check(msg) {
    loaded();
    // Every graph once, on two windows of silence: the first run compiles
    // the graphics chip's programs, and a device that cannot run them says
    // so now.
    const u = new Utterance(msg.slot, () => {});
    await u.push(new Float32Array(WIN));
    u.close();
    const s = new Stream(msg.slot);
    await s.push(new Float32Array(2 * WIN));
    await s.end();
    return {};
}

async function start(msg) {
    loaded();
    if (listening) listening.close();
    generation++;
    listening = new Utterance(msg.slot, emitter(generation));
    return {};
}

async function finish() {
    const u = listening;
    listening = null;
    generation++;
    return { text: u ? await u.finish() : "" };
}

async function stop() {
    if (listening) listening.close();
    listening = null;
    generation++;
    return {};
}

const HANDLERS = { init, load, check, start, finish, stop, unload };

// One thing at a time, in the order it came: audio, and the requests between
// it (audio sent before a stop is heard before it).
let queue = Promise.resolve();

self.onmessage = (event) => {
    const msg = event.data || {};
    if (msg.type === "audio") {
        queue = queue.then(async () => {
            const u = listening;
            if (!u || !S) return;
            try {
                await u.push(msg.samples);
            } catch (e) {
                if (listening === u) {
                    u.close();
                    listening = null;
                    generation++;
                    self.postMessage({ event: "error", text: describe(e) });
                }
            }
        });
        return;
    }
    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) });
        }
    });
};