/**
 * Transcript-translation speed on onnxruntime-node (CPU), D-WS4b.
 *
 *   flock /tmp/whisper-heavy.lock bun scripts/bench-translate.ts [batchSizes…]
 *
 * Translates the reference transcripts in tests/audio_files/ with the same
 * Opus-MT loader and translation loop the browser worker uses, once per
 * batch size, and prints wall-clock time, pieces, and whether the output
 * depends on the batch size (padding). Downloads ~113 MB per direction into
 * the Transformers.js cache on first run. Not browser numbers: WASM in a
 * tab is slower than onnxruntime-node.
 */

import { join } from "node:path";
import { buildTranslationUnits, getTranslationModel, translateUnits } from "../src/lib/translation";
import { loadOpusMt } from "../src/worker/opusMt";

const AUDIO = join(import.meta.dir, "../tests/audio_files");
const batchSizes = process.argv.slice(2).map(Number).filter((n) => n > 0);
const sizes = batchSizes.length > 0 ? batchSizes : [1, 8];

for (const [file, source, target] of [
  ["en_alien.txt", "en", "da"],
  ["da_festuge.txt", "da", "en"],
] as const) {
  const model = getTranslationModel(source, target)!;
  const text = await Bun.file(join(AUDIO, file)).text();
  const units = buildTranslationUnits({ text, chunks: [] }).map((unit) => unit.text);
  const loadStart = performance.now();
  const translator = await loadOpusMt(model.id);
  const loadMs = performance.now() - loadStart;
  console.log(`\n${file} (${text.length} chars, ${units.length} lines) ${model.label} — load ${loadMs.toFixed(0)} ms`);
  let baseline: string[] | null = null;
  for (const batchSize of sizes) {
    let pieces = 0;
    const start = performance.now();
    const out = await translateUnits(units, {
      batchSize,
      countTokens: translator.countTokens,
      translateBatch: (batch) => translator.translateBatch(batch),
      onProgress: (_done, total) => (pieces = total),
    });
    const seconds = (performance.now() - start) / 1000;
    const same = baseline === null ? "baseline" : JSON.stringify(out) === JSON.stringify(baseline) ? "identical" : "differs";
    baseline ??= out;
    console.log(
      `  batch ${String(batchSize).padStart(2)}: ${seconds.toFixed(2)} s, ${pieces} pieces, ` +
        `${(text.length / seconds).toFixed(0)} chars/s, output ${same}`,
    );
  }
  console.log(`  sample: ${baseline!.slice(0, 3).join(" | ")}`);
  await translator.dispose();
}