/*
 * kokoro-sidecar.js — the natural voices, without Python.
 *
 * WHY THIS EXISTS. On macOS the natural voices come from `mlx_audio.server`, a
 * Python sidecar that runs Kokoro on MLX. MLX is Apple-Silicon-only, so on
 * Windows that process cannot exist at all — and the built-in Piper voices are
 * a quality drop the church rejected by name.
 *
 * So this serves the SAME HTTP API on the SAME port using the ONNX runtime the
 * app already ships. Everything upstream — sidecarWanted(), the warm-up, the
 * readiness handshake, the circuit breaker, the per-language voice choices —
 * is untouched, because from tts-service.js's point of view nothing changed.
 *
 * MEASURED against the live MLX sidecar, same text and voice:
 *   duration ratio  1.0000, 1.0000, 1.0000  (identical audio length)
 *   rms ratio       1.476 / 1.471 / 1.457   -> this build is ~3.3 dB HOTTER
 *   phonemes        byte-identical to misaki's EspeakG2P
 * The gain below cancels that level difference so a Windows campus and a Mac
 * campus sound the same. Duration being exactly equal is what matters most:
 * the HLS window pins a language's clips to one rate and length, and a
 * different duration would change the stream's pacing.
 *
 * dtype MUST be fp32. fp16 returns 100% NaN on the CPU execution provider —
 * silent audio with a flattering speed number — and q8 changes clip length.
 */
'use strict';
const http = require('http');
const fs = require('fs');
const path = require('path');

const PORT = Number(process.argv[2] || process.env.CT_KOKORO_PORT || 43111);
const MODEL_ID = 'onnx-community/Kokoro-82M-v1.0-ONNX';
const RATE = 24000;

/* Cancels the level difference against MLX (measured 1.457-1.476, mean 1.468). */
const MLX_MATCH_GAIN = 1 / 1.468;

/* The app speaks Kokoro's single-letter codes; espeak-ng wants its own names. */
const ESPEAK_VOICE = {
  a: 'en-us', b: 'en-gb', e: 'es', p: 'pt-br', f: 'fr-fr',
  i: 'it', h: 'hi', j: 'ja', z: 'cmn'
};

function log(msg) { process.stdout.write('[kokoro] ' + msg + '\n'); }

/* ---------------------------------------------------------------- G2P
 * A reimplementation of misaki.espeak.EspeakG2P — the exact preprocessing the
 * Python side does — on top of the espeak-ng WASM build. Verified phoneme-for-
 * phoneme identical to misaki across punctuation, «», numbers and Devanagari.
 */
const ZWJ = '‍';                       // espeak --ipa=3 tie character
const PUNCT = ';:,.!?¡¿—…"«»“”()';
const E2M = [
  ['a^ɪ', 'I'], ['a^ʊ', 'W'], ['d^z', 'ʣ'], ['d^ʒ', 'ʤ'], ['e^ɪ', 'A'],
  ['o^ʊ', 'O'], ['ə^ʊ', 'Q'], ['s^s', 'S'], ['t^s', 'ʦ'], ['t^ʃ', 'ʧ'], ['ɔ^ɪ', 'Y']
];

let ESpeakNg = null;
async function espeakIPA(text, voice) {
  if (!ESpeakNg) ESpeakNg = (await import('espeak-ng')).default;
  const e = await ESpeakNg({ arguments: ['--phonout', 'out', '-q', '--ipa=3', '-v', voice, text] });
  return e.FS.readFile('out', { encoding: 'utf8' });
}

async function g2p(text, espeakVoice) {
  text = String(text).replace(/«/g, '“').replace(/»/g, '”').replace(/\(/g, '«').replace(/\)/g, '»');
  const re = new RegExp('([' + PUNCT.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') + '])');
  let out = '';
  for (const p of text.split(re)) {
    if (!p) continue;
    if (p.length === 1 && PUNCT.includes(p)) { out += p; continue; }
    if (!p.trim()) { out += p; continue; }
    out += p.match(/^\s*/)[0] + (await espeakIPA(p.trim(), espeakVoice)).trim() + p.match(/\s*$/)[0];
  }
  let ps = out.replaceAll(ZWJ, '^');
  for (const [o, n] of E2M) ps = ps.replaceAll(o, n);
  return ps.replaceAll('^', '').replaceAll('-', '')
           .replace(/«/g, '(').replace(/»/g, ')').replace(/\s+/g, ' ').trim();
}

/* ------------------------------------------------------------- the model */
let model = null, tokenizer = null, Tensor = null, loading = null;
const styleCache = {};

function voicesDir() {
  const candidates = [
    process.env.CT_KOKORO_VOICES,
    path.join(__dirname, 'node_modules', 'kokoro-js', 'voices'),
    path.join(__dirname, '..', 'node_modules', 'kokoro-js', 'voices')
  ].filter(Boolean);
  for (const c of candidates) { try { if (fs.existsSync(c)) return c; } catch (_) {} }
  throw new Error('voice style vectors not found');
}

/*
 * One 256-float style vector per voice, chosen by TOKEN COUNT — Kokoro stores
 * 510 of them per voice and picks by length, which is how it varies delivery
 * with sentence length. Same indexing as the reference implementation.
 */
function styleFor(voice, nTokens) {
  if (!styleCache[voice]) {
    const buf = fs.readFileSync(path.join(voicesDir(), voice + '.bin'));
    styleCache[voice] = new Float32Array(buf.buffer.slice(buf.byteOffset, buf.byteOffset + buf.length));
  }
  const off = 256 * Math.min(Math.max(nTokens - 2, 0), 509);
  return styleCache[voice].slice(off, off + 256);
}

async function ensureModel() {
  if (model) return;
  if (loading) return loading;
  loading = (async () => {
    const tf = await import('@huggingface/transformers');
    Tensor = tf.Tensor;
    tf.env.allowLocalModels = false;
    if (process.env.CT_MODELS_DIR) tf.env.cacheDir = process.env.CT_MODELS_DIR;
    const t0 = Date.now();
    const [m, tk] = await Promise.all([
      tf.StyleTextToSpeech2Model.from_pretrained(MODEL_ID, { dtype: 'fp32', device: 'cpu' }),
      tf.AutoTokenizer.from_pretrained(MODEL_ID)
    ]);
    model = m; tokenizer = tk;
    log('model ready in ' + ((Date.now() - t0) / 1000).toFixed(1) + 's');
  })();
  return loading;
}

/*
 * ONE generation at a time.
 *
 * ONNX inference saturates the CPU, and the app fires several languages at
 * once per sentence. Run concurrently they all slow down together and every
 * one of them risks blowing its budget and falling back to a built-in voice —
 * which is the voice-switching failure this whole engine exists to avoid. A
 * queue makes the first clip fast and the rest predictable.
 */
let chain = Promise.resolve();
function serialize(fn) {
  const run = chain.then(fn, fn);
  chain = run.then(() => {}, () => {});
  return run;
}

function wavOf(float32) {
  const n = float32.length, b = Buffer.alloc(44 + n * 2);
  b.write('RIFF', 0); b.writeUInt32LE(36 + n * 2, 4); b.write('WAVE', 8); b.write('fmt ', 12);
  b.writeUInt32LE(16, 16); b.writeUInt16LE(1, 20); b.writeUInt16LE(1, 22);
  b.writeUInt32LE(RATE, 24); b.writeUInt32LE(RATE * 2, 28);
  b.writeUInt16LE(2, 32); b.writeUInt16LE(16, 34);
  b.write('data', 36); b.writeUInt32LE(n * 2, 40);
  for (let i = 0; i < n; i++) {
    let v = float32[i] * MLX_MATCH_GAIN;
    v = v > 0.999 ? 0.999 : (v < -0.999 ? -0.999 : v);
    b.writeInt16LE(Math.round(v * 32767), 44 + i * 2);
  }
  return b;
}

async function synth(voice, text, langCode, speed) {
  await ensureModel();
  const ev = ESPEAK_VOICE[langCode] || langCode || 'en-us';
  const phonemes = await g2p(text, ev);
  if (!phonemes) throw new Error('nothing to say');
  const { input_ids } = await tokenizer(phonemes, { truncation: true });
  const n = input_ids.dims.at(-1);
  const out = await model({
    input_ids,
    style: new Tensor('float32', styleFor(voice, n), [1, 256]),
    speed: new Tensor('float32', [Number(speed) || 1], [1])
  });
  return wavOf(out.waveform.data);
}

const server = http.createServer((req, res) => {
  const url = (req.url || '').split('?')[0];

  if (req.method === 'GET' && (url === '/healthz' || url === '/')) {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ ok: true, engine: 'kokoro-onnx', ready: !!model }));
    return;
  }
  /* Same shape mlx_audio returns, so anything probing it sees what it expects. */
  if (req.method === 'GET' && url === '/v1/models') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ object: 'list', data: [{ id: MODEL_ID, object: 'model', owned_by: 'system' }] }));
    return;
  }
  if (req.method === 'POST' && url === '/v1/audio/speech') {
    let body = '';
    req.on('data', (d) => { body += d; if (body.length > 1 << 20) req.destroy(); });
    req.on('end', () => {
      let j;
      try { j = JSON.parse(body); } catch (_) { res.writeHead(400).end('bad json'); return; }
      const voice = String(j.voice || '');
      const text = String(j.input || '');
      const lang = String(j.lang_code || 'a');
      if (!voice || !text) { res.writeHead(400).end('voice and input are required'); return; }
      /*
       * Only Kokoro lives here. A request for the Qwen3 expressive model (an
       * operator picking a "Local" voice) is refused plainly rather than
       * answered with the wrong voice — tts-service then falls back to the
       * built-in voice for that language, which is the honest outcome on a
       * platform where Qwen3/MLX cannot run.
       */
      if (j.model && String(j.model).indexOf('Kokoro') === -1) {
        res.writeHead(501, { 'Content-Type': 'text/plain' });
        res.end('this build serves Kokoro only');
        return;
      }
      serialize(() => synth(voice, text, lang, j.speed))
        .then((wav) => {
          res.writeHead(200, { 'Content-Type': 'audio/wav', 'Content-Length': wav.length });
          res.end(wav);
        })
        .catch((err) => {
          log('synth failed (' + voice + '/' + lang + '): ' + err.message);
          res.writeHead(500, { 'Content-Type': 'text/plain' });
          res.end(String(err.message || err));
        });
    });
    return;
  }
  res.writeHead(404).end('not found');
});

server.listen(PORT, '127.0.0.1', () => {
  log('listening on 127.0.0.1:' + PORT + ' (Kokoro via ONNX, no Python)');
  /* Load now rather than on the first sentence of the service. */
  ensureModel().catch((e) => log('model load failed: ' + e.message));
});
server.on('error', (e) => { log('cannot listen on ' + PORT + ': ' + e.message); process.exit(1); });