/*
 * mt-service.js — the translation model, in a process of its own.
 *
 * Three runtimes were tried before this one, and the reasons matter:
 *
 *   • ONNX Runtime *Web* in the renderer: its weights go through a 32-bit WASM
 *     heap capped at 4 GB. NLLB aborts at fp16 and at q4f16. The only variant
 *     that fits is int8, which emits fluent-looking nonsense stitched together
 *     from a dozen unrelated languages.
 *
 *   • onnxruntime-node inside Electron's main process: crashes outright —
 *     EXC_BREAKPOINT in CPUAllocator::Alloc on CrBrowserMain while building the
 *     session. A multi-gigabyte synchronous allocation on the browser thread is
 *     not something that thread will tolerate.
 *
 *   • onnxruntime-node in a plain Node process: works. Correct output in every
 *     script tested, about 0.8 s per language on this M2 Max.
 *
 * So translation lives here, in a forked plain-Node process (the parent runs it
 * with ELECTRON_RUN_AS_NODE), with a normal 64-bit heap and nothing else
 * sharing it. Electron's own utilityProcess was tried first and exits with a
 * bare code 5 and no output. If this dies, the parent restarts it — with
 * backoff — and the app keeps running.
 */

'use strict';

const fs = require('fs');
const path = require('path');
const {
  pipeline, env, LogitsProcessor, LogitsProcessorList
} = require('@huggingface/transformers');

let pipe = null;

/*
 * ─────────────────────────── batched multi-target decode ───────────────────
 *
 * NLLB picks its target language with `forced_bos_token_id` — the FIRST token
 * the decoder is made to emit. The encoder input carries only the SOURCE
 * language tag (tokenizers.js `_build_translation_inputs` rewrites the
 * post-processor's special token to the source, and nothing else). So one
 * encoder input serves every target, and the only per-language difference is
 * one token at decode step 1.
 *
 * That makes N languages one batched decode instead of N passes over the same
 * 600M weights. This model is memory-bandwidth-bound — three worker processes
 * measured SLOWER than one (3.4s vs 1.9s for five languages) — and batching is
 * the right shape of fix for that: the weights are read once per step for the
 * whole batch instead of once per step per language.
 *
 * MEASURED on this M2 Max, one sentence, fp32, versus the same work serially:
 *
 *     2 langs  1.59x     3 langs  1.93x     5 langs  3.06x     9 langs  3.89x
 *
 * with output identical to the serial path in 9 of 9 languages.
 *
 * Two things had to be built for it:
 *
 *   • PER-ROW target forcing. The stock ForcedBOSTokenLogitsProcessor holds a
 *     single id for the whole batch, so every row would come out in the same
 *     language.
 *
 *   • STICKY EOS. transformers.js only leaves the decode loop once EVERY row
 *     is done (`if (stop.every(x => x)) break`) and nothing pads a row that
 *     already finished — its own source says `TODO: ensure all_input_ids is
 *     padded correctly`. Untreated, a short language keeps generating free
 *     text while a long one catches up: Spanish finished correctly and then
 *     drifted through Czech and Romanian for another hundred tokens. Pinning
 *     a finished row to EOS makes the tail strippable as special tokens.
 */
class BatchTargets extends LogitsProcessor {
  constructor(bosIds, eos) { super(); this.bosIds = bosIds; this.eos = eos; }
  _call(input_ids, logits) {
    for (let i = 0; i < input_ids.length; ++i) {
      const row = input_ids[i];
      const d = logits[i].data;
      if (row.length === 1) {
        /* Decode step 1: this row's target language, and nothing else. */
        d.fill(-Infinity);
        d[this.bosIds[i]] = 0;
      } else if (row[row.length - 1] === this.eos) {
        /* Already finished — hold it there rather than let it wander. */
        d.fill(-Infinity);
        d[Number(this.eos)] = 0;
      }
    }
    return logits;
  }
}

/* The batched path needs tokenizer internals. If a future version of the
   library moves them, say so once and keep translating one at a time. */
function batchable(tok) {
  return !!(tok &&
    typeof tok.lang_to_token === 'function' &&
    tok.languageRegex instanceof RegExp &&
    tok.post_processor && tok.post_processor.config &&
    Array.isArray(tok.post_processor.config.single) &&
    tok.model && typeof tok.model.convert_tokens_to_ids === 'function');
}

let warnedNoBatch = false;

async function translateMulti(text, src, tgts) {
  const tok = pipe.tokenizer, model = pipe.model;
  if (!batchable(tok)) {
    if (!warnedNoBatch) {
      warnedNoBatch = true;
      reply({ type: 'note', message: 'batched translation unavailable; using one pass per language' });
    }
    return null;
  }

  /* Pin the source language into the post-processor, the way the library's own
     translation path does. */
  for (const item of tok.post_processor.config.single) {
    if ('SpecialToken' in item && tok.languageRegex.test(item.SpecialToken.id)) {
      item.SpecialToken.id = tok.lang_to_token(src);
      break;
    }
  }

  const bosIds = tgts.map((t) => tok.model.convert_tokens_to_ids([tok.lang_to_token(t)])[0]);
  if (bosIds.some((id) => id === undefined || id === null)) return null;

  const enc = tok(new Array(tgts.length).fill(text), { padding: true, truncation: true });
  const inLen = enc.input_ids.dims[1];
  const eos = BigInt(model.generation_config.eos_token_id);

  const lp = new LogitsProcessorList();
  lp.push(new BatchTargets(bosIds, eos));

  const out = await model.generate({
    ...enc,
    logits_processor: lp,
    forced_bos_token_id: null,
    /* A row that never emits EOS would otherwise hold the whole batch to the
       model's 200-token ceiling. No translation of a sentence runs to four
       times its source length. */
    max_new_tokens: Math.min(200, inLen * 4 + 24)
  });

  const rows = tok.batch_decode(out, { skip_special_tokens: true });
  const texts = {};
  tgts.forEach((t, i) => { texts[t] = (rows[i] || '').trim(); });
  return texts;
}

function reply(msg) {
  try { if (process.send) process.send(msg); } catch (_) {}
}

process.on('message', async (msg) => {
  msg = msg || {};

  try {
    if (msg.type === 'load') {
      env.allowLocalModels = false;
      if (msg.cacheDir) env.cacheDir = msg.cacheDir;

      const load = () => pipeline('translation', msg.model, {
        dtype: msg.dtype || 'fp32',
        progress_callback: (p) => reply({ type: 'progress', p })
      });

      try {
        pipe = await load();
      } catch (err) {
        /*
         * A download interrupted partway leaves a truncated .onnx in the cache,
         * and every launch after that dies on "Protobuf parsing failed" — the
         * app is bricked until somebody knows to delete a folder. Throw the
         * cached copy away and fetch it once more instead.
         */
        const looksCorrupt = /protobuf|parsing failed|invalid model|corrupt|unexpected end/i
          .test((err && err.message) || '');
        if (!looksCorrupt) throw err;

        const dir = path.join(env.cacheDir, msg.model.replace(/\\/g, '/'));
        reply({ type: 'progress', p: { status: 'initiate', file: 're-downloading damaged model files' } });
        try { fs.rmSync(dir, { recursive: true, force: true }); } catch (_) {}
        pipe = await load();
      }

      reply({ type: 'ready' });
      return;
    }

    if (msg.type === 'translateMulti') {
      if (!pipe) throw new Error('translation model is not loaded yet');
      const text = String(msg.text || '');
      const src = String(msg.src || 'eng_Latn');
      const tgts = (msg.tgts || []).map(String);
      if (!tgts.length) { reply({ type: 'multi', id: msg.id, texts: {} }); return; }

      let texts = null;
      if (tgts.length > 1) texts = await translateMulti(text, src, tgts);

      /* One language, or a library that no longer exposes what batching needs:
         the ordinary path still answers, just without the speed-up. */
      if (!texts) {
        texts = {};
        for (const t of tgts) {
          const o = await pipe(text, { src_lang: src, tgt_lang: t });
          const g = Array.isArray(o) ? o[0] : o;
          texts[t] = (g && g.translation_text || '').trim();
        }
      }
      reply({ type: 'multi', id: msg.id, texts });
      return;
    }

    if (msg.type === 'translate') {
      if (!pipe) throw new Error('translation model is not loaded yet');
      const out = await pipe(String(msg.text || ''), {
        src_lang: String(msg.src || 'eng_Latn'),
        tgt_lang: String(msg.tgt || 'eng_Latn')
      });
      const got = Array.isArray(out) ? out[0] : out;
      reply({ type: 'result', id: msg.id, text: (got && got.translation_text || '').trim() });
      return;
    }
  } catch (err) {
    reply({
      type: msg.type === 'load' ? 'loadError' : 'error',
      id: msg.id,
      message: (err && err.message) || String(err)
    });
  }
});

/* A crash here must be a visible, recoverable event in the parent — never a
   silent stall that looks like a slow model. */
process.on('uncaughtException', (e) => {
  reply({ type: 'fatal', message: e.message });
  process.exit(1);
});