const fs = require("fs");
const path = require("path");
const { spawn } = require("child_process");

const options = parseArguments(process.argv.slice(2));
if (!options.worker || !options.audio || !options.model) {
  console.error("Usage: node benchmark-worker.js --worker <exe> --audio <16-kHz-mono.wav> --model <ggml.bin> [--max-latency-ms 2160] [--expect-transcript \"...\"]");
  process.exit(2);
}

run().catch((error) => {
  console.error(error.stack || error.message);
  process.exit(1);
});

async function run() {
  const pcm = readPcmWave(options.audio);
  const child = spawn(options.worker, [], {
    stdio: ["pipe", "pipe", "pipe"],
    windowsHide: true,
    env: process.env
  });
  const messages = [];
  let stdout = "";
  let stderr = "";
  child.stderr.on("data", (chunk) => { stderr = `${stderr}${chunk}`.slice(-8000); });
  child.stdout.on("data", (chunk) => {
    stdout += chunk;
    let newline;
    while ((newline = stdout.indexOf("\n")) >= 0) {
      const line = stdout.slice(0, newline).trim();
      stdout = stdout.slice(newline + 1);
      if (!line) continue;
      try { messages.push(JSON.parse(line)); } catch { }
    }
  });
  const exited = new Promise((resolve) => child.once("exit", (code) => resolve(code)));
  await waitFor(messages, (message) => message.type === "ready", 20000, exited, stderr);
  send(child, { type: "load_model", model: { id: "benchmark-model", path: path.resolve(options.model) } });
  const loaded = await waitFor(messages, (message) => ["model_loaded", "model_error"].includes(message.type), 30000, exited, stderr);
  if (loaded.type !== "model_loaded") throw new Error(loaded.message || "Model loading failed.");

  const sessionId = "benchmark-session";
  const trackId = "benchmark-track";
  send(child, { type: "start_session", session: { id: sessionId, mode: "benchmark" } });
  send(child, { type: "add_track", session_id: sessionId, track: { source_uuid: trackId } });
  const frameBytes = 640;
  const started = Date.now();
  let sequence = 0;
  for (let offset = 0; offset < pcm.length; offset += frameBytes) {
    const dueAt = started + sequence * 20;
    const wait = dueAt - Date.now();
    if (wait > 0) await delay(wait);
    const frame = Buffer.alloc(frameBytes);
    pcm.copy(frame, 0, offset, Math.min(pcm.length, offset + frameBytes));
    const capturedAt = Date.now();
    send(child, {
      type: "audio", session_id: sessionId, track_id: trackId, sequence,
      capture_timestamp_us: capturedAt * 1000, captured_at: capturedAt
    }, frame);
    sequence += 1;
  }
  for (let index = 0; index < 50; index += 1) {
    const dueAt = started + sequence * 20;
    const wait = dueAt - Date.now();
    if (wait > 0) await delay(wait);
    const capturedAt = Date.now();
    send(child, {
      type: "audio", session_id: sessionId, track_id: trackId, sequence,
      capture_timestamp_us: capturedAt * 1000, captured_at: capturedAt
    }, Buffer.alloc(frameBytes));
    sequence += 1;
  }
  await waitFor(messages, (message) => message.type === "hypothesis" && message.final, 15000, exited, stderr);
  send(child, { type: "stop_session", session_id: sessionId });
  await waitFor(messages, (message) => message.type === "session_stopped" && message.session_id === sessionId, 5000, exited, stderr);
  send(child, { type: "shutdown" });
  await Promise.race([exited, delay(3000)]);

  const hypotheses = messages.filter((message) => message.type === "hypothesis");
  const selected = hypotheses.findLast((message) => message.final) || hypotheses.at(-1);
  if (!selected) throw new Error("The worker did not produce a hypothesis.");
  const words = mergeFirstSeen(selected.words || [], hypotheses);
  const latencies = words.map((word) => word.latency_ms).filter(Number.isFinite);
  const inference = hypotheses.map((message) => message.inference_ms).filter(Number.isFinite);
  const result = {
    backend: loaded.backend,
    audio_ms: Math.round(pcm.length / 32),
    revisions: hypotheses.length,
    transcript: selected.text,
    word_count: words.length,
    latency_ms: statistics(latencies),
    inference_ms: statistics(inference),
    control_tokens_present: words.some((word) => /\[_(?:BEG_|TT_\d+)\]/i.test(word.text)),
    transcript_matches_expected: options.expectTranscript === undefined ? null
      : comparableTranscript(selected.text) === comparableTranscript(options.expectTranscript)
  };
  console.log(JSON.stringify(result, null, 2));
  if (result.control_tokens_present) throw new Error("Whisper control tokens leaked into word analysis.");
  if (result.transcript_matches_expected === false)
    throw new Error(`Transcript did not match the expected acceptance phrase: ${selected.text}`);
  if (Number.isFinite(options.maxLatencyMs) && result.latency_ms.max > options.maxLatencyMs)
    throw new Error(`Maximum first-seen word latency ${result.latency_ms.max} ms exceeds ${options.maxLatencyMs} ms.`);
}

function mergeFirstSeen(selected, hypotheses) {
  const chosen = withOccurrences(selected);
  const candidates = hypotheses.flatMap((hypothesis) => withOccurrences(Array.isArray(hypothesis.words) ? hypothesis.words : []));
  return chosen.map((word) => {
    const matches = candidates.filter((candidate) => candidate.key === word.key && candidate.occurrence === word.occurrence &&
      Math.abs(Number(candidate.captured_at_ms) - Number(word.captured_at_ms)) <= 1500);
    const { key: _key, occurrence: _occurrence, ...plain } = word;
    return matches.length ? { ...plain, latency_ms: Math.min(...matches.map((candidate) => Number(candidate.latency_ms))) } : plain;
  }).filter((word) => comparable(word.text));
}
function withOccurrences(words) {
  const seen = new Map();
  return words.map((word) => {
    const key = comparable(word.text);
    const occurrence = seen.get(key) || 0;
    seen.set(key, occurrence + 1);
    return { ...word, key, occurrence };
  });
}

function comparable(value) { return String(value || "").toLocaleLowerCase().replace(/[^\p{L}\p{N}]+/gu, ""); }
function comparableTranscript(value) { return String(value || "").toLocaleLowerCase().match(/[\p{L}\p{N}]+/gu)?.join(" ") || ""; }
function statistics(values) {
  if (!values.length) return { count: 0, min: null, average: null, median: null, p99: null, max: null };
  const sorted = [...values].sort((left, right) => left - right);
  return {
    count: sorted.length, min: sorted[0],
    average: sorted.reduce((sum, value) => sum + value, 0) / sorted.length,
    median: percentile(sorted, .5), p99: percentile(sorted, .99), max: sorted.at(-1)
  };
}
function percentile(values, ratio) {
  const position = (values.length - 1) * ratio;
  const lower = Math.floor(position); const upper = Math.ceil(position);
  return values[lower] + (values[upper] - values[lower]) * (position - lower);
}
function send(child, metadata, pcm = Buffer.alloc(0)) {
  const json = Buffer.from(JSON.stringify(metadata));
  const header = Buffer.alloc(8);
  header.writeUInt32LE(json.length, 0);
  header.writeUInt32LE(pcm.length, 4);
  child.stdin.write(Buffer.concat([header, json, pcm]));
}
async function waitFor(messages, predicate, timeoutMs, exited, stderr) {
  const started = Date.now();
  while (Date.now() - started < timeoutMs) {
    const match = messages.find(predicate);
    if (match) return match;
    const exit = await Promise.race([exited, delay(20).then(() => null)]);
    if (exit !== null) throw new Error(`Worker exited with code ${exit}.${stderr ? ` ${stderr.trim()}` : ""}`);
  }
  throw new Error(`Worker response timed out after ${timeoutMs} ms.${stderr ? ` ${stderr.trim()}` : ""}`);
}
function readPcmWave(filename) {
  const body = fs.readFileSync(filename);
  if (body.toString("ascii", 0, 4) !== "RIFF" || body.toString("ascii", 8, 12) !== "WAVE") throw new Error("Audio input is not a WAV file.");
  let offset = 12;
  let format = null;
  let pcm = null;
  while (offset + 8 <= body.length) {
    const type = body.toString("ascii", offset, offset + 4);
    const size = body.readUInt32LE(offset + 4);
    const start = offset + 8;
    if (type === "fmt ") format = {
      codec: body.readUInt16LE(start), channels: body.readUInt16LE(start + 2),
      sampleRate: body.readUInt32LE(start + 4), bits: body.readUInt16LE(start + 14)
    };
    if (type === "data") pcm = body.subarray(start, start + size);
    offset = start + size + (size % 2);
  }
  if (!format || format.codec !== 1 || format.channels !== 1 || format.sampleRate !== 16000 || format.bits !== 16 || !pcm)
    throw new Error("Audio input must be 16-kHz mono PCM16.");
  return pcm;
}
function parseArguments(args) {
  const result = {};
  for (let index = 0; index < args.length; index += 2) {
    const key = args[index]?.replace(/^--/, "").replace(/-([a-z])/g, (_match, letter) => letter.toUpperCase());
    if (key) result[key] = args[index + 1];
  }
  result.maxLatencyMs = result.maxLatencyMs === undefined ? null : Number(result.maxLatencyMs);
  return result;
}
function delay(ms) { return new Promise((resolve) => setTimeout(resolve, ms)); }