/**
 * Parakeet vs Whisper Large v3 Turbo benchmark (D-WS7a) — CPU only.
 *
 *   flock /tmp/whisper-heavy.lock bun scripts/bench-parakeet.ts [parakeet-node|parakeet-web|turbo ...]
 *
 * Whole-file WER (same normalization as tests/live-real-model.test.ts) and
 * wall-clock per English clip in tests/audio_files. Parakeet runs the
 * worker's own `createParakeetEngine` on onnxruntime-node (CPU EP) and on
 * onnxruntime-web (WASM, the browser engine's baseline); Turbo runs the
 * Transformers.js pipeline with the worker's file-flow options on
 * onnxruntime-node at its WASM-default q4 dtype. Weights come from / go to
 * node_modules/@huggingface/transformers/.cache (Parakeet must already be
 * there — `PARAKEET_REAL=1 bun test tests/parakeet-real-model.test.ts`
 * downloads it). Wall-clock excludes model load.
 */

import { join } from "node:path";
import ffmpegPath from "ffmpeg-static";
import { pipeline } from "@huggingface/transformers";
import { PARAKEET_MODEL_ID, resolveDefaultDtype } from "../src/lib/models";
import { PARAKEET_VOCAB_FILE, parakeetFileName } from "../src/lib/parakeet";
import { createParakeetEngine, parakeetSessionOptions, type OrtLike } from "../src/worker/parakeetEngine";
import { wordErrorRate } from "../tests/dsp";

const ROOT = join(import.meta.dir, "..");
const AUDIO = join(ROOT, "tests/audio_files");
const PARAKEET_DIR = join(ROOT, "node_modules/@huggingface/transformers/.cache", PARAKEET_MODEL_ID);
const CLIPS = ["en_aiproblem.m4a", "en_alien.opus"];
const THREADS = 12;

const normalize = (text: string) =>
  text.toLowerCase().replace(/['’]/g, "").replace(/[.,!?;:"“”‘‚«»()\-–—…]/g, " ").replace(/\s+/g, " ").trim();

async function decode(file: string): Promise<Float32Array> {
  const proc = Bun.spawn(
    [ffmpegPath as unknown as string, "-hide_banner", "-loglevel", "error", "-i", join(AUDIO, file),
      "-ac", "1", "-ar", "16000", "-f", "f32le", "-"],
    { stdout: "pipe" },
  );
  return new Float32Array(await new Response(proc.stdout).arrayBuffer());
}

type Transcribe = (audio: Float32Array) => Promise<string>;

async function parakeet(runtime: "node" | "web"): Promise<Transcribe> {
  let ort: OrtLike;
  let options = parakeetSessionOptions("wasm");
  if (runtime === "web") {
    const web = await import("onnxruntime-web");
    web.env.wasm.numThreads = THREADS;
    ort = web as unknown as OrtLike;
  } else {
    ort = (await import("onnxruntime-node")) as unknown as OrtLike;
    const cpu = { executionProviders: ["cpu"], intraOpNumThreads: THREADS, graphOptimizationLevel: "all" };
    options = { encoder: cpu, decoder: cpu };
  }
  const read = async (file: string) => new Uint8Array(await Bun.file(join(PARAKEET_DIR, file)).arrayBuffer());
  const engine = await createParakeetEngine(
    ort,
    {
      encoder: await read(parakeetFileName("encoder", "q8")),
      decoder: await read(parakeetFileName("decoder", "q8")),
      vocab: await Bun.file(join(PARAKEET_DIR, PARAKEET_VOCAB_FILE)).text(),
    },
    options,
  );
  return async (audio) => (await engine.transcribe(audio)).text;
}

async function turbo(): Promise<Transcribe> {
  const id = "onnx-community/whisper-large-v3-turbo";
  const asr = await pipeline("automatic-speech-recognition", id, { dtype: resolveDefaultDtype(id, "wasm") });
  return async (audio) => {
    const output = await asr(audio, {
      chunk_length_s: 30, stride_length_s: 5, return_timestamps: true, max_new_tokens: 440,
      language: "en", task: "transcribe",
    });
    return ((Array.isArray(output) ? output[0] : output) as { text?: string }).text ?? "";
  };
}

const engines: Record<string, () => Promise<Transcribe>> = {
  "parakeet-node": () => parakeet("node"),
  "parakeet-web": () => parakeet("web"),
  turbo,
};

for (const name of process.argv.slice(2).length ? process.argv.slice(2) : Object.keys(engines)) {
  const transcribe = await engines[name]!();
  for (const clip of CLIPS) {
    const audio = await decode(clip);
    const reference = await Bun.file(join(AUDIO, clip.replace(/\.[^.]+$/, ".txt"))).text();
    const start = performance.now();
    const text = await transcribe(audio);
    const seconds = (performance.now() - start) / 1000;
    const wer = wordErrorRate(normalize(reference), normalize(text));
    console.log(
      `${name}\t${clip}\taudio ${(audio.length / 16000).toFixed(1)} s\twall ${seconds.toFixed(1)} s\tWER ${wer.toFixed(3)}`,
    );
  }
}
process.exit(0);