#!/usr/bin/env node
// Smoke test for the built voice packs: verifies every FileRef in the current
// catalogue (size + SHA-256), then synthesizes a phrase and a single word with
// every pack exactly as the voice worker will: the bytes of each FileRef are
// handed to Moonshine's TextToSpeech under their canonical names, from the
// output directory only. Nothing is fetched, so a pass also shows the packs
// work offline.
//
//   node tools/voices/smoke.mjs                      # every pack
//   node tools/voices/smoke.mjs en-kokoro-af_heart   # some packs
//   node tools/voices/smoke.mjs --wav /tmp/voices    # also write WAVs to audition
//   node tools/voices/smoke.mjs --out DIR            # another output dir
//
// Node >= 22 (the runtime's pthreads run on worker_threads). Timings are Node on
// this machine, not a phone or iPad: measure those on the devices.
import crypto from 'node:crypto';
import fs from 'node:fs';
import path from 'node:path';
import { fileURLToPath, pathToFileURL } from 'node:url';

const HERE = path.dirname(fileURLToPath(import.meta.url));
// The packs the site builds ship: beside the open-standards checkout, or $AAC_VOICES_DIR.
let out = process.env.AAC_VOICES_DIR || path.resolve(HERE, '../../../../../aac-voice-packs');
let wavDir = null;
const only = [];
const argv = process.argv.slice(2);
for (let i = 0; i < argv.length; i++) {
  if (argv[i] === '--out') out = path.resolve(argv[++i]);
  else if (argv[i] === '--wav') wavDir = path.resolve(argv[++i]);
  else only.push(argv[i]);
}

// A request, and a word said alone, per language.
const PHRASES = {
  en: ['I want more juice, please.', 'more'],
  es: ['Quiero más jugo, por favor.', 'más'],
  fr: ["Je veux encore du jus, s'il te plaît.", 'encore'],
  de: ['Ich möchte mehr Saft, bitte.', 'mehr'],
  pt: ['Eu quero mais suco, por favor.', 'mais'],
  it: ['Voglio ancora del succo, per favore.', 'ancora'],
  nl: ['Ik wil meer sap, alsjeblieft.', 'meer'],
  zh: ['我想要更多果汁。', '要'],
  ja: ['もっとジュースがほしいです。', 'もっと'],
  hi: ['मुझे और जूस चाहिए।', 'और'],
  uk: ['Я хочу ще соку, будь ласка.', 'ще'],
};

const pointer = JSON.parse(fs.readFileSync(path.join(out, 'index.json'), 'utf8'));
const cat = JSON.parse(fs.readFileSync(path.join(out, pointer.current), 'utf8'));
if (cat.schema !== 'aac-voice-index/1') throw new Error(`unknown schema ${cat.schema}`);

// 1. Every FileRef is present with the right size and hash.
const refs = [];
for (const r of Object.values(cat.runtimes)) refs.push(...r.files, r.licence_text);
for (const c of Object.values(cat.components)) refs.push(...c.files, c.licence_text);
for (const l of Object.values(cat.languages)) for (const p of l.packs) refs.push(...p.files, p.licence_text, ...(p.sample ? [p.sample] : []));
let bad = 0;
for (const f of refs) {
  const p = path.join(out, f.url);
  const data = fs.existsSync(p) ? fs.readFileSync(p) : null;
  const sha = data && crypto.createHash('sha256').update(data).digest('hex');
  if (!data || data.length !== f.bytes || sha !== f.sha256) { bad++; console.error(`BAD FileRef ${f.url}`); }
}
console.log(`filerefs: ${refs.length - bad}/${refs.length} verified`);
if (bad) process.exit(1);

// 2. Synthesis with every pack.
const bytes = (f) => new Uint8Array(fs.readFileSync(path.join(out, f.url)));
const failures = [], warned = [];
const runtimes = new Map();
async function runtime(id) {
  if (!runtimes.has(id)) {
    const rt = cat.runtimes[id];
    const { default: createModule } = await import(pathToFileURL(path.join(out, rt.module)).href);
    runtimes.set(id, await createModule({
      wasmBinary: fs.readFileSync(path.join(out, rt.wasm)), print: () => {}, printErr: () => {},
    }));
  }
  return runtimes.get(id);
}
const errText = (m, e) => (typeof e === 'number' && m.getExceptionMessage ? m.getExceptionMessage(e) : String(e));

function wav(file, audio, sr) {
  const b = Buffer.alloc(44 + audio.length * 2);
  b.write('RIFF', 0); b.writeUInt32LE(36 + audio.length * 2, 4); b.write('WAVEfmt ', 8);
  b.writeUInt32LE(16, 16); b.writeUInt16LE(1, 20); b.writeUInt16LE(1, 22); b.writeUInt32LE(sr, 24);
  b.writeUInt32LE(sr * 2, 28); b.writeUInt16LE(2, 32); b.writeUInt16LE(16, 34); b.write('data', 36);
  b.writeUInt32LE(audio.length * 2, 40);
  for (let i = 0; i < audio.length; i++) b.writeInt16LE(Math.max(-32768, Math.min(32767, Math.round(audio[i] * 32767))), 44 + i * 2);
  fs.writeFileSync(file, b);
}

for (const [code, lang] of Object.entries(cat.languages)) {
  for (const pack of lang.packs) {
    if (only.length && !only.includes(pack.id) && !only.includes(code)) continue;
    const m = await runtime(pack.runtime);
    const files = [...pack.requires.flatMap((c) => cat.components[c].files), ...pack.files];
    const keys = files.map((f) => f.name);
    const problems = [], warnings = [];
    // The canonical asset list of this runtime must be covered (Piper stage files excepted: the
    // pack carries one monolithic ORT model instead).
    try {
      const voiceArg = pack.engine === 'kokoro' ? pack.options.voice : '';
      const deps = JSON.parse(m.ttsDependencies(pack.lang_key, voiceArg)).groups.flatMap((g) => g.files.map((f) => f.name))
        .filter((n) => pack.engine === 'kokoro' || !(n.includes('/piper-voices/') || n.startsWith('kokoro/')));
      const missing = deps.filter((n) => !keys.includes(n));
      if (missing.length) problems.push(`missing canonical assets: ${missing.join(', ')}`);
    } catch (e) { problems.push(`ttsDependencies: ${errText(m, e)}`); }

    const names = Object.keys(pack.options), values = Object.values(pack.options);
    let t0 = performance.now(), tts;
    try { tts = new m.TextToSpeech(pack.lang_key, keys, files.map(bytes), names, values); }
    catch (e) { failures.push(`${pack.id}: load: ${errText(m, e)}`); console.log(`FAIL ${pack.id} load: ${errText(m, e)}`); continue; }
    const loadMs = performance.now() - t0;
    const lines = [];
    for (const [i, text] of (PHRASES[code] ?? ['test']).entries()) {
      t0 = performance.now();
      let r;
      try { r = tts.say(text); } catch (e) { problems.push(`say "${text}": ${errText(m, e)}`); continue; }
      const ms = performance.now() - t0;
      const a = r.audio, secs = a.length / r.sampleRate;
      let sum = 0, peak = 0, finite = true;
      for (const x of a) { if (!Number.isFinite(x)) finite = false; sum += x * x; peak = Math.max(peak, Math.abs(x)); }
      const rms = Math.sqrt(sum / Math.max(1, a.length));
      if (r.sampleRate !== pack.sample_rate) problems.push(`sample rate ${r.sampleRate} != ${pack.sample_rate}`);
      if (!finite) problems.push(`non-finite samples for "${text}"`);
      if (secs < 0.2 || rms < 0.005) problems.push(`silent or empty audio for "${text}" (${secs.toFixed(2)} s, rms ${rms.toFixed(4)})`);
      // A word said alone should not take seconds: VITS voices can babble on short input.
      if (i === 1 && secs > 2.5) warnings.push(`"${text}" alone took ${secs.toFixed(2)} s: listen for babble`);
      lines.push(`"${text}" ${secs.toFixed(2)} s in ${ms.toFixed(0)} ms (RTF ${(ms / 1000 / secs).toFixed(2)}, rms ${rms.toFixed(3)}, peak ${peak.toFixed(2)})`);
      if (wavDir && i === 0) { fs.mkdirSync(wavDir, { recursive: true }); wav(path.join(wavDir, `${pack.id}.wav`), a, r.sampleRate); }
    }
    tts.delete?.();
    const rss = (process.memoryUsage().rss / 1e6).toFixed(0);
    console.log(`${problems.length ? 'FAIL' : 'ok  '} ${pack.id}  load ${loadMs.toFixed(0)} ms, rss ${rss} MB`);
    for (const l of lines) console.log(`       ${l}`);
    for (const p of problems) console.log(`       ! ${p}`);
    for (const w of warnings) console.log(`       ? ${w}`);
    if (warnings.length) warned.push(pack.id);
    if (problems.length) failures.push(`${pack.id}: ${problems.join('; ')}`);
  }
}
if (failures.length) { console.error(`\n${failures.length} pack(s) failed`); process.exit(1); }
console.log(`\nall packs synthesized${warned.length ? `; audition ${warned.join(', ')}` : ''}`);
process.exit(0);