// AksharaLP sentence mode (v3): word-mode candidate generation from the v2 package, then a tiny on-device Hindi neural
// language model re-scores whole sentences by choosing one candidate per word jointly (search3).
// HARD GUARDRAIL: the search only ever returns a word that is in that slot's own candidate list; fixed slots pass through.
// Every result carries guard: {ok, words, outside}, so a caller can audit the guardrail at runtime.
import { NLM } from "./nlm.mjs";
import { search3 } from "./v3.mjs";

const RE_DEVA = /^[\u0900-\u0963\u0971-\u097F]+$/;

/** Chosen on validation/dev data only (PROGRESS.md '2026-11-01 07:05–07:30 IST — v3 FINAL'): Kaggle 'S' LM, q8, norm self. */
export const V3SET = {
  w_nlm: 0.4, w_lm: 0.25, w_ch_s: 1.5, w_agn: 2, w_agr: 0, nlm_margin: 2, nlm_k: 6, beam_s: 8,
  lat_k: 12, rec_n: 3, use_lm: true, oov_pen_s: 4.0, nlm_norm: "self",
};

/**
 * Attach sentence mode to a loaded BhashaCore instance.
 * `bytes` is the nlm.q8 file. Word mode keeps working: engine.sentenceMode toggles between the two decoders.
 */
export function attachSentence(engine, bytes) {
  const buf = bytes instanceof Uint8Array ? bytes : new Uint8Array(bytes);
  const dv = new DataView(buf.buffer, buf.byteOffset, buf.byteLength);
  const n = dv.getUint32(0, true);
  const meta = JSON.parse(new TextDecoder().decode(buf.subarray(4, 4 + n)));
  const nParams = meta.tensors.reduce((a, t) => a + t[1].reduce((x, y) => x * y, 1), 0);
  const nlm = new NLM(buf, V3SET.nlm_norm);

  const wordSlot = engine.decodeSlots.bind(engine);
  /** v3 build_lattice (bhc/translit3.py) + search3 (release/space/js/v3.mjs) */
  const sentSlot = (slots, context = true) => {
    const p = { ...engine.p, ...V3SET };
    const lat = [];
    for (const [kind, t] of slots) {
      if (kind === "x") {
        let sc = engine.candFeats(t).map(([s, fv, wid]) => [s, engine.score(fv), wid]);
        if (sc.length) {
          const mx = Math.max(...sc.map((x) => x[1])), lse = mx + Math.log(sc.reduce((a, x) => a + Math.exp(x[1] - mx), 0));
          sc = sc.map(([s, v, wid]) => [s, v - lse, wid]).sort((a, b) => b[1] - a[1]).slice(0, p.lat_k);
        }
        lat.push(["x", sc.length ? sc : [[t, 0, null]]]);
      } else lat.push(["fix", [[t, 0, RE_DEVA.test(t) ? (engine.lm.id(t) ?? -1) : -2]]]);
    }
    const post = lat.map(([, opts]) => opts.map(([s, v]) => [s, v]));   // per-slot candidate list (for the UI chips)
    let out, guard;
    if (!context) {                                               // context off: top-ranked candidate per slot (= word mode)
      out = lat.map((o) => o.reduce((a, b) => (b[1] > a[1] ? b : a))[0]);
      guard = { ok: true, words: 0, outside: 0 };
    } else {
      out = search3(lat, p, engine.lm, nlm);
      let words = 0, outside = 0;
      out.forEach((w, i) => { if (lat[i][0] !== "x") return; words++; if (!lat[i][1].some((o) => o[0] === w)) outside++; });
      guard = { ok: outside === 0, words, outside };
    }
    engine._lastGuard = guard;
    return { out, lat: post, guard };
  };

  engine.nlm = nlm;
  engine.v3 = { params: { ...V3SET }, params_count: nParams, cfg: meta.cfg, vocab: meta.vocab.length, bytes: buf.byteLength };
  engine.sentenceMode = true;
  engine._lastGuard = null;
  engine.decodeSlots = (slots, context) => (engine.sentenceMode ? sentSlot(slots, context) : wordSlot(slots, context));

  const base = engine.transliterate.bind(engine);
  engine.transliterate = function (text, opt) {
    this._lastGuard = null;
    const r = base(text, opt);
    r.guard = this._lastGuard || { ok: true, words: 0, outside: 0 };
    return r;
  };
  return { nlm, info: engine.v3 };
}