#!/usr/bin/env node
// Standalone parakeet.cpp transcription with overlapped splitting + stitching.
//
// parakeet-cli only accepts a single WAV and chokes on very long audio (dr_wav's
// data-chunk size is 32-bit, so >4GB PCM overflows and yields an empty
// transcript). This script works around both: it slices the source into
// OVERLAPPING 16kHz-mono WAV windows (the format parakeet-cli resamples to
// internally anyway), transcribes each with `--json --timestamps`, then stitches
// the per-window word timestamps back into one transcript. The overlap means a
// word clipped at one window boundary is always captured whole by the
// neighbouring window; a per-boundary cut point in the middle of each overlap
// region decides which window owns each word, so nothing is dropped or doubled.
//
// Dual-use, by design:
//   * In the app: common/lib/transcriptionApps.ts registers a `parakeet` app
//     whose binary IS this script. transcribeOne runs it with cwd == the video
//     dir and an explicit output path, then renames + normalizes the result.
//   * On the CLI: run it directly. With no output path it prints JSON to stdout.
//
// Output is chough-native JSON ({ chunk_data: [{ start_time, end_time, text }] },
// seconds) so it flows through the existing normalizeTranscript("chough-json")
// path with zero downstream changes.
//
// Usage:
//   parakeet-stitch.mjs --model <gguf> [options] <audioFile> [outputFile]
//
// Options (env fallback in parens):
//   --model <path>      parakeet .gguf model                  (PARAKEET_MODEL) [required]
//   --cli <path>        parakeet-cli binary                   (PARAKEET_CLI, default "parakeet-cli")
//   --ffmpeg <path>     ffmpeg binary                         (FFMPEG_BIN, default "ffmpeg")
//   --ffprobe <path>    ffprobe binary                        (FFPROBE_BIN, default derived from ffmpeg)
//   --segment <sec>     window length                         (PARAKEET_SEGMENT_SEC, default 480)
//   --overlap <sec>     window overlap                        (PARAKEET_OVERLAP_SEC, default 6)
//   --decoder <ctc|tdt> passed through to parakeet-cli        (PARAKEET_DECODER)
//   --lang <locale>     passed through to parakeet-cli        (PARAKEET_LANG)
//   --device <dev>      compute device, exported to parakeet-cli as the
//                       PARAKEET_DEVICE env var (parakeet.cpp has no --device
//                       flag): "cpu", "CUDA0", "Vulkan1", ...   (PARAKEET_DEVICE)
//
// Resumable, yt-dlp style: each window's raw parakeet-cli JSON is written to a
// per-audio work dir (".<audio>.parakeet/" beside the audio) as it completes,
// and the final transcript is stitched only once ALL windows are done (then the
// work dir is removed). Re-running the same command picks up the cached window
// JSONs and only transcribes what's missing. On SIGTERM/SIGINT it finishes the
// in-flight window, leaves the completed window JSONs in place, and exits WITHOUT
// stitching — so a long run can be paused and resumed later with no lost work.
//   --gap <sec>         start a new cue when the inter-word gap exceeds this (default 0.8)
//   --max-cue <sec>     cap a single cue's duration                          (default 8)
//   -o, --output <f>    output file (alternative to the positional arg)
//   --keep-temp         keep the temp working dir (for debugging)
//   --words             also write the stitched words, as a top-level
//                       `words: [{w, start, end, conf}]` beside chunk_data
//                       (seconds; readers of chunk_data ignore it)
//   -h, --help          show this help

import { execFile } from "node:child_process";
import { promisify } from "node:util";
import { mkdir, mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { existsSync } from "node:fs";
import os from "node:os";
import path from "node:path";

const execFileP = promisify(execFile);

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

function progress(msg) {
  // stderr so it never contaminates stdout JSON, and so the app's progress
  // parser (createParakeetProgressParser) can read it off the merged stream.
  process.stderr.write(`parakeet-stitch: ${msg}\n`);
}

const HELP = `parakeet-stitch.mjs --model <gguf> [options] <audioFile> [outputFile]

Slices audio into overlapping 16kHz-mono windows, transcribes each with
parakeet-cli --json --timestamps, and stitches the word timestamps into one
chough-native JSON transcript. With no outputFile, prints JSON to stdout.

See the header of this file for the full option list.`;

function parseArgs(argv) {
  const opts = {
    model: process.env.PARAKEET_MODEL ?? "",
    cli: process.env.PARAKEET_CLI ?? "parakeet-cli",
    ffmpeg: process.env.FFMPEG_BIN ?? "ffmpeg",
    ffprobe: process.env.FFPROBE_BIN ?? "",
    segment: numEnv(process.env.PARAKEET_SEGMENT_SEC, 480),
    overlap: numEnv(process.env.PARAKEET_OVERLAP_SEC, 6),
    decoder: process.env.PARAKEET_DECODER ?? "",
    lang: process.env.PARAKEET_LANG ?? "",
    device: process.env.PARAKEET_DEVICE ?? "",
    gap: 0.8,
    maxCue: 8,
    keepTemp: false,
    words: false,
    output: undefined,
    audio: undefined,
  };
  const positional = [];
  for (let i = 0; i < argv.length; i += 1) {
    const a = argv[i];
    const next = () => {
      const v = argv[++i];
      if (v === undefined) fail(`missing value for ${a}`);
      return v;
    };
    switch (a) {
      case "-h": case "--help": process.stdout.write(HELP + "\n"); process.exit(0);
      case "--model": opts.model = next(); break;
      case "--cli": opts.cli = next(); break;
      case "--ffmpeg": opts.ffmpeg = next(); break;
      case "--ffprobe": opts.ffprobe = next(); break;
      case "--segment": opts.segment = Number(next()); break;
      case "--overlap": opts.overlap = Number(next()); break;
      case "--decoder": opts.decoder = next(); break;
      case "--lang": opts.lang = next(); break;
      case "--device": opts.device = next(); break;
      case "--gap": opts.gap = Number(next()); break;
      case "--max-cue": opts.maxCue = Number(next()); break;
      case "-o": case "--output": opts.output = next(); break;
      case "--keep-temp": opts.keepTemp = true; break;
      case "--words": opts.words = true; break;
      default:
        if (a.startsWith("-")) fail(`unknown option ${a}`);
        positional.push(a);
    }
  }
  opts.audio = positional[0];
  if (opts.output === undefined) opts.output = positional[1];
  return opts;
}

function numEnv(v, dflt) {
  if (v === undefined || v === "") return dflt;
  const n = Number(v);
  return Number.isFinite(n) ? n : dflt;
}

const round3 = (n) => Math.round(n * 1000) / 1000;

// The stitched words as --words writes them: absolute seconds, empty words
// dropped, conf kept only when the engine gave one.
function wordsOut(stitched) {
  return stitched
    .filter((w) => String(w.w ?? "").trim() !== "")
    .map((w) => ({
      w: String(w.w).trim(),
      start: round3(w.start),
      end: round3(w.end),
      ...(typeof w.conf === "number" ? { conf: w.conf } : {}),
    }));
}

// Seconds -> m:ss (or h:mm:ss). Used for the per-video ETA in progress lines.
function formatClock(totalSeconds) {
  const s = Math.max(0, Math.floor(totalSeconds));
  const hh = Math.floor(s / 3600);
  const mm = Math.floor((s % 3600) / 60);
  const ss = s % 60;
  const pad = (n) => String(n).padStart(2, "0");
  return hh > 0 ? `${hh}:${pad(mm)}:${pad(ss)}` : `${mm}:${pad(ss)}`;
}

// Per-video ETA suffix for a segment progress line. `times` holds the measured
// transcription wall-time (ms) of each completed segment; `remaining` is the
// count of segments still to transcribe (the current one included). Returns ""
// until at least one segment has completed (no average to project from yet).
function etaSuffix(times, remaining) {
  if (times.length === 0 || remaining <= 0) return "";
  const avgMs = times.reduce((a, b) => a + b, 0) / times.length;
  const etaSec = (avgMs * remaining) / 1000;
  return ` (avg ${round3(avgMs / 1000)}s/seg, ETA ${formatClock(etaSec)})`;
}

// Locate ffprobe: explicit flag/env wins, else swap a trailing "ffmpeg" in the
// ffmpeg path for "ffprobe" (handles absolute paths), else fall back to PATH.
function resolveFfprobe(opts) {
  if (opts.ffprobe) return opts.ffprobe;
  if (/ffmpeg$/.test(opts.ffmpeg)) return opts.ffmpeg.replace(/ffmpeg$/, "ffprobe");
  return "ffprobe";
}

async function probeDuration(ffprobe, audio) {
  const { stdout } = await execFileP(ffprobe, [
    "-v", "error",
    "-show_entries", "format=duration",
    "-of", "default=noprint_wrappers=1:nokey=1",
    audio,
  ]);
  const dur = Number.parseFloat(stdout.trim());
  if (!Number.isFinite(dur) || dur <= 0) {
    throw new Error(`could not determine duration of ${audio} (got "${stdout.trim()}")`);
  }
  return dur;
}

// Slice [start, start+len) of the source into a 16kHz-mono s16 WAV. Input
// seeking (-ss before -i) is fast and accurate enough given the overlap.
async function sliceWav(ffmpeg, audio, start, len, outWav) {
  const args = ["-nostdin", "-v", "error", "-y"];
  if (start > 0) args.push("-ss", String(start));
  args.push("-i", audio);
  if (len !== undefined) args.push("-t", String(len));
  args.push("-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", "-f", "wav", outWav);
  await execFileP(ffmpeg, args);
}

// Run parakeet-cli on one window and return its parsed JSON document.
async function transcribeWindow(opts, wav) {
  const args = ["transcribe", "--model", opts.model, "--timestamps", "--json", "--input", wav];
  if (opts.decoder) args.push("--decoder", opts.decoder);
  if (opts.lang) args.push("--lang", opts.lang);
  // parakeet.cpp's parakeet-cli has NO --device flag; the compute device is
  // chosen solely via the PARAKEET_DEVICE env var (e.g. "cpu", "CUDA0",
  // "Vulkan1"). Its default is the first GPU the ggml registry reports, so we
  // MUST set the env var to honor a "cpu" (or specific-GPU) request.
  const env = opts.device ? { ...process.env, PARAKEET_DEVICE: opts.device } : process.env;
  const { stdout } = await execFileP(opts.cli, args, { env, maxBuffer: 256 * 1024 * 1024 });
  return parseCliJson(stdout);
}

// parakeet-cli prints the JSON document to stdout; model-load chatter goes to
// stderr. Be defensive anyway and carve out the outermost { ... }.
function parseCliJson(stdout) {
  const s = stdout.indexOf("{");
  const e = stdout.lastIndexOf("}");
  if (s < 0 || e <= s) throw new Error(`no JSON in parakeet-cli output: ${stdout.slice(0, 200)}`);
  return JSON.parse(stdout.slice(s, e + 1));
}

// Group absolute-timestamped words into cues: break on a large inter-word gap or
// when a cue would exceed maxCue seconds.
function groupCues(words, { gap, maxCue }) {
  const cues = [];
  let cur = null;
  for (const w of words) {
    const text = (w.w ?? w.text ?? "").trim();
    if (!text) continue;
    if (cur && (w.start - cur.end > gap || w.end - cur.start > maxCue)) {
      cues.push(cur);
      cur = null;
    }
    if (!cur) cur = { start: w.start, end: w.end, words: [text] };
    else { cur.end = w.end; cur.words.push(text); }
  }
  if (cur) cues.push(cur);
  return cues.map((c) => ({
    start_time: round3(c.start),
    end_time: round3(c.end),
    text: c.words.join(" "),
  }));
}

// Resumable per-audio work dir: ".<audio>.parakeet/" beside the audio file. Holds
// one raw parakeet-cli JSON per completed window plus a meta.json describing the
// windowing, so a re-run reuses finished windows and a pause loses nothing.
function workDirFor(absAudio) {
  return path.join(
    path.dirname(absAudio),
    `.${path.basename(absAudio)}.parakeet`,
  );
}

const PARTIAL_PAUSE_CODE = 0;

async function main() {
  const opts = parseArgs(process.argv.slice(2));
  if (!opts.audio) fail("missing <audioFile>\n\n" + HELP);
  if (!opts.model) fail("missing --model (or PARAKEET_MODEL)");
  if (!existsSync(opts.audio)) fail(`audio file not found: ${opts.audio}`);
  if (!existsSync(opts.model)) fail(`model not found: ${opts.model}`);
  if (!(opts.segment > 0)) fail(`--segment must be > 0 (got ${opts.segment})`);
  if (!(opts.overlap >= 0)) fail(`--overlap must be >= 0 (got ${opts.overlap})`);
  // Overlap must be strictly less than the window or the windows never advance.
  const overlap = Math.min(opts.overlap, opts.segment - 1);
  const step = opts.segment - overlap;

  const ffprobe = resolveFfprobe(opts);
  const duration = await probeDuration(ffprobe, opts.audio);

  // Window offsets: 0, step, 2*step, ... while there is audio left to cover.
  const offsets = [];
  for (let t = 0; t < duration; t += step) {
    offsets.push(t);
    if (t + opts.segment >= duration) break; // this window reaches the end
  }
  const N = offsets.length;
  progress(`${path.basename(opts.audio)}: ${round3(duration)}s -> ${N} window(s) of ${opts.segment}s (overlap ${overlap}s)`);

  // Resumable work dir. Validate cached windows against the current windowing
  // (duration/segment/overlap); if anything changed, start fresh.
  const absAudio = path.resolve(opts.audio);
  const workDir = workDirFor(absAudio);
  const metaPath = path.join(workDir, "meta.json");
  const meta = { v: 1, duration: round3(duration), segment: opts.segment, overlap, n: N };
  await mkdir(workDir, { recursive: true });
  let cachedMeta = null;
  try {
    cachedMeta = JSON.parse(await readFile(metaPath, "utf8"));
  } catch {
    cachedMeta = null;
  }
  const compatible =
    cachedMeta &&
    cachedMeta.v === meta.v &&
    cachedMeta.duration === meta.duration &&
    cachedMeta.segment === meta.segment &&
    cachedMeta.overlap === meta.overlap &&
    cachedMeta.n === meta.n;
  if (!compatible) {
    await rm(workDir, { recursive: true, force: true });
    await mkdir(workDir, { recursive: true });
  }
  await writeFile(metaPath, JSON.stringify(meta));
  const winPath = (i) => path.join(workDir, `win-${String(i).padStart(4, "0")}.json`);
  // A cached window counts only if it's non-empty and parses back to a
  // { start, doc } payload. A 0-byte or truncated file (e.g. from a prior run
  // interrupted mid-write) is treated as absent so the window is re-transcribed,
  // rather than being trusted here and crashing the final stitch on JSON.parse.
  const cachedWindow = async (i) => {
    if (!existsSync(winPath(i))) return false;
    try {
      const raw = await readFile(winPath(i), "utf8");
      if (!raw.trim()) return false;
      const parsed = JSON.parse(raw);
      return parsed != null && typeof parsed === "object" && "doc" in parsed;
    } catch {
      return false;
    }
  };

  // Graceful pause: on SIGTERM/SIGINT, finish the in-flight window then stop.
  // Completed window JSONs stay in the work dir; we do NOT stitch (the final
  // transcript is produced only when ALL windows are done). Re-running resumes.
  let stopRequested = false;
  const onStop = () => {
    if (!stopRequested) progress("pause requested — finishing the current window, then stopping (re-run to resume)");
    stopRequested = true;
  };
  process.on("SIGTERM", onStop);
  process.on("SIGINT", onStop);

  const tmp = await mkdtemp(path.join(os.tmpdir(), "parakeet-stitch-"));
  const segTimes = []; // measured wall-time (ms) per freshly-transcribed window, for ETA
  let done = 0; // windows present (cached or freshly done)
  try {
    for (let i = 0; i < N; i += 1) {
      const start = offsets[i];
      // Resume: reuse a window already transcribed in a prior run.
      if (await cachedWindow(i)) {
        progress(`segment ${i + 1}/${N} @${round3(start)}s — cached`);
        done += 1;
        continue;
      }
      if (stopRequested) {
        progress(`paused after ${done}/${N} window(s) — re-run to resume`);
        process.exit(PARTIAL_PAUSE_CODE);
      }
      const isLast = i === N - 1;
      // Last window runs to EOF (no -t) so we never miss a trailing fragment.
      const len = isLast ? undefined : opts.segment;
      const wav = path.join(tmp, `win-${String(i).padStart(4, "0")}.wav`);
      const eta = etaSuffix(segTimes, N - i);
      progress(`segment ${i + 1}/${N} @${round3(start)}s — slicing${eta}`);
      await sliceWav(opts.ffmpeg, opts.audio, start, len, wav);
      progress(`segment ${i + 1}/${N} @${round3(start)}s — transcribing${eta}`);
      const t0 = Date.now();
      const doc = await transcribeWindow(opts, wav);
      segTimes.push(Date.now() - t0);
      // Persist the raw parakeet-cli payload for this window (the "original json"),
      // tagged with its absolute start offset. This is the durable, resumable
      // artifact; stitching happens only at the end.
      await writeFile(winPath(i), JSON.stringify({ start, doc }));
      done += 1;
      if (!opts.keepTemp) await rm(wav, { force: true });
    }
  } finally {
    if (!opts.keepTemp) await rm(tmp, { recursive: true, force: true });
  }

  // All windows present — load the raw payloads and stitch the final transcript.
  const windowWords = [];
  for (let i = 0; i < N; i += 1) {
    const { start, doc } = JSON.parse(await readFile(winPath(i), "utf8"));
    const words = Array.isArray(doc.words) ? doc.words : [];
    windowWords.push(
      words.map((w) => ({
        w: w.w ?? w.text ?? "",
        start: (Number(w.start) || 0) + start,
        end: (Number(w.end) || 0) + start,
        conf: w.conf,
      })),
    );
  }

  // Boundary cut points: the midpoint of each overlap region. Window i owns words
  // whose start is in [cut[i-1], cut[i]); the first/last windows are open-ended.
  // This keeps exactly one copy of every word across the seams.
  const cuts = [];
  for (let i = 0; i < N - 1; i += 1) {
    cuts.push(offsets[i + 1] + overlap / 2);
  }
  const stitched = [];
  for (let i = 0; i < N; i += 1) {
    const lo = i > 0 ? cuts[i - 1] : -Infinity;
    const hi = i < N - 1 ? cuts[i] : Infinity;
    for (const w of windowWords[i]) {
      if (w.start >= lo && w.start < hi) stitched.push(w);
    }
  }
  stitched.sort((a, b) => a.start - b.start);

  const chunkData = groupCues(stitched, { gap: opts.gap, maxCue: opts.maxCue });
  const text = chunkData.map((c) => c.text).join(" ");
  const doc = {
    duration_seconds: round3(duration),
    chunks: chunkData.length,
    text,
    chunk_data: chunkData,
    ...(opts.words ? { words: wordsOut(stitched) } : {}),
  };

  const json = JSON.stringify(doc);
  if (opts.output) {
    await writeFile(opts.output, json);
    progress(`wrote ${chunkData.length} cues (${stitched.length} words) -> ${opts.output}`);
  } else {
    process.stdout.write(json + "\n");
    progress(`emitted ${chunkData.length} cues (${stitched.length} words) to stdout`);
  }
  // Stitch succeeded — drop the per-window cache.
  if (!opts.keepTemp) await rm(workDir, { recursive: true, force: true });
}

main().catch((err) => {
  fail(err?.stack || String(err), 1);
});