const $ = id => document.getElementById(id);
let worker, loaded = false, language = "en", generating = false;
const words = {
  en: {loading: "Loading model", working: "Writing on your device…", ready: "Ready when you are.", done: "Generation finished.", stopped: "Stopped.", bytes: "bytes/s"},
  fi: {loading: "Ladataan mallia", working: "Kirjoitetaan tällä laitteella…", ready: "Valmiina, kun sinä olet.", done: "Tekstin luonti päättyi.", stopped: "Pysäytetty.", bytes: "tavua/s"}
};
function setLanguage(value) {
  language = value;
  document.documentElement.lang = value;
  $("intro-en").hidden = value !== "en";
  $("intro-fi").hidden = value !== "fi";
  for (const node of document.querySelectorAll("[data-en]")) node.textContent = node.dataset[value];
  for (const lang of ["en", "fi"]) $(lang).setAttribute("aria-pressed", lang === value);
  updateModelInfo();
  if (generating) $("status").textContent = words[language].working;
}
$("en").onclick = () => setLanguage("en");
$("fi").onclick = () => setLanguage("fi");
$("theme").onclick = () => {
  document.documentElement.dataset.theme = document.documentElement.dataset.theme === "dark" ? "light" : "dark";
};
function finish() {
  generating = false;
  $("generate").disabled = false;
  $("stop").hidden = true;
  $("model-choice").disabled = false;
  $("loading").hidden = true;
}
function generate() {
  generating = true;
  $("model-choice").disabled = true;
  $("generate").disabled = true;
  $("stop").hidden = false;
  $("status").textContent = words[language].working;
  $("output").textContent = $("prompt").value;
  $("prefill").textContent = $("speed").textContent = "—";
  worker.postMessage({type: "generate", prompt: $("prompt").value,
    count: Number($("count").value), temperature: Number($("temperature").value), seed: Number($("seed").value)});
}
function load() {
  $("model-choice").disabled = true;
  $("generate").disabled = true;
  $("status").textContent = words[language].loading;
  $("loading").value = 0;
  $("loading").hidden = false;
  $("stop").hidden = false;
  worker = new Worker("worker.js?v=final-20000-1");
  worker.onmessage = ({data}) => {
    if (data.type === "loading") {
      $("loading").value = data.loaded / data.total;
      $("status").textContent = `${words[language].loading}: ${(data.loaded / 1e6).toFixed(1)} / ${(data.total / 1e6).toFixed(1)} MB`;
    }
    if (data.type === "ready") {
      loaded = true;
      $("loading").hidden = true;
      $("backend").textContent = data.backend === "wasm" ? "WebAssembly" : "WebGPU";
      generate();
    }
    if (data.type === "token") {
      $("output").textContent = data.text;
      $("prefill").textContent = `${data.prefillMs.toFixed(0)} ms`;
      $("speed").textContent = data.bytesPerSecond ? `${data.bytesPerSecond.toFixed(1)} ${words[language].bytes}` : "—";
      $("memory").textContent = `${(data.cacheBytes / 1048576).toFixed(2)} MiB`;
      $("cache-fill").style.width = `${data.positions / 512 * 100}%`;
      $("cache-fill").parentElement.setAttribute("aria-valuenow", data.positions);
      $("positions").textContent = `${data.positions} / 512`;
    }
    if (data.type === "done") {
      $("output").textContent = data.text;
      $("status").textContent = data.stopped ? words[language].stopped : words[language].done;
      finish();
    }
    if (data.type === "error") {
      $("status").textContent = data.message;
      finish();
      if (!loaded) { worker.terminate(); worker = null; }
    }
  };
  worker.onerror = event => {
    $("status").textContent = event.message || "Browser worker failed";
    finish(); worker.terminate(); worker = null; loaded = false;
  };
  worker.postMessage({type: "load", model: $("model-choice").value, backend: new URLSearchParams(location.search).get("backend") || "auto"});
}
$("form").onsubmit = event => {
  event.preventDefault();
  const length = new TextEncoder().encode($("prompt").value).length;
  if (!length || length >= 512) {
    $("status").textContent = language === "en" ? "Use a prompt of 1–511 UTF-8 bytes." : "Kehotteen pituuden on oltava 1–511 UTF-8-tavua.";
    return;
  }
  if (loaded) generate(); else load();
};
$("stop").onclick = () => {
  if (!loaded) {
    worker.terminate();
    worker = null;
    $("status").textContent = words[language].stopped;
    finish();
  } else worker.postMessage({type: "stop"});
};

function updateModelInfo() {
  const option = $("model-choice").selectedOptions[0];
  $("model-size").textContent = option.dataset.parameters;
  $("download-info").textContent = option.dataset[language === "en" ? "downloadEn" : "downloadFi"];
  $("model-state").textContent = language === "en" ? "Trained model" : "Koulutettu malli";
}
$("model-choice").onchange = () => {
  if (worker) worker.terminate();
  worker = null;
  loaded = false;
  $("output").textContent = "";
  $("prefill").textContent = $("speed").textContent = $("backend").textContent = "—";
  $("memory").textContent = "0 MiB";
  $("cache-fill").style.width = "0%";
  $("cache-fill").parentElement.setAttribute("aria-valuenow", "0");
  $("positions").textContent = "0 / 512";
  $("status").textContent = words[language].ready;
  updateModelInfo();
};
setLanguage("en");