#!/usr/bin/env node
// Sortformer diarization engine for scripts/diarize.mjs.
//
// The exact counterpart of scripts/diarize-sherpa.py, and deliberately the same
// shape: decode audio, run the model, print raw turns as JSON on stdout. All
// policy (filenames, provenance shape, atomic writes) lives in the .mjs wrapper.
//
// stdout is JSON ONLY. Progress goes to stderr so the wrapper can stream it.
//
// WHAT THIS ADDS OVER THE BINARY: audio decoding. scripts/sortformer/
// diarize-file.cpp only speaks 16 kHz mono PCM, and the corpus holds mp3, m4a,
// webm and whole video containers. ffmpeg does that here, exactly as the sherpa
// engine shells out to ffmpeg for the same reason.
//
// NO TEMPORARY WAV, ON PURPOSE. A decoded 8-hour VOD is ~900 MB, this box's disk
// runs at 98%, and /tmp is frequently a tmpfs — writing one would be the largest
// single risk in this lane. ffmpeg's raw output is piped straight into the
// engine, which consumes it in blocks, so neither process ever holds the file.
//
// Usage (invoked by diarize.mjs; runnable directly for debugging):
//   diarize-sortformer.mjs --bin <diarize-file> --model <gguf> [options] <audio>
//
//   --bin <path>       engine binary            [required]
//   --model <path>     sortformer .gguf         [required]
//   --backend <b>      vulkan | cpu             (default vulkan)
//   --threads <n>      CPU backend threads only (default 6 — see below)
//   --ffmpeg <path>    ffmpeg binary            (FFMPEG_BIN, "ffmpeg")

import { spawn } from "node:child_process";
import path from "node:path";

const SAMPLE_RATE = 16000; // what the model's frontend expects; not a knob

function fail(msg, code = 2) {
  process.stderr.write(`diarize-sortformer: ${msg}\n`);
  process.exit(code);
}

const argv = process.argv.slice(2);
if (argv.includes("-h") || argv.includes("--help")) {
  process.stdout.write(
    "diarize-sortformer.mjs --bin <diarize-file> --model <gguf> [options] <audio>\n",
  );
  process.exit(0);
}

function arg(flag, fallback) {
  const i = argv.indexOf(flag);
  return i < 0 || i === argv.length - 1 ? fallback : argv[i + 1];
}

const positional = [];
for (let i = 0; i < argv.length; i++) {
  if (argv[i].startsWith("-")) {
    i++; // every flag here takes a value
    continue;
  }
  positional.push(argv[i]);
}

const audio = positional[0];
if (!audio) fail("missing <audioFile>");

const bin = arg("--bin", process.env.SORTFORMER_BIN ?? "");
const model = arg("--model", process.env.SORTFORMER_MODEL ?? "");
const backend = String(arg("--backend", "vulkan")).toLowerCase();
// 6, not 4 and not 8: measured on this box at 1305 s/audio-hour against 1514 for
// ggml's default of 4, and 1401 for 8 — which is WORSE than 4 through
// oversubscription on a 4c/8t part. Ignored by the Vulkan backend.
const threads = Number(arg("--threads", "6"));
const ffmpeg = arg("--ffmpeg", process.env.FFMPEG_BIN ?? "ffmpeg");

if (!bin) fail("missing --bin (or SORTFORMER_BIN)");
if (!model) fail("missing --model (or SORTFORMER_MODEL)");

// Settings speak "vulkan"/"cpu"; ggml speaks device names. One translation, here,
// so nothing upstream of this file has to know ggml's spelling.
const ggmlBackend =
  backend === "cpu" ? "CPU" : backend === "vulkan" ? "Vulkan0" : null;
if (!ggmlBackend) fail(`--backend must be 'vulkan' or 'cpu' (got '${backend}')`);

// ffmpeg: anything in -> raw s16le mono 16k on stdout. `-vn` because a source
// container carries a video stream that would otherwise be negotiated first.
const dec = spawn(
  ffmpeg,
  [
    "-nostdin",
    "-v", "error",
    "-i", audio,
    "-vn",
    "-ac", "1",
    "-ar", String(SAMPLE_RATE),
    "-f", "s16le",
    "pipe:1",
  ],
  { stdio: ["ignore", "pipe", "inherit"] },
);

const engine = spawn(
  bin,
  [
    "--model", model,
    "--audio-raw", "-",
    "--sample-rate", String(SAMPLE_RATE),
    "--backend", ggmlBackend,
    "--threads", String(threads),
  ],
  {
    stdio: ["pipe", "pipe", "inherit"],
    env: {
      ...process.env,
      // The build sets RPATH=$ORIGIN when patchelf is available; this covers the
      // case where it was not, so the binary finds its ggml .so files either way.
      LD_LIBRARY_PATH: [path.dirname(path.resolve(bin)), process.env.LD_LIBRARY_PATH]
        .filter(Boolean)
        .join(":"),
    },
  },
);

dec.on("error", (err) => fail(`failed to run ${ffmpeg}: ${err.message}`, 3));
engine.on("error", (err) => fail(`failed to run ${bin}: ${err.message}`, 3));

dec.stdout.pipe(engine.stdin);
// The engine exiting early (bad model, unusable backend) closes the pipe under
// ffmpeg. That is an expected race, not a crash to report.
engine.stdin.on("error", () => dec.kill("SIGTERM"));

let stdout = "";
engine.stdout.setEncoding("utf8");
engine.stdout.on("data", (d) => {
  stdout += d;
});

let decCode = null;
let engineCode = null;

function finish() {
  // The ENGINE's verdict is what gates everything: it is the one that can fail
  // for a reason worth reporting, and waiting on ffmpeg first is what used to
  // hang here forever.
  if (engineCode === null) return;
  if (engineCode !== 0) fail(`engine exited ${engineCode}`, engineCode || 1);
  // Only on the success path does ffmpeg's exit matter — a non-zero decode with
  // a healthy engine means the audio was truncated, which must not pass
  // silently as a short diarization.
  if (decCode === null) return;
  if (decCode !== 0) fail(`ffmpeg exited ${decCode} — audio may be truncated`, decCode || 1);
  if (!stdout.trim()) fail("engine produced no output");
  process.stdout.write(stdout.endsWith("\n") ? stdout : stdout + "\n");
}

dec.on("close", (code) => {
  decCode = code ?? 0;
  finish();
});

engine.on("close", (code) => {
  engineCode = code ?? 0;
  // THE ENGINE IS THE ONLY READER OF ffmpeg's OUTPUT. Once it is gone — and it
  // can go abruptly: a Vulkan device loss aborts it mid-stream — ffmpeg is
  // writing into a pipe nobody will ever drain, so it blocks on a full pipe and
  // never exits, and this process never exits either because it is still
  // waiting for that close. Three of these accumulated during development, each
  // holding a decoder resident for hours. The engine's exit is therefore the
  // signal to tear the decoder down rather than wait for an EOF that cannot come.
  if (dec.exitCode === null && !dec.killed) dec.kill("SIGKILL");
  finish();
});