// Runs the original ear.tflite: it is fetched from the source repo, translated
// to an ONNX graph in memory (tflite2onnx.js) and executed by onnxruntime-web
// on WebGPU. LiteRT.js's own WebGPU delegate rejects this graph, hence the detour.
importScripts("tflite2onnx.js");
const ORT_VER = "1.30.0";
const ORT_DIR = `https://cdn.jsdelivr.net/npm/onnxruntime-web@${ORT_VER}/dist/`;
const REPO = "https://huggingface.co/desert-ant-labs/ear/resolve/main/";
const CACHE = "ear-webgpu-v1";

const SR = 16000, N_FFT = 400, HOP = 160, N_FRAMES = 3000, N_SAMPLES = 480000, N_BINS = 201, N_MELS = 80;

let ort = null, sess = null, feSess = null, backend = null, inName = null, outName = null;
let melF = null, melLo = null, melHi = null;

const post = (m, t) => self.postMessage(m, t || []);

async function fetchCached(url, label) {
  let cache = null;
  try { cache = await caches.open(CACHE); } catch {}
  if (cache) {
    const hit = await cache.match(url);
    if (hit) {
      const buf = new Uint8Array(await hit.arrayBuffer());
      post({ type: "progress", file: label, loaded: buf.length, total: buf.length, cached: true });
      return buf;
    }
  }
  const res = await fetch(url);
  if (!res.ok) throw new Error(`${label}: HTTP ${res.status}`);
  const total = +res.headers.get("content-length") || 0;
  const reader = res.body.getReader();
  const chunks = []; let loaded = 0, last = 0;
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;
    chunks.push(value); loaded += value.length;
    const now = performance.now();
    if (now - last > 60) { last = now; post({ type: "progress", file: label, loaded, total }); }
  }
  const buf = new Uint8Array(loaded); let o = 0;
  for (const c of chunks) { buf.set(c, o); o += c.length; }
  post({ type: "progress", file: label, loaded, total: total || loaded });
  if (cache) { try { await cache.put(url, new Response(buf)); } catch {} }
  return buf;
}

// ---------- front end: Whisper log-mel (80 x 3000), per ear_meta.json ----------
const HANN = new Float32Array(N_FFT);
for (let i = 0; i < N_FFT; i++) HANN[i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / N_FFT); // periodic

// Mixed-radix FFT for N=400 (2^4 * 5^2), recursive decimation in time.
const FACT = [5, 5, 2, 2, 2, 2];
const TW_RE = new Float64Array(N_FFT), TW_IM = new Float64Array(N_FFT);
for (let k = 0; k < N_FFT; k++) { TW_RE[k] = Math.cos(-2 * Math.PI * k / N_FFT); TW_IM[k] = Math.sin(-2 * Math.PI * k / N_FFT); }
function fftRec(inR, inI, off, stride, n, outR, outI, outOff, fi) {
  if (n === 1) { outR[outOff] = inR[off]; outI[outOff] = inI[off]; return; }
  const p = FACT[fi], m = n / p;
  for (let q = 0; q < p; q++) fftRec(inR, inI, off + q * stride, stride * p, m, outR, outI, outOff + q * m, fi + 1);
  const step = N_FFT / n; // twiddle stride for size n
  const tr = new Float64Array(p), ti = new Float64Array(p);
  for (let k = 0; k < m; k++) {
    for (let q = 0; q < p; q++) {
      const idx = (q * k * step) % N_FFT;
      const ar = outR[outOff + q * m + k], ai = outI[outOff + q * m + k];
      tr[q] = ar * TW_RE[idx] - ai * TW_IM[idx];
      ti[q] = ar * TW_IM[idx] + ai * TW_RE[idx];
    }
    for (let s = 0; s < p; s++) {
      let sr = 0, si = 0;
      for (let q = 0; q < p; q++) {
        const idx = (q * s * m * step) % N_FFT;
        sr += tr[q] * TW_RE[idx] - ti[q] * TW_IM[idx];
        si += tr[q] * TW_IM[idx] + ti[q] * TW_RE[idx];
      }
      SCR_R[s * m + k] = sr; SCR_I[s * m + k] = si;
    }
  }
  for (let j = 0; j < n; j++) { outR[outOff + j] = SCR_R[j]; outI[outOff + j] = SCR_I[j]; }
}
const SCR_R = new Float64Array(N_FFT), SCR_I = new Float64Array(N_FFT);
const FR = new Float64Array(N_FFT), FI = new Float64Array(N_FFT), ZI = new Float64Array(N_FFT);
const OR = new Float64Array(N_FFT), OI = new Float64Array(N_FFT);
const POW = new Float64Array(N_BINS);

function logMel(audio) {
  // audio: Float32Array <= 30 s at 16 kHz; zero-padded to 30 s, reflect-padded by n_fft/2
  const a = new Float32Array(N_SAMPLES);
  a.set(audio.length > N_SAMPLES ? audio.subarray(0, N_SAMPLES) : audio);
  const pad = N_FFT / 2, L = N_SAMPLES;
  const at = (i) => (i < 0 ? a[-i] : i >= L ? a[2 * L - 2 - i] : a[i]);
  const mel = new Float32Array(N_MELS * N_FRAMES);
  const lastActive = Math.min(N_FRAMES, Math.ceil((audio.length + pad) / HOP) + 1);
  let mx = -Infinity;
  for (let t = 0; t < N_FRAMES; t++) {
    if (t > lastActive) { for (let m = 0; m < N_MELS; m++) mel[m * N_FRAMES + t] = 1e-10; continue; }
    const s0 = t * HOP - pad;
    for (let i = 0; i < N_FFT; i++) FR[i] = at(s0 + i) * HANN[i];
    fftRec(FR, ZI, 0, 1, N_FFT, OR, OI, 0, 0);
    for (let k = 0; k < N_BINS; k++) POW[k] = OR[k] * OR[k] + OI[k] * OI[k];
    for (let m = 0; m < N_MELS; m++) {
      let s = 0; const base = m * N_BINS;
      for (let k = melLo[m]; k < melHi[m]; k++) s += melF[base + k] * POW[k];
      mel[m * N_FRAMES + t] = s;
    }
  }
  for (let i = 0; i < mel.length; i++) { const v = Math.log10(Math.max(mel[i], 1e-10)); mel[i] = v; if (v > mx) mx = v; }
  const floor = mx - 8;
  for (let i = 0; i < mel.length; i++) mel[i] = (Math.max(mel[i], floor) + 4) / 4;
  return mel;
}

function softmax(x) {
  let m = -Infinity; for (const v of x) if (v > m) m = v;
  const e = new Float32Array(x.length); let s = 0;
  for (let i = 0; i < x.length; i++) { e[i] = Math.exp(x[i] - m); s += e[i]; }
  for (let i = 0; i < x.length; i++) e[i] /= s;
  return e;
}

// Framing only (cheap); the DFT, mel projection and log run on the GPU.
const FRAMES = new Float32Array(N_FRAMES * N_FFT);
function frames(audio) {
  const n = Math.min(audio.length, N_SAMPLES), pad = N_FFT / 2, L = N_SAMPLES;
  const at = (i) => { const j = i < 0 ? -i : i >= L ? 2 * L - 2 - i : i; return j < n ? audio[j] : 0; };
  for (let t = 0; t < N_FRAMES; t++) {
    const s0 = t * HOP - pad, o = t * N_FFT;
    if (s0 >= 0 && s0 + N_FFT <= n) FRAMES.set(audio.subarray(s0, s0 + N_FFT), o);
    else for (let i = 0; i < N_FFT; i++) FRAMES[o + i] = at(s0 + i);
  }
  return FRAMES;
}

async function infer(audio) {
  const t0 = performance.now();
  let mel;
  if (feSess) {
    const r = await feSess.run({ frames: new ort.Tensor("float32", frames(audio), [N_FRAMES, N_FFT]) });
    mel = r.logmel.data;
  } else mel = logMel(audio);
  const t1 = performance.now();
  const res = await sess.run({ [inName]: new ort.Tensor("float32", mel, [1, N_MELS, N_FRAMES]) });
  const logits = new Float32Array(res[outName].data);
  const t2 = performance.now();
  return { logits, probs: softmax(logits), melMs: t1 - t0, runMs: t2 - t1 };
}

async function load(opts) {
  const tStart = performance.now();
  const [ortMod, modelBytes, melBytes] = await Promise.all([
    import(ORT_DIR + "ort.webgpu.bundle.min.mjs"),
    fetchCached(REPO + "ear.tflite", "ear.tflite"),
    fetchCached(REPO + "mel_filters.f32", "mel_filters.f32"),
  ]);
  ort = ortMod;
  ort.env.wasm.wasmPaths = ORT_DIR;
  ort.env.wasm.numThreads = 1;
  const tDl = performance.now();
  const dv = new DataView(melBytes.buffer, melBytes.byteOffset, melBytes.byteLength);
  const nm = dv.getUint32(0, true), nb = dv.getUint32(4, true);
  if (nm !== N_MELS || nb !== N_BINS) throw new Error(`unexpected mel table ${nm}x${nb}`);
  melF = new Float32Array(melBytes.slice(8).buffer);
  melLo = new Int32Array(N_MELS); melHi = new Int32Array(N_MELS);
  for (let m = 0; m < N_MELS; m++) {
    let lo = N_BINS, hi = 0;
    for (let k = 0; k < N_BINS; k++) if (melF[m * N_BINS + k] !== 0) { if (k < lo) lo = k; hi = k + 1; }
    melLo[m] = lo; melHi[m] = Math.max(hi, lo);
  }
  post({ type: "stage", msg: "Translating graph…" });
  const conv = self.tflite2onnx.convert(modelBytes);
  inName = conv.input; outName = conv.output;
  const tConv = performance.now();
  post({ type: "stage", msg: "Compiling for GPU…" });
  let adapterInfo = "";
  if (self.navigator.gpu && opts.prefer !== "wasm") {
    try {
      const ad = await navigator.gpu.requestAdapter({ powerPreference: "high-performance" });
      if (ad) {
        adapterInfo = ad.info ? [ad.info.vendor, ad.info.architecture].filter(Boolean).join(" ") : "";
        sess = await ort.InferenceSession.create(conv.bytes, { executionProviders: ["webgpu"], graphOptimizationLevel: "all" });
        backend = "webgpu";
        if (opts.gpuFrontend !== false) {
          try {
            feSess = await ort.InferenceSession.create(self.tflite2onnx.frontend(melF, N_MELS, N_FFT, N_FRAMES), { executionProviders: ["webgpu"] });
          } catch (e) { post({ type: "log", msg: "GPU front end failed, using JS: " + (e && e.message || e) }); feSess = null; }
        }
      }
    } catch (e) {
      post({ type: "log", msg: "WebGPU session failed: " + (e && e.message || e) });
      sess = null;
    }
  }
  if (!sess) { sess = await ort.InferenceSession.create(conv.bytes, { executionProviders: ["wasm"] }); backend = "wasm"; }
  const w = await infer(new Float32Array(SR)); // warm-up: shader compilation
  const tEnd = performance.now();
  post({ type: "ready", backend, adapter: adapterInfo, loadMs: tEnd - tStart, downloadMs: tDl - tStart, convertMs: tConv - tDl, sessionMs: tEnd - tConv, warmMs: w.melMs + w.runMs, graph: conv.stats, gpuFrontend: !!feSess });
}

let busy = Promise.resolve();
self.onmessage = (ev) => {
  const d = ev.data;
  busy = busy.then(async () => {
    try {
      if (d.type === "load") await load(d);
      else if (d.type === "infer") {
        const r = await infer(d.audio);
        post({ type: "result", id: d.id, ...r }, [r.logits.buffer, r.probs.buffer]);
      } else if (d.type === "batch") {
        // d.audio: whole clip; d.segs: [[start,len],...] in samples
        for (let i = 0; i < d.segs.length; i++) {
          const [s, n] = d.segs[i];
          const r = await infer(d.audio.subarray(s, s + n));
          post({ type: "seg", id: d.id, i, ...r });
        }
        post({ type: "batchdone", id: d.id });
      }
    } catch (e) {
      post({ type: "error", id: d.id, msg: String(e && e.message || e) });
    }
  });
};