// Benchmark worker: runs one backend over the test sentences and reports timings.
import { loadAssets } from './assets.ts';
import type { LoadedAssets } from './assets.ts';
import { parseLexicon } from './frontend.ts';
import type { Voice } from './frontend.ts';
import { PiVoEngine } from './engine.ts';
import type { Backend, RenderStats } from './engine.ts';
import { createLiteBackend, createOrtBackend } from './backends.ts';
import type { OrtModule } from './backends.ts';

export interface BenchOptions {
  ortModule?: string;
  ortWasmPaths?: string;
  threads?: number;
}
export interface BenchSentence {
  text: string;
  voice: Voice;
  stem: string;
}
export interface BenchRequest {
  kind: string;
  opts: BenchOptions;
  sentences: BenchSentence[];
  reps: number;
  refBase: string | null;
  baseUrl: string;
}
export interface BenchRow {
  stem: string;
  seconds: number;
  planMs: number;
  totalMs: number;
  firstMs: number;
  stats: RenderStats | null;
  parity: { corr: number; snr: number; sameLength: boolean } | null;
}
export type BenchResponse =
  | { ok: true; kind: string; name: string; initMs: number; rows: BenchRow[]; crossOriginIsolated: boolean; threads: number }
  | { ok: false; kind?: string; error: string };

const scope = self as unknown as {
  postMessage(message: BenchResponse): void;
  onmessage: ((event: MessageEvent<BenchRequest>) => void) | null;
  crossOriginIsolated: boolean;
};

let assets: LoadedAssets | null = null;

async function readWav(url: string): Promise<Float32Array | null> {
  const r = await fetch(url);
  if (!r.ok) return null;
  const v = new DataView(await r.arrayBuffer());
  const n = (v.byteLength - 44) / 2;
  const out = new Float32Array(n);
  for (let i = 0; i < n; i++) out[i] = v.getInt16(44 + i * 2, true) / 32768;
  return out;
}

function snr(a: Float32Array, ref: Float32Array) {
  const n = Math.min(a.length, ref.length);
  let dot = 0;
  let aa = 0;
  let bb = 0;
  let err = 0;
  for (let i = 0; i < n; i++) {
    // compare after the same PCM16 quantization as the reference WAV
    const q = Math.round(Math.max(-1, Math.min(1, a[i])) * 32767) / 32768;
    dot += q * ref[i];
    aa += q * q;
    bb += ref[i] * ref[i];
    err += (q - ref[i]) ** 2;
  }
  return { corr: dot / Math.sqrt(aa * bb), snr: 10 * Math.log10(bb / Math.max(err, 1e-30)), sameLength: a.length === ref.length };
}

async function makeBackend(kind: string, opts: BenchOptions, loaded: LoadedAssets): Promise<Backend<never>> {
  if (kind.startsWith('ort-')) {
    const ort = (await import(/* @vite-ignore */ opts.ortModule as string)) as OrtModule;
    ort.env.wasm.wasmPaths = opts.ortWasmPaths;
    ort.env.wasm.numThreads = opts.threads ?? 1;
    ort.env.wasm.simd = true;
    return (await createOrtBackend(ort, loaded.buffers, { ep: kind === 'ort-webgpu' ? 'webgpu' : 'wasm' })) as unknown as Backend<never>;
  }
  return createLiteBackend(loaded.buffers, { vocoder: kind.replace('lite-', '') as 'wasm' | 'js' | 'interpreted' }) as unknown as Backend<never>;
}

scope.onmessage = async (e) => {
  const { kind, opts, sentences, reps, refBase, baseUrl } = e.data;
  try {
    assets ??= await loadAssets(baseUrl);
    const t0 = performance.now();
    const backend = await makeBackend(kind, opts, assets);
    const initMs = performance.now() - t0;
    const engine = new PiVoEngine(backend, parseLexicon(new TextDecoder().decode(assets.buffers.cmudict)));
    // warm-up (not timed)
    await engine.render(await engine.plan(sentences[0].text, { voice: 'Pi' }));
    const rows: BenchRow[] = [];
    for (const s of sentences) {
      for (let r = 0; r < reps; r++) {
        const tStart = performance.now();
        const plan = await engine.plan(s.text, { voice: s.voice });
        const tPlan = performance.now() - tStart;
        let first: number | null = null;
        let stats: RenderStats | null = null;
        const audio = await engine.render(plan, {
          onFrames: () => {
            first ??= performance.now() - tStart;
          },
          onStats: (x) => (stats = x),
        });
        const total = performance.now() - tStart;
        let parity: BenchRow['parity'] = null;
        if (refBase && r === 0) {
          const ref = await readWav(`${refBase}/${s.stem}.wav`).catch(() => null);
          if (ref) parity = snr(audio, ref);
        }
        rows.push({ stem: s.stem, seconds: audio.length / 16000, planMs: tPlan, totalMs: total, firstMs: first ?? total, stats, parity });
      }
    }
    scope.postMessage({ ok: true, kind, name: backend.name, initMs, rows, crossOriginIsolated: scope.crossOriginIsolated, threads: opts.threads ?? 1 });
    backend.release?.();
  } catch (error) {
    scope.postMessage({ ok: false, kind, error: String((error as Error)?.stack || error) });
  }
};