// Lab: window-level speaker embeddings for a recording.
// Reads sherpa's raw turn boundaries, slices the 16k mono wav into fixed windows
// (no ffmpeg: direct PCM slicing), embeds each window with a given sherpa
// embedding model via the spk_embed helper, and dumps JSON for offline
// clustering experiments (cluster-lab.mjs).
//
// Usage: node lab-embed.mjs <segs.txt> <wav> <model.onnx> <out.json> [winSec] [hopSec]
import fs from 'node:fs';
import { spawn } from 'node:child_process';
import readline from 'node:readline';
import os from 'node:os';

const [segsPath, wavPath, model, outPath, winSecArg, hopSecArg] = process.argv.slice(2);
const WIN = +(winSecArg || 3.0), HOP = +(hopSecArg || 1.5);
const MIN_WIN = 0.6;                       // shortest slice we bother embedding
const DIR = '/opt/sherpa/sherpa-onnx-v1.13.3-linux-aarch64-shared-cpu';
const TMP = fs.mkdtempSync(`${os.tmpdir()}/labemb_`);

// ---- wav (PCM s16 mono) ----------------------------------------------------
function loadWav(p) {
  const b = fs.readFileSync(p);
  if (b.toString('ascii', 0, 4) !== 'RIFF' || b.toString('ascii', 8, 12) !== 'WAVE') throw new Error('not RIFF/WAVE');
  let off = 12, fmt = null, data = null;
  while (off + 8 <= b.length) {
    const id = b.toString('ascii', off, off + 4), size = b.readUInt32LE(off + 4);
    if (id === 'fmt ') fmt = { ch: b.readUInt16LE(off + 10), rate: b.readUInt32LE(off + 12), bits: b.readUInt16LE(off + 22) };
    if (id === 'data') { data = { start: off + 8, size: Math.min(size, b.length - off - 8) }; break; }
    off += 8 + size + (size & 1);
  }
  if (!fmt || !data) throw new Error('wav chunks missing');
  if (fmt.bits !== 16) throw new Error('need s16, got ' + fmt.bits);
  return { buf: b, fmt, data };
}
function writeSlice(w, startSec, endSec, path) {
  const bytesPerFrame = 2 * w.fmt.ch;
  const s = Math.max(0, Math.floor(startSec * w.fmt.rate) * bytesPerFrame);
  const e = Math.min(w.data.size, Math.ceil(endSec * w.fmt.rate) * bytesPerFrame);
  const n = Math.max(0, e - s);
  const h = Buffer.alloc(44);
  h.write('RIFF', 0); h.writeUInt32LE(36 + n, 4); h.write('WAVE', 8);
  h.write('fmt ', 12); h.writeUInt32LE(16, 16); h.writeUInt16LE(1, 20);
  h.writeUInt16LE(w.fmt.ch, 22); h.writeUInt32LE(w.fmt.rate, 24);
  h.writeUInt32LE(w.fmt.rate * bytesPerFrame, 28); h.writeUInt16LE(bytesPerFrame, 32); h.writeUInt16LE(16, 34);
  h.write('data', 36); h.writeUInt32LE(n, 40);
  fs.writeFileSync(path, Buffer.concat([h, w.buf.subarray(w.data.start + s, w.data.start + e)]));
  return n / (w.fmt.rate * bytesPerFrame);
}

// ---- spk_embed pipe --------------------------------------------------------
function embedder(modelPath) {
  const proc = spawn('taskset', ['-c', '4-7', '/opt/callrec/spk_embed', modelPath],
    { env: { ...process.env, LD_LIBRARY_PATH: `${DIR}/lib` }, stdio: ['pipe', 'pipe', 'pipe'] });
  let pending = null;
  const rl = readline.createInterface({ input: proc.stdout });
  const ready = new Promise((res, rej) => {
    const to = setTimeout(() => rej(new Error('READY timeout')), 180000);
    rl.on('line', line => {
      if (line.startsWith('READY')) { clearTimeout(to); return res(); }
      if (!pending) return;
      const p = pending; pending = null;
      if (line.startsWith('OK ')) p.res(line.slice(3).split(',').map(Number));
      else p.rej(new Error(line));
    });
  });
  proc.stderr.on('data', d => process.stderr.write('[spk] ' + d));
  let chain = ready;
  return {
    embed(path) {
      chain = chain.then(() => new Promise((res, rej) => { pending = { res, rej }; proc.stdin.write(path + '\n'); }));
      return chain;
    },
    close() { try { proc.kill(); } catch {} },
  };
}

// ---- main ------------------------------------------------------------------
const turns = [];
for (const line of fs.readFileSync(segsPath, 'utf8').split('\n')) {
  const m = line.match(/([\d.]+)\s*--\s*([\d.]+)\s*(speaker_\d+)/);
  if (m) turns.push({ start: +m[1], end: +m[2], raw: m[3] });
}
turns.sort((a, b) => a.start - b.start);
console.log(`turns=${turns.length} rawLabels=${new Set(turns.map(t => t.raw)).size} speech=${turns.reduce((a, t) => a + t.end - t.start, 0).toFixed(0)}s`);

const wav = loadWav(wavPath);
console.log('wav', wav.fmt, (wav.data.size / (wav.fmt.rate * 2 * wav.fmt.ch)).toFixed(1) + 's');

const wins = [];
turns.forEach((t, ti) => {
  const dur = t.end - t.start;
  if (dur < MIN_WIN) return;
  if (dur <= WIN) { wins.push({ ti, start: t.start, end: t.end }); return; }
  for (let s = t.start; s + MIN_WIN <= t.end; s += HOP) wins.push({ ti, start: s, end: Math.min(s + WIN, t.end) });
});
console.log(`windows=${wins.length}`);

const emb = embedder(model);
const t0 = Date.now();
let done = 0;
for (const w of wins) {
  const p = `${TMP}/w.wav`;
  writeSlice(wav, w.start, w.end, p);
  try { w.vec = await emb.embed(p); } catch (e) { w.err = e.message; }
  if (++done % 100 === 0) {
    const rate = done / ((Date.now() - t0) / 1000);
    console.log(`  ${done}/${wins.length} ${rate.toFixed(1)}/s eta ${((wins.length - done) / rate).toFixed(0)}s`);
  }
}
emb.close();
const ok = wins.filter(w => w.vec);
fs.writeFileSync(outPath, JSON.stringify({
  model, wav: wavPath, win: WIN, hop: HOP,
  dim: ok.length ? ok[0].vec.length : 0,
  turns, wins: wins.map(w => ({ ti: w.ti, s: +w.start.toFixed(3), e: +w.end.toFixed(3), v: w.vec ? w.vec.map(x => +x.toFixed(5)) : null })),
}));
console.log(`done ok=${ok.length}/${wins.length} dim=${ok.length ? ok[0].vec.length : 0} in ${((Date.now() - t0) / 1000).toFixed(0)}s -> ${outPath}`);
try { fs.rmSync(TMP, { recursive: true, force: true }); } catch {}