/**
 * Before/after measurement for the long-form pipeline (D-WS3a): wall-clock,
 * WER against tests/audio_files/*.txt, generated tokens and decode calls, for
 * the worker's current chunked path ("old") and long-form configurations.
 *
 *   flock /tmp/whisper-heavy.lock bun scripts/measure-longform.ts \
 *     [--clips en_aiproblem,da_festuge,da_festuge:auto,en_alien,meeting] \
 *     [--configs old,lf,lf-ngram5] [--repeat 1] [--out dir]
 *
 * Runs on onnxruntime-node (CPU), Whisper Large v3 Turbo at its WASM-default
 * dtype (q4) — the same weights and quantization as the browser's WASM
 * engine. "meeting" is the synthetic stand-in (tests/meetingFixture.ts).
 */

import { mkdir } from "node:fs/promises";
import { join } from "node:path";
import { parseArgs } from "node:util";
import { pipeline, type AutomaticSpeechRecognitionPipeline } from "@huggingface/transformers";
import { resolveDefaultDtype } from "../src/lib/models";
import {
  transcribeLongForm,
  type LongFormOptions,
} from "../src/worker/longform";
import { createWhisperDeps, loadSileroVad } from "../src/worker/longform/transformers";
import { decodeClip, normalize, readReference } from "../tests/clips";
import { wordErrorRate } from "../tests/dsp";
import { buildMeetingFixture } from "../tests/meetingFixture";

const MODEL_ID = "onnx-community/whisper-large-v3-turbo";

/** Long-form configurations under test; "old" is the worker's chunked path. */
const CONFIGS: Record<string, Omit<LongFormOptions, "language" | "task"> | "old"> = {
  old: "old",
  lf: {},
  "lf-ngram3": { guard: { noRepeatNgramSize: 3, repetitionPenalty: 1 } },
  "lf-ngram5": { guard: { noRepeatNgramSize: 5, repetitionPenalty: 1 } },
  "lf-rep1.1": { guard: { noRepeatNgramSize: 0, repetitionPenalty: 1.1 } },
  "lf-fallback": { fallback: {} },
  "lf-sil160": { vad: { minSilenceDurationMs: 160 } },
};

const { values } = parseArgs({
  options: {
    clips: { type: "string", default: "en_aiproblem,da_festuge,da_festuge:auto,meeting,en_alien" },
    configs: { type: "string", default: "old,lf" },
    repeat: { type: "string", default: "1" },
    out: { type: "string" },
  },
});

interface Clip {
  name: string;
  audio: Float32Array;
  reference: string;
  /** undefined = the UI's "Auto-detect". */
  language?: string;
}

async function loadClip(spec: string): Promise<Clip> {
  const [base, mode] = spec.split(":");
  if (base === "meeting") {
    const { audio, reference } = await buildMeetingFixture();
    return { name: "meeting (synthetic)", audio, reference, language: "en" };
  }
  const file = (await Array.fromAsync(new Bun.Glob(`${base}.*`).scan(join(import.meta.dir, "../tests/audio_files"))))
    .find((name) => !name.endsWith(".txt"));
  if (!file) throw new Error(`no audio for ${base}`);
  return {
    name: spec,
    audio: await decodeClip(file),
    reference: await readReference(base!),
    ...(mode !== "auto" && { language: base!.split("_")[0] }),
  };
}

const asr = (await pipeline("automatic-speech-recognition", MODEL_ID, {
  dtype: resolveDefaultDtype(MODEL_ID, "wasm"),
})) as AutomaticSpeechRecognitionPipeline;

// Count generated tokens on every generate call, prompt prefix excluded
// (leading special tokens below the first timestamp token).
const model = asr.model as unknown as { generate(args: unknown): Promise<unknown> };
const originalGenerate = model.generate.bind(model);
let generated = 0;
let generateCalls = 0;
model.generate = async (args: unknown) => {
  const out = (await originalGenerate(args)) as { tolist(): bigint[][] };
  const ids = out.tolist()[0]!.map(Number);
  let prefix = 0;
  while (prefix < ids.length && ids[prefix]! >= 50257 && ids[prefix]! <= 50364) prefix += 1;
  generated += ids.length - prefix;
  generateCalls += 1;
  return out;
};

const vad = await loadSileroVad({ localOnly: false });
const whisper = createWhisperDeps(asr, { multilingual: true });

// Warm-up so the first measured run doesn't pay session initialization.
await asr(new Float32Array(16_000), { max_new_tokens: 8, language: "en" });

const rows: string[] = [];
const outDir = values.out;
if (outDir) await mkdir(outDir, { recursive: true });

for (const spec of values.clips!.split(",")) {
  const clip = await loadClip(spec);
  for (let r = 0; r < Number(values.repeat); r++) {
    for (const name of values.configs!.split(",")) {
      const config = CONFIGS[name];
      if (config === undefined) throw new Error(`unknown config ${name}`);
      generated = 0;
      generateCalls = 0;
      const audio = clip.audio.slice(); // the pipeline may keep views; give each run its own
      const started = performance.now();
      let text: string;
      let chunks: unknown;
      let extra = "";
      if (config === "old") {
        const output = (await asr(audio, {
          chunk_length_s: 30,
          stride_length_s: 5,
          return_timestamps: true,
          max_new_tokens: 440,
          ...(clip.language !== undefined && { language: clip.language }),
          task: "transcribe",
        })) as { text: string; chunks: unknown };
        text = output.text;
        chunks = output.chunks;
      } else {
        let vadMs = 0;
        const result = await transcribeLongForm(
          audio,
          {
            speechProbabilities: async (pcm) => {
              const t = performance.now();
              const probs = await vad.speechProbabilities(pcm);
              vadMs = performance.now() - t;
              return probs;
            },
            ...whisper,
          },
          { ...config, ...(clip.language !== undefined && { language: clip.language }), task: "transcribe" },
        );
        text = result.text;
        chunks = result.chunks;
        extra = `vad ${(vadMs / 1000).toFixed(1)} s, speech ${result.speechSeconds.toFixed(0)} s, lang ${result.language}`;
      }
      const seconds = (performance.now() - started) / 1000;
      const wer = wordErrorRate(normalize(clip.reference), normalize(text));
      const row = `| ${clip.name} | ${name} | ${seconds.toFixed(1)} | ${wer.toFixed(3)} | ${generated} | ${generateCalls} | ${extra} |`;
      console.log(row);
      rows.push(row);
      if (outDir) {
        await Bun.write(
          join(outDir, `${spec.replace(":", "-")}.${name}.${r}.json`),
          JSON.stringify({ text, chunks }, null, 1),
        );
      }
    }
  }
}

console.log("\n| clip | config | wall s | WER | gen tokens | decodes | notes |\n|---|---|---|---|---|---|---|");
console.log(rows.join("\n"));
await asr.dispose();
await vad.dispose();