// The demo page: pick a language, download that model once (the browser keeps it), run gector.js on the text.
import { Gector } from "./gector.js";

const MODELS = {
  en: { name: "English", tag: "v1.0.0", licence: "Apache-2.0",
        text: "My brother have two cat. The results was very good.", tries: ["He go to school every day.", "There is many people here."] },
  it: { name: "Italiano", tag: "v1.0.0", licence: "Apache-2.0",
        text: "La ragazza sono molto stanco. Le mie sorelle è molto simpatico.", tries: ["I bambini gioca nel giardino.", "Io sono andato a il mare."] },
  es: { name: "Español", tag: "v1.2.0", licence: "CC BY 4.0",
        text: "La casa son muy bonito. Mis hermanos tiene dos gato.", tries: ["Los niños juega en el parque.", "¿Donde está la estación?"] },
  fr: { name: "Français", tag: "v1.1.0", licence: "Apache-2.0",
        text: "Ma sœur sont très content. Les enfants joue dans le jardin.", tries: ["Nous est arrivés en retard.", "Je vais à le marché."] },
  de: { name: "Deutsch", tag: "v2.1.0", licence: "Apache-2.0",
        text: "Meine Schwester sind sehr mude. Du hat heute keine zeit.", tries: ["Die Kinder spielt im Garten.", "Ich hoffe, das du kommst."] },
};
const NAMES = { en: "English", it: "Italian", es: "Spanish", fr: "French", de: "German" };
const CACHE = "proofbird-models-v1";
const $ = (id) => document.getElementById(id);

ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.22.0/dist/";

let lang = pickLanguage();
const loaded = {};                                        // lang -> Promise<Gector>
let edited = false;

function pickLanguage() {
  for (const l of navigator.languages || [navigator.language || "en"]) {
    const code = l.slice(0, 2).toLowerCase();
    if (MODELS[code]) return code;
  }
  return "en";
}

const url = (l, file) => `https://huggingface.co/proofbird/gector-bert-${l}/resolve/${MODELS[l].tag}/${file}`;

function renderLanguages() {
  $("langs").replaceChildren(...Object.entries(MODELS).map(([code, m]) => {
    const b = document.createElement("button");
    b.className = "lang";
    b.textContent = code === lang ? `${m.name} · ${m.tag.slice(1)}` : m.name;
    b.setAttribute("aria-pressed", String(code === lang));
    b.onclick = () => selectLanguage(code);
    return b;
  }));
  const m = MODELS[lang];
  $("chips").replaceChildren(...m.tries.map((t) => {
    const c = document.createElement("button");
    c.className = "chip";
    c.textContent = t;
    c.onclick = () => { $("text").value = t; edited = true; check(); };
    return c;
  }));
  $("card").href = `https://huggingface.co/proofbird/gector-bert-${lang}`;
  $("licence").textContent = m.licence;
  document.documentElement.lang = lang;
}

function selectLanguage(code) {
  lang = code;
  if (!edited) $("text").value = MODELS[code].text;
  $("out").hidden = true;
  $("err").hidden = true;
  renderLanguages();
}

// --- files: from the browser's cache when we have them, else from Hugging Face (with progress for the big one) ---------
async function fetchFile(u, onProgress) {
  let cache = null;
  try { cache = await caches.open(CACHE); } catch { /* no Cache Storage (private window): download every time */ }
  const hit = cache && await cache.match(u);
  if (hit) return hit.arrayBuffer();
  const res = await fetch(u);
  if (!res.ok) throw new Error(`${res.status} for ${u}`);
  const total = Number(res.headers.get("content-length")) || 0;
  let body;
  if (onProgress && res.body) {
    const reader = res.body.getReader();
    const parts = [];
    let got = 0;
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      parts.push(value);
      got += value.length;
      onProgress(got, total);
    }
    body = new Uint8Array(got);
    let at = 0;
    for (const p of parts) { body.set(p, at); at += p.length; }
    body = body.buffer;
  } else {
    body = await res.arrayBuffer();
  }
  if (cache) {
    try { await cache.put(u, new Response(body, { headers: { "content-type": "application/octet-stream" } })); }
    catch { /* quota: fine, the next visit downloads again */ }
  }
  return body;
}

const mb = (n) => Math.round(n / 1e6);

function load(l) {
  if (loaded[l]) return loaded[l];
  loaded[l] = (async () => {
    const dec = new TextDecoder();
    const text = async (f) => dec.decode(await fetchFile(url(l, f)));
    const [config, tokenizer, rules, verbs] = await Promise.all([
      text("config.json").then(JSON.parse), text("tokenizer.json").then(JSON.parse),
      text("language-rules.json").then(JSON.parse).catch(() => ({})), text("verb-form-vocab.txt").catch(() => ""),
    ]);
    let cachedModel = false;
    try { cachedModel = !!(await (await caches.open(CACHE)).match(url(l, "model.onnx"))); } catch { /* no cache */ }
    if (!cachedModel) {
      $("dl").hidden = false;
      $("dl-label").textContent = `Downloading the ${NAMES[l]} model, once`;
      $("dl-count").textContent = "";
      $("dl-bar").style.width = "0";
    }
    const model = await fetchFile(url(l, "model.onnx"), (got, total) => {
      $("dl-count").textContent = total ? `${mb(got)} of ${mb(total)} MB` : `${mb(got)} MB`;
      $("dl-bar").style.width = total ? `${Math.min(100, (got / total) * 100)}%` : "50%";
    });
    $("dl-label").textContent = `Starting the ${NAMES[l]} model`;
    const session = await ort.InferenceSession.create(model, { executionProviders: ["wasm"] });
    $("dl").hidden = true;
    const run = async (ids) => {
      const big = (xs) => BigInt64Array.from(xs, (x) => BigInt(x));
      const r = await session.run({
        input_ids: new ort.Tensor("int64", big(ids), [1, ids.length]),
        attention_mask: new ort.Tensor("int64", big(ids.map(() => 1)), [1, ids.length]),
      });
      return { labels: r.logits_labels.data, d: r.logits_d.data };
    };
    return new Gector({ config, tokenizer, rules, verbs }, l, run);
  })();
  loaded[l].catch(() => { delete loaded[l]; });
  return loaded[l];
}

// --- check -------------------------------------------------------------------------------------------------------------
async function check() {
  const text = $("text").value;
  $("err").hidden = true;
  if (!text.trim()) { showError("Type or paste a sentence first."); return; }
  const l = lang;
  $("check").disabled = true;
  $("check").textContent = "Checking…";
  try {
    const g = await load(l);
    const t0 = performance.now();
    // one sentence at a time, as the Proofbird app does (the model sees one sentence; two at once can hide an error)
    const results = [];
    for (const sent of sentences(text)) {
      const r = await g.correct(sent.text);
      results.push({ ...sent, passes: g.passes,
                     suggestions: r.suggestions.map((x) => ({ ...x, start: x.start + sent.start })) });
    }
    const ms = Math.round(performance.now() - t0);
    if (l === lang) show(g, text, results, ms);
  } catch (e) {
    $("dl").hidden = true;
    showError(`The ${NAMES[l]} model couldn't be loaded (${e.message || e}). Check your connection and try again.`);
  } finally {
    $("check").disabled = false;
    $("check").textContent = "Check";
  }
}

function showError(msg) {
  $("err").textContent = msg;
  $("err").hidden = false;
}

// sentence boundaries: end punctuation, then space, then something that can start a sentence
function sentences(text) {
  const out = [];
  const re = /[.!?…]+["»”')\]]*\s+(?=[\p{Lu}\p{N}¿¡"„«(])/gu;
  let start = 0, m;
  while ((m = re.exec(text))) {
    const end = m.index + m[0].trimEnd().length;
    if (text.slice(start, end).trim()) out.push({ start, text: text.slice(start, end) });
    start = m.index + m[0].length;
  }
  const rest = text.slice(start).trimEnd();
  if (rest.trim()) out.push({ start, text: rest });
  return out;
}

function show(g, original, results, ms) {
  const suggestions = results.flatMap((r) => r.suggestions);
  const passes = Math.max(...results.map((r) => r.passes.length));
  $("meta").textContent = `Result · ${suggestions.length} ${suggestions.length === 1 ? "change" : "changes"} · `
    + `${passes} ${passes === 1 ? "pass" : "passes"} · ${ms} ms`;
  // the corrected text with each replacement marked (positions shift as replacements are applied left to right)
  const frag = document.createDocumentFragment();
  let at = 0;
  for (const s of suggestions) {
    frag.append(original.slice(at, s.start));
    if (s.replacement.trim()) {
      const span = document.createElement("span");
      span.className = "ch";
      span.textContent = s.replacement;
      frag.append(span);
    }
    at = s.start + s.length;
  }
  frag.append(original.slice(at));
  $("result").replaceChildren(frag);

  $("sugs").replaceChildren(...results.flatMap((r) => r.suggestions.map((s) => {
    const row = document.createElement("div");
    row.className = "sug";
    const old = document.createElement("span");
    old.className = "old";
    old.textContent = s.original.trim() || "(nothing)";
    const arrow = document.createElement("span");
    arrow.textContent = "→";
    arrow.setAttribute("aria-hidden", "true");
    const neu = document.createElement("span");
    neu.className = "new";
    neu.textContent = s.replacement.trim() || "(removed)";
    const tag = document.createElement("span");
    tag.className = "tag";
    tag.textContent = tagsFor(g, s, r.passes);
    row.append(old, arrow, neu, tag);
    return row;
  })));

  const many = results.length > 1;
  $("passes").replaceChildren(...results.flatMap((r, si) => r.passes.flatMap((p) => {
    const head = document.createElement("p");
    head.className = "pass";
    const lead = many ? `Sentence ${si + 1}, pass ${p.pass}` : `Pass ${p.pass}`;
    if (!p.changed) {
      head.textContent = `${lead} · ${p.input} (no change, stops)`;
      return [head];
    }
    head.textContent = `${lead} · ${p.input}`;
    const tab = document.createElement("div");
    tab.className = "ptab";
    for (const h of ["Word", "Wrong", "Edit", "Prob."]) tab.append(cell(h, "h"));
    for (const w of p.words.filter((w) => w.edit !== null || w.wrong >= 0.05)) {
      tab.append(cell(w.word), cell(pct(w.wrong)), cell(w.tag, w.tag !== "$KEEP" ? "e" : ""), cell(pct(w.probability)));
    }
    return [head, tab];
  })));
  $("out").hidden = false;
}

// which tags made this change: follow the original words through the passes
function tagsFor(g, s, passes) {
  let cur = new Set(g.words(s.original));
  const used = [];
  for (const p of passes) {
    const next = new Set(cur);
    p.words.forEach((w) => {
      if (w.edit === null) return;
      const appendHere = w.tag.startsWith("$APPEND_") && g.words(s.replacement).includes(w.tag.slice(8));
      if (cur.has(w.word) || appendHere) {
        used.push(`${w.tag} · ${pct(w.probability)}`);
        for (const x of g.words(w.edit)) next.add(x);
      }
    });
    cur = next;
  }
  return [...new Set(used)].join(", ") || "";
}

function cell(text, cls = "") {
  const d = document.createElement("span");
  if (cls) d.className = cls;
  d.textContent = text;
  return d;
}

const pct = (x) => `${(x * 100).toFixed(1)}%`;

$("text").addEventListener("input", () => { edited = true; $("err").hidden = true; });
$("text").addEventListener("keydown", (e) => { if (e.key === "Enter" && (e.ctrlKey || e.metaKey)) check(); });
$("check").onclick = check;
$("text").value = MODELS[lang].text;
renderLanguages();