/*
 * Does one batched decode across N target languages beat N serial calls?
 *
 * The premise (verified in the installed tokenizer source): NLLB's tokenizer
 * prefixes the ENCODER input with the SOURCE language only — the target is
 * applied purely as `forced_bos_token_id` on the decoder. So one encoder input
 * serves every target, and the only thing that differs per row is the first
 * decoder token.
 */
import { pipeline, env, LogitsProcessor, LogitsProcessorList } from '@huggingface/transformers';
import os from 'os'; import path from 'path';
env.allowLocalModels = false;
env.cacheDir = path.join(os.homedir(), 'Library/Application Support/Church Translate Models');

const TEXT = 'And I want you to understand that God has a plan for your life';
const LANGS = ['spa_Latn','por_Latn','fra_Latn','zho_Hans','hat_Latn','hin_Deva','kor_Hang','vie_Latn','arb_Arab'];
const SRC = 'eng_Latn';

/* Forces a DIFFERENT language token on each row of the batch. The stock
   ForcedBOSTokenLogitsProcessor holds a single id for the whole batch. */
class PerRowForcedBOS extends LogitsProcessor {
  constructor(ids, eos) { super(); this.ids = ids; 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) {
        /* First decoder step: this row's target language. */
        d.fill(-Infinity);
        d[this.ids[i]] = 0;
        continue;
      }
      /*
       * transformers.js only breaks the decode loop once EVERY row is done
       * (models.js: `if (stop.every(x => x)) break`) and nothing pads a row
       * that already finished — so a short language keeps generating free
       * text while a long one catches up. Pin finished rows to EOS; the
       * trailing tokens are then stripped by skip_special_tokens.
       */
      if (row[row.length - 1] === this.eos) {
        d.fill(-Infinity);
        d[Number(this.eos)] = 0;
      }
    }
    return logits;
  }
}

const pipe = await pipeline('translation', 'Xenova/nllb-200-distilled-600M', { dtype: 'fp32' });
const tok = pipe.tokenizer, model = pipe.model;

console.log('tokenizer:', tok.constructor.name,
  '| lang_to_token:', typeof tok.lang_to_token,
  '| post_processor:', !!(tok.post_processor && tok.post_processor.config));

/* Pin the source language into the post-processor exactly as the library 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 EOS = BigInt(model.generation_config.eos_token_id ?? tok.model.convert_tokens_to_ids(['</s>'])[0]);
console.log('eos token id:', EOS.toString());
const bosOf = (l) => tok.model.convert_tokens_to_ids([tok.lang_to_token(l)])[0];

/* Confirm the encoder input really is target-independent. */
const a = tok(TEXT).input_ids.tolist()[0];
console.log('encoder ids (src-prefixed):', a.slice(0, 4).join(','), '… len', a.length);

async function batched(langs) {
  const enc = tok(Array(langs.length).fill(TEXT), { padding: true, truncation: true });
  const lp = new LogitsProcessorList();
  lp.push(new PerRowForcedBOS(langs.map(bosOf), EOS));
  const out = await model.generate({
    ...enc,
    logits_processor: lp,
    forced_bos_token_id: null,
    max_new_tokens: 128,
  });
  const rows = tok.batch_decode(out, { skip_special_tokens: true });
  const o = {}; langs.forEach((l, i) => o[l] = rows[i].trim());
  return o;
}

async function serial(langs) {
  const o = {};
  for (const l of langs) {
    const r = await pipe(TEXT, { src_lang: SRC, tgt_lang: l });
    o[l] = (Array.isArray(r) ? r[0] : r).translation_text.trim();
  }
  return o;
}

/* warm both paths */
await serial(['spa_Latn']); await batched(['spa_Latn']);

console.log('\n N |  serial  |  batched |  speedup');
console.log('---+----------+----------+---------');
const sample = {};
for (const n of [1, 2, 3, 5, 9]) {
  const langs = LANGS.slice(0, n);
  let s = Date.now(); const so = await serial(langs); s = Date.now() - s;
  let b = Date.now(); const bo = await batched(langs); b = Date.now() - b;
  console.log(String(n).padStart(2) + ' | ' + (s + 'ms').padStart(8) + ' | ' + (b + 'ms').padStart(8) +
    ' | ' + (s / b).toFixed(2) + '×');
  if (n === 9) { sample.serial = so; sample.batched = bo; }
}

console.log('\n--- identical output? ---');
let same = 0;
for (const l of LANGS) {
  const ok = sample.serial[l] === sample.batched[l];
  if (ok) same++; else console.log('  DIFFER ' + l + '\n    serial : ' + sample.serial[l] + '\n    batched: ' + sample.batched[l]);
}
console.log('  ' + same + '/' + LANGS.length + ' identical');
console.log('  spa batched: ' + sample.batched.spa_Latn);
console.log('  zho batched: ' + sample.batched.zho_Hans);