// AntSpeaker (MECT-B2) speaker embeddings + Silero VAD, fully in the browser.
// Embedding model on WebGPU (onnxruntime-web), Kaldi fbank front end in JS, VAD on WASM.
import { fbank } from "./fbank.js";

const ORT_VERSION = "1.30.0";
const ORT_BASE = `https://cdn.jsdelivr.net/npm/onnxruntime-web@${ORT_VERSION}/dist/`;
const REPO = "https://huggingface.co/mrfakename/AntSpeaker-ONNX/resolve/main/";
const VAD_URL = "https://huggingface.co/onnx-community/silero-vad/resolve/main/onnx/model.onnx";
const CACHE_NAME = "antspeaker-onnx-v1";
const PART = 8 * 1024 * 1024;
const PARALLEL = 6;

let ort = null, enc = null, vad = null, backend = "wasm", precision = "fp32";
const post = (type, data = {}) => self.postMessage({ type, ...data });

async function headSize(url) {
  const r = await fetch(url, { method: "HEAD" });
  const n = Number(r.headers.get("content-length") || 0);
  return n;
}

async function fetchRange(url, start, end, onDelta) {
  for (let attempt = 0; ; attempt++) {
    try {
      const res = await fetch(url, { headers: { Range: `bytes=${start}-${end - 1}` } });
      if (res.status !== 206 && !(res.status === 200 && start === 0)) throw new Error(`HTTP ${res.status}`);
      const out = new Uint8Array(end - start);
      const reader = res.body.getReader();
      let off = 0;
      for (;;) {
        const { done, value } = await reader.read();
        if (done) break;
        const n = Math.min(value.byteLength, out.length - off);
        out.set(value.subarray(0, n), off);
        off += n; onDelta(n);
        if (off >= out.length) { reader.cancel().catch(() => {}); break; }
      }
      if (off !== out.length) throw new Error("short read");
      return out;
    } catch (e) {
      if (attempt >= 3) throw e;
      await new Promise((r) => setTimeout(r, 400 * (attempt + 1)));
    }
  }
}

async function fetchFile(url, size, onBytes) {
  let cache = null;
  try {
    cache = await caches.open(CACHE_NAME);
    const hit = await cache.match(url);
    if (hit) { const b = new Uint8Array(await hit.arrayBuffer()); onBytes(b.byteLength, b.byteLength, true); return b; }
  } catch (_) { cache = null; }
  let got = 0, buf;
  if (size > PART) {
    buf = new Uint8Array(size);
    const parts = [];
    for (let s = 0; s < size; s += PART) parts.push([s, Math.min(size, s + PART)]);
    let next = 0;
    const lane = async () => {
      while (next < parts.length) {
        const [a, b] = parts[next++];
        let pg = 0;
        try { buf.set(await fetchRange(url, a, b, (n) => { got += n; pg += n; onBytes(got, size); }), a); }
        catch (e) { got -= pg; throw e; }
      }
    };
    await Promise.all(Array.from({ length: PARALLEL }, lane));
  } else {
    const res = await fetch(url);
    if (!res.ok) throw new Error(`HTTP ${res.status} for ${url}`);
    buf = new Uint8Array(await res.arrayBuffer());
    onBytes(buf.byteLength, buf.byteLength);
  }
  if (cache) { try { await cache.put(url, new Response(buf)); } catch (_) {} }
  return buf;
}

async function load({ webgpu = true, fp16 = null }) {
  const t0 = performance.now();
  post("status", { text: "Loading runtime" });
  ort = await import(ORT_BASE + "ort.webgpu.bundle.min.mjs");
  ort.env.wasm.wasmPaths = ORT_BASE;
  ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.min(4, navigator.hardwareConcurrency || 1) : 1;

  let gpuOk = false, hasF16 = false;
  if (webgpu && self.navigator.gpu) {
    try {
      const ad = await navigator.gpu.requestAdapter();
      if (ad) { gpuOk = true; hasF16 = ad.features.has("shader-f16"); }
    } catch (_) {}
  }
  const useF16 = fp16 === null ? gpuOk && hasF16 : fp16;
  precision = useF16 ? "fp16" : "fp32";
  const encUrl = REPO + (useF16 ? "onnx/model_fp16.onnx" : "onnx/model.onnx");

  post("status", { text: "Downloading model" });
  const [encSize, vadSize] = await Promise.all([headSize(encUrl).catch(() => 0), headSize(VAD_URL).catch(() => 0)]);
  const prog = { enc: [0, encSize || 1], vad: [0, vadSize || 1] };
  const rep = () => post("progress", { got: prog.enc[0] + prog.vad[0], total: prog.enc[1] + prog.vad[1] });
  const [encBuf, vadBuf] = await Promise.all([
    fetchFile(encUrl, encSize, (g, t) => { prog.enc = [g, t]; rep(); }),
    fetchFile(VAD_URL, vadSize, (g, t) => { prog.vad = [g, t]; rep(); }),
  ]);
  const tDl = performance.now();

  post("status", { text: "Compiling" });
  if (gpuOk) {
    try {
      enc = await ort.InferenceSession.create(encBuf, { executionProviders: ["webgpu"], graphOptimizationLevel: "all" });
      backend = "webgpu";
    } catch (e) {
      post("warn", { text: "WebGPU session failed, falling back to WASM: " + e.message });
      enc = null;
    }
  }
  if (!enc) {
    enc = await ort.InferenceSession.create(encBuf, { executionProviders: ["wasm"] });
    backend = "wasm";
  }
  vad = await ort.InferenceSession.create(vadBuf, { executionProviders: ["wasm"] });

  // warm up the usual shapes (1.5 s windows, batch 1 and 8)
  post("status", { text: "Warming up" });
  await runEmbed([new Float32Array(24000)]);
  const tReady = performance.now();
  post("ready", { backend, precision, downloadMs: tDl - t0, totalMs: tReady - t0 });
}

// clips: Float32Array[] at 16 kHz. Equal frame counts are batched together.
async function runEmbed(clips, batch = 8) {
  const feats = clips.map((c) => fbank(c));
  const out = new Array(clips.length);
  const groups = new Map();
  feats.forEach((f, i) => { if (!groups.has(f.frames)) groups.set(f.frames, []); groups.get(f.frames).push(i); });
  for (const [T, idx] of groups) {
    for (let s = 0; s < idx.length; s += batch) {
      const ids = idx.slice(s, s + batch);
      const B = ids.length;
      const data = new Float32Array(B * T * 80);
      ids.forEach((id, k) => data.set(feats[id].data, k * T * 80));
      const r = await enc.run({ fbank: new ort.Tensor("float32", data, [B, T, 80]) });
      const e = r.embedding.data;
      ids.forEach((id, k) => {
        const v = Float32Array.from(e.subarray(k * 192, (k + 1) * 192));
        let n = 0; for (const x of v) n += x * x; n = Math.sqrt(n) || 1;
        for (let j = 0; j < 192; j++) v[j] /= n;
        out[id] = v;
      });
    }
  }
  return out;
}

async function runVad(pcm) {
  let state = new ort.Tensor("float32", new Float32Array(2 * 128), [2, 1, 128]);
  const sr = new ort.Tensor("int64", BigInt64Array.from([16000n]), []);
  const probs = new Float32Array(Math.floor(pcm.length / 512));
  const inp = new Float32Array(576);
  for (let i = 0; i < probs.length; i++) {
    inp.copyWithin(0, 512, 576); // last 64 samples of previous chunk = context
    if (i === 0) inp.fill(0, 0, 64);
    inp.set(pcm.subarray(i * 512, i * 512 + 512), 64);
    const r = await vad.run({ input: new ort.Tensor("float32", inp.slice(), [1, 576]), state, sr });
    probs[i] = r.output.data[0];
    state = r.stateN;
    if (i % 200 === 0) post("vadProgress", { done: i, total: probs.length });
  }
  return probs;
}

// onnxruntime sessions must not run concurrently: serialise all jobs.
let queue = Promise.resolve();
self.onmessage = (ev) => { queue = queue.then(() => handle(ev)); };
async function handle(ev) {
  const { type, id } = ev.data;
  try {
    if (type === "load") await load(ev.data);
    else if (type === "embed") {
      const t0 = performance.now();
      const embs = await runEmbed(ev.data.clips, ev.data.batch || 8);
      post("result", { id, embs, ms: performance.now() - t0 });
    } else if (type === "vad") {
      const t0 = performance.now();
      const probs = await runVad(ev.data.pcm);
      post("result", { id, probs, ms: performance.now() - t0 });
    }
  } catch (e) {
    post("error", { id, text: e && e.message ? e.message : String(e) });
  }
}