// Offline acoustic speaker diarization via sherpa-onnx (arm64 CPU, no pytorch).
// Splits a 16 kHz mono wav into speaker segments; AUTO speaker count (cluster
// threshold) because a conference phone ("Raumspinne") can put an unknown number
// of people on the single BT channel. Pinned to the A76 cores.
import { execFile } from 'node:child_process';
import { promisify } from 'node:util';
import fs from 'node:fs';
import * as spk from './spk.js';
const pexec = promisify(execFile);

// Auto-mode over-segmentation fix: sherpa's clustering splits ONE speaker into
// many clusters (e.g. 8-13 on a 2-person call, 61 on a 46-min one), and raising
// the cluster threshold barely helps. So after sherpa, we re-embed each cluster
// with the SAME TitaNet model used for voiceprint identity and MERGE clusters
// whose centroids are cosine-similar — i.e. the same voice that got split. This
// keeps the count dynamic (speakers can join/leave) without a fixed guess.
const MERGE_THRESHOLD = +(process.env.DIAR_MERGE_THRESHOLD || 0.5);
const REP_SEGMENTS = 3;    // longest segments per cluster averaged into its centroid
const MIN_REP_SEC = 0.4;   // skip ultra-short cuts for embedding when a longer one exists
// Stage 2: sub-second stray clusters (overlap crosstalk, noise) survive the
// centroid merge because they're too short to embed reliably. Any speaker with
// less than this much total speech is folded into the temporally-nearest real
// speaker. A genuine join/leave participant speaks well over this.
const MIN_SPEAKER_SEC = +(process.env.DIAR_MIN_SPEAKER_SEC || 3);
// Stage 3: a single >=MIN_SPEAKER_SEC stray cluster can still OPEN an anchor
// when its (short, unreliable) embedding lands below MERGE_THRESHOLD to every
// real speaker (#70: 3 phantom speakers with 2.5-4.8s total talk next to real
// ones with 53-736s). Total talk time over the whole call separates real
// participants from such noise: any FINAL speaker below this is folded into
// the acoustically nearest strong speaker. NOTE: measured on SEGMENT times
// (incl. intra-segment pauses), which run ~20-40% above summed word times —
// 10s let a 9.2s-word-time phantom through on #70, hence 20.
const MIN_TOTAL_SEC = +(process.env.DIAR_MIN_TOTAL_SEC || 20);

// Diarization on the capped A76 cores runs at ~1x wall-clock (a 46-min wav ≈
// 40 min). The old fixed 10-min timeout killed every recording over ~11 min and
// surfaced as an opaque "Command failed" (only sherpa's argv echo in stderr).
// Scale the timeout with the audio length: ~2.5x its duration, floor 20 min.
function diarTimeoutMs(wavPath) {
  let sec = 0;
  try { sec = Math.max(0, (fs.statSync(wavPath).size - 44) / 32000); } catch {} // 16kHz mono s16
  return Math.max(20 * 60 * 1000, Math.ceil(sec * 2.5) * 1000);
}

const DIR = '/opt/sherpa/sherpa-onnx-v1.13.3-linux-aarch64-shared-cpu';
const BIN = `${DIR}/bin/sherpa-onnx-offline-speaker-diarization`;
const SEG = process.env.DIAR_SEG || '/opt/sherpa/sherpa-onnx-pyannote-segmentation-3-0/model.onnx';
const EMB = process.env.DIAR_EMB || '/opt/sherpa/nemo_titanet_small.onnx';
// 0.7 (not 0.5): a lower threshold over-segments — a single (esp. child) voice
// split into several clusters. 0.7 merges more → fewer spurious speakers while
// still auto-detecting the count (needed for a conference phone / "Raumspinne").
const THRESHOLD = process.env.DIAR_THRESHOLD || '0.7';

function env() { return { ...process.env, LD_LIBRARY_PATH: `${DIR}/lib` }; }

// Returns [{ start, end, speaker }] (seconds, speaker = "speaker_00"…).
// numSpeakers > 0 fixes the count (--clustering.num-clusters) when the user knows
// it; otherwise auto via the cluster threshold.
export async function diarize(wavPath, numSpeakers) {
  const n = parseInt(numSpeakers, 10);
  const clusterArg = n > 0 ? `--clustering.num-clusters=${n}` : `--clustering.cluster-threshold=${THRESHOLD}`;
  const args = [
    'taskset', '-c', '4-7', BIN,
    `--segmentation.pyannote-model=${SEG}`,
    `--segmentation.num-threads=4`,
    `--embedding.model=${EMB}`,
    `--embedding.num-threads=4`,
    clusterArg,
    wavPath,
  ];
  const { stdout } = await pexec(args[0], args.slice(1), { env: env(), maxBuffer: 1 << 22, timeout: diarTimeoutMs(wavPath) });
  const segs = [];
  for (const line of stdout.split('\n')) {
    const m = line.match(/([\d.]+)\s*--\s*([\d.]+)\s*(speaker_\d+)/);
    if (m) segs.push({ start: +m[1], end: +m[2], speaker: m[3] });
  }
  // Only auto mode over-segments; a user-pinned count is already trusted.
  if (!(n > 0) && segs.length) {
    try { return await consolidateSpeakers(wavPath, segs); }
    catch (e) { console.log('diar: consolidation failed, using raw clusters:', e.message); }
  }
  return segs;
}

// Average voiceprint of a cluster: embed its longest few segments and mean them
// (short cuts give noisy embeddings, so prefer the longest and average out noise).
async function clusterCentroid(wavPath, clusterSegs) {
  const sorted = [...clusterSegs].sort((a, b) => (b.end - b.start) - (a.end - a.start));
  const pick = sorted.filter(s => s.end - s.start >= MIN_REP_SEC).slice(0, REP_SEGMENTS);
  const use = pick.length ? pick : sorted.slice(0, 1);
  const vecs = [];
  for (const s of use) {
    const seg = `/tmp/diarcons_${process.pid}_${Math.round(s.start * 1000)}.wav`;
    try {
      await pexec('ffmpeg', ['-y', '-i', wavPath, '-ss', String(s.start),
        '-to', String(Math.min(s.end, s.start + 8)), '-ar', '16000', '-ac', '1', seg]);
      const v = await spk.embed(seg);
      if (Array.isArray(v) && v.length) vecs.push(v);
    } catch { /* skip this cut */ }
    finally { try { fs.unlinkSync(seg); } catch {} }
  }
  if (!vecs.length) return null;
  const dim = vecs[0].length, c = new Array(dim).fill(0);
  for (const v of vecs) for (let i = 0; i < dim; i++) c[i] += v[i];
  for (let i = 0; i < dim; i++) c[i] /= vecs.length;
  return c;
}

// Merge sherpa's over-segmented clusters (anchor-based, see fn body),
// then renumber the surviving speakers by first appearance. Segments whose
// centroid couldn't be embedded keep their own label (safe: never lose a turn).
export async function consolidateSpeakers(wavPath, segs) {
  const labels = [...new Set(segs.map(s => s.speaker))];
  if (labels.length <= 1) return segs;

  // Centroid + total speech duration per raw sherpa cluster.
  const cents = {}, dur = {};
  for (const lab of labels) {
    const cs = segs.filter(s => s.speaker === lab);
    cents[lab] = await clusterCentroid(wavPath, cs);
    dur[lab] = cs.reduce((a, s) => a + (s.end - s.start), 0);
  }

  // ANCHOR-based merge (replaces the old pairwise union-find, which chained
  // transitively — A~B, B~C folded A and C together — so one borderline link
  // near the 0.5 cosine could cascade and the count flipped run-to-run under
  // load/thermal-cap FP noise: the SAME 2-3 person call came out 2, 3, or 7).
  // Only long clusters (reliable embeddings) become anchors and, processed
  // LONGEST-FIRST, either join an existing speaker or open a new one — a stable,
  // order-deterministic decision. Every short/stray cluster is then assigned to
  // its NEAREST anchor (never opens a speaker), so noise can't fragment.
  const anchors = labels
    .filter(l => cents[l] && dur[l] >= MIN_SPEAKER_SEC)
    .sort((a, b) => (dur[b] - dur[a]) || (a < b ? -1 : 1));

  const speakerOf = {};   // rawLabel -> canonical speaker rep (a rawLabel)
  const reps = [];        // accepted speaker representatives
  for (const a of anchors) {
    let joined = null;
    for (const r of reps) if (spk.cosine(cents[a], cents[r]) >= MERGE_THRESHOLD) { joined = r; break; }
    if (joined) speakerOf[a] = joined; else { speakerOf[a] = a; reps.push(a); }
  }

  const nearestRep = (lab) => {
    if (!reps.length) return lab;                 // no anchors -> keep own label
    if (!cents[lab]) return reps[0];              // unembeddable -> dominant speaker
    let best = reps[0], bs = -Infinity;
    for (const r of reps) { const c = spk.cosine(cents[lab], cents[r]); if (c > bs) { bs = c; best = r; } }
    return best;
  };

  const merged = segs.map(s => ({ ...s, speaker: speakerOf[s.speaker] || nearestRep(s.speaker) }));

  // Stage 3: fold weak FINAL speakers (< MIN_TOTAL_SEC total talk over the
  // whole call) into the acoustically nearest strong one — see MIN_TOTAL_SEC.
  const total = {};
  for (const s of merged) total[s.speaker] = (total[s.speaker] || 0) + (s.end - s.start);
  const strong = reps.filter(r => (total[r] || 0) >= MIN_TOTAL_SEC);
  if (strong.length && strong.length < reps.length) {
    const fold = {};
    for (const r of Object.keys(total)) {
      if (strong.includes(r)) continue;
      let best = strong[0], bs = -Infinity;
      for (const t of strong) {
        const c = (cents[r] && cents[t]) ? spk.cosine(cents[r], cents[t]) : -1;
        if (c > bs) { bs = c; best = t; }
      }
      fold[r] = best;
    }
    for (const s of merged) if (fold[s.speaker]) s.speaker = fold[s.speaker];
    console.log(`diar: folded ${Object.keys(fold).length} weak speaker(s) <${MIN_TOTAL_SEC}s total talk`);
  }

  // Renumber surviving speakers by first appearance.
  const rootMap = {}; let n = 0;
  for (const s of merged) if (!(s.speaker in rootMap)) rootMap[s.speaker] = `speaker_${String(n++).padStart(2, '0')}`;
  if (n !== labels.length) console.log(`diar: consolidated ${labels.length} -> ${n} speakers (anchors ${reps.length}, merge>=${MERGE_THRESHOLD}, floor ${MIN_SPEAKER_SEC}s)`);
  return merged.map(s => ({ ...s, speaker: rootMap[s.speaker] }));
}

// Assign each word (t_start_ms/t_end_ms) the speaker whose segment overlaps it
// most. Words with no overlap inherit the previous word's speaker.
export function assignSpeakers(words, segs) {
  if (!segs.length) return words;
  let last = segs[0].speaker;
  for (const w of words) {
    const ws = w.t_start_ms / 1000, we = w.t_end_ms / 1000;
    let best = null, bestOv = 0;
    for (const s of segs) {
      const ov = Math.min(we, s.end) - Math.max(ws, s.start);
      if (ov > bestOv) { bestOv = ov; best = s.speaker; }
    }
    w.speaker = best || last;
    last = w.speaker;
  }
  return words;
}

// Group consecutive same-speaker words into turns [{ speaker, text }].
export function toTurns(words) {
  const turns = [];
  for (const w of words) {
    const sp = w.speaker || 'speaker_00';
    if (!turns.length || turns[turns.length - 1].speaker !== sp) turns.push({ speaker: sp, text: w.word });
    else turns[turns.length - 1].text += ' ' + w.word;
  }
  return turns;
}