// Vintage Chat frontend
const $ = (id) => document.getElementById(id);
const el = {
  modelInfo: $("model-info"),
  footerCkpt: $("footer-checkpoint"),
  footerDevice: $("footer-device"),
  registerTabs: $("register-tabs"),
  examplePrompts: $("example-prompts"),
  prompt: $("prompt"),
  generateBtn: $("generate-btn"),
  cancelBtn: $("cancel-btn"),
  status: $("status"),
  output: $("output"),
  copyBtn: $("copy-btn"),
  translateControls: $("translate-controls"),
  translateInput: $("translate-input"),
  translateOutput: $("translate-output"),
  resolvedPrompt: $("resolved-prompt"),
  gloss: $("gloss"),
  glossText: $("gloss-text"),
  maxNewTokens: $("max-new-tokens"),
  temperature: $("temperature"),
  topK: $("top-k"),
  topP: $("top-p"),
  repPenalty: $("rep-penalty"),
  seed: $("seed"),
};

let examples = { registers: [] };
let activeRegister = null;
let abortCtrl = null;
let translationAvailable = false;

async function loadConfig() {
  try {
    const r = await fetch("/api/config");
    const cfg = await r.json();
    el.footerCkpt.textContent = cfg.model.checkpoint || "—";
    el.footerDevice.textContent = `${cfg.model.device}${cfg.model.dtype ? " · " + cfg.model.dtype : ""}`;
    el.modelInfo.textContent = `${cfg.model.checkpoint?.split("/").slice(-2).join("/") || "model"} · ${cfg.model.device}`;
    translationAvailable = !!cfg.translation_available;
    el.translateControls.hidden = !translationAvailable;
  } catch (e) {
    el.modelInfo.textContent = "(config unavailable)";
  }
}

async function loadExamples() {
  const r = await fetch("/api/examples");
  examples = await r.json();
  renderRegisterTabs();
  if (examples.registers?.length) setActiveRegister(examples.registers[0].id);
}

function renderRegisterTabs() {
  el.registerTabs.innerHTML = "";
  for (const reg of examples.registers || []) {
    const btn = document.createElement("button");
    btn.type = "button";
    btn.className = "register-tab";
    btn.dataset.id = reg.id;
    btn.textContent = `${reg.label} · ${reg.label_en}`;
    btn.addEventListener("click", () => setActiveRegister(reg.id));
    el.registerTabs.appendChild(btn);
  }
}

function setActiveRegister(id) {
  activeRegister = id;
  for (const btn of el.registerTabs.querySelectorAll(".register-tab")) {
    btn.classList.toggle("active", btn.dataset.id === id);
  }
  const reg = examples.registers.find((r) => r.id === id);
  if (!reg) return;
  el.examplePrompts.innerHTML = "";
  const desc = document.createElement("div");
  desc.className = "register-description";
  desc.textContent = reg.description || "";
  el.examplePrompts.appendChild(desc);
  for (const p of reg.prompts || []) {
    const b = document.createElement("button");
    b.type = "button";
    b.className = "example-prompt";
    b.textContent = p;
    b.addEventListener("click", () => {
      el.prompt.value = p;
      el.translateInput.checked = false;
      el.prompt.focus();
    });
    el.examplePrompts.appendChild(b);
  }
}

function readParams() {
  const numeric = (input, fallback) => {
    const v = parseFloat(input.value);
    return Number.isFinite(v) ? v : fallback;
  };
  return {
    max_new_tokens: parseInt(el.maxNewTokens.value, 10) || 200,
    temperature: numeric(el.temperature, 0.8),
    top_k: parseInt(el.topK.value, 10) ?? 50,
    top_p: numeric(el.topP, 0.9),
    repetition_penalty: numeric(el.repPenalty, 1.3),
    seed: el.seed.value ? parseInt(el.seed.value, 10) : null,
  };
}

function setStatus(msg, kind = "info") {
  el.status.textContent = msg;
  el.status.className = "status" + (kind === "err" ? " err" : "");
}

async function translateInput(text) {
  setStatus("Translating input → 文言文…");
  const r = await fetch("/api/translate", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ text, direction: "to_classical", language: "auto" }),
  });
  if (!r.ok) throw new Error(`Translation failed (${r.status})`);
  const data = await r.json();
  return data.text;
}

async function translateOutput(text) {
  setStatus("Glossing output → modern…");
  const r = await fetch("/api/translate", {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({ text, direction: "to_modern", language: "Modern Chinese" }),
  });
  if (!r.ok) throw new Error(`Gloss failed (${r.status})`);
  const data = await r.json();
  return data.text;
}

async function runGenerate() {
  const rawPrompt = el.prompt.value.trim();
  if (!rawPrompt) return;
  el.generateBtn.disabled = true;
  el.cancelBtn.disabled = false;
  el.copyBtn.disabled = true;
  el.output.textContent = "";
  el.gloss.hidden = true;
  el.glossText.textContent = "";
  el.resolvedPrompt.hidden = true;
  el.resolvedPrompt.textContent = "";
  setStatus("");

  let prompt = rawPrompt;
  try {
    if (translationAvailable && el.translateInput.checked) {
      prompt = await translateInput(rawPrompt);
      el.resolvedPrompt.textContent = prompt;
      el.resolvedPrompt.hidden = false;
    }
  } catch (e) {
    setStatus(e.message || String(e), "err");
    el.generateBtn.disabled = false;
    el.cancelBtn.disabled = true;
    return;
  }

  const params = readParams();
  abortCtrl = new AbortController();

  setStatus("Generating…");

  try {
    const r = await fetch("/api/generate", {
      method: "POST",
      headers: { "Content-Type": "application/json", Accept: "text/event-stream" },
      body: JSON.stringify({ prompt, ...params }),
      signal: abortCtrl.signal,
    });
    if (!r.ok || !r.body) {
      const t = await r.text().catch(() => "");
      throw new Error(`Server error ${r.status}: ${t.slice(0, 200)}`);
    }

    const reader = r.body.getReader();
    const dec = new TextDecoder();
    let buf = "";
    let acc = "";
    el.output.textContent = prompt;

    const FRAME_RE = /\r?\n\r?\n/;

    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      buf += dec.decode(value, { stream: true });
      let match;
      while ((match = FRAME_RE.exec(buf))) {
        const frame = buf.slice(0, match.index);
        buf = buf.slice(match.index + match[0].length);
        const lines = frame.split(/\r?\n/);
        let event = "message";
        let data = "";
        for (const line of lines) {
          if (line.startsWith("event:")) event = line.slice(6).trim();
          else if (line.startsWith("data:")) data += line.slice(5).trim();
        }
        if (event === "delta" && data) {
          try {
            const parsed = JSON.parse(data);
            acc += parsed.text || "";
            el.output.textContent = prompt + acc;
          } catch {}
        } else if (event === "error") {
          try { setStatus(JSON.parse(data).message || "error", "err"); } catch { setStatus("error", "err"); }
        } else if (event === "done") {
          /* end */
        }
      }
    }

    el.copyBtn.disabled = !acc;
    setStatus(`Done · ${acc.length} chars`);

    if (translationAvailable && el.translateOutput.checked && acc) {
      try {
        const gloss = await translateOutput(prompt + acc);
        el.glossText.textContent = gloss;
        el.gloss.hidden = false;
        setStatus(`Done · ${acc.length} chars · glossed`);
      } catch (e) {
        setStatus(e.message || String(e), "err");
      }
    }
  } catch (e) {
    if (e.name === "AbortError") {
      setStatus("Cancelled");
    } else {
      setStatus(e.message || String(e), "err");
    }
  } finally {
    el.generateBtn.disabled = false;
    el.cancelBtn.disabled = true;
    abortCtrl = null;
  }
}

el.generateBtn.addEventListener("click", runGenerate);
el.cancelBtn.addEventListener("click", () => abortCtrl?.abort());
el.copyBtn.addEventListener("click", async () => {
  try { await navigator.clipboard.writeText(el.output.textContent); setStatus("Copied"); } catch {}
});
el.prompt.addEventListener("keydown", (e) => {
  if ((e.metaKey || e.ctrlKey) && e.key === "Enter") {
    e.preventDefault();
    runGenerate();
  }
});

loadConfig();
loadExamples();