/* ============================================================
   VLPS · Vision Lab — frontend logic
   ============================================================ */
const $  = (s, r = document) => r.querySelector(s);
const $$ = (s, r = document) => [...r.querySelectorAll(s)];

const TOOLS = [
  { id: "retrieve", label: "Retrieve" },
  { id: "vqa",      label: "Visual QA" },
  { id: "ocr",      label: "OCR Lab" },
  { id: "browse",   label: "Index Browser" },
  { id: "agent",    label: "Agent" },
];

/* ---------- tool switching ---------- */
const nav = $("#nav"), rail = $("#toolRail"), readout = $("#vpReadout");
TOOLS.forEach((t, i) => {
  const n = document.createElement("button");
  n.textContent = t.label; n.dataset.tool = t.id;
  nav.appendChild(n);
  const r = document.createElement("button");
  r.dataset.tool = t.id;
  r.innerHTML = `<span>${t.label}</span><span class="idx">0${i + 1}</span>`;
  rail.appendChild(r);
});

function selectTool(id) {
  $$(".panel").forEach(p => p.classList.toggle("active", p.dataset.tool === id));
  $$("#nav button, #toolRail button").forEach(b => b.classList.toggle("on", b.dataset.tool === id));
  const label = TOOLS.find(t => t.id === id)?.label || id;
  readout.textContent = `${label.toUpperCase()} · idle`;
  window.scrollTo({ top: $(".console").offsetTop - 70, behavior: "smooth" });
}
$$("#nav button, #toolRail button").forEach(b => b.onclick = () => selectTool(b.dataset.tool));

/* ---------- segmented controls ---------- */
$$(".seg").forEach(seg => seg.onclick = e => {
  const b = e.target.closest("button"); if (!b) return;
  $$("button", seg).forEach(x => x.classList.toggle("on", x === b));
});
const segVal = seg => $(".on", seg)?.dataset.v;

/* ---------- busy / fetch helpers ---------- */
function busy(on, label) {
  readout.classList.toggle("busy", on);
  if (on && label) readout.textContent = `${label} · working…`;
}
async function postJSON(path, body) {
  const r = await fetch(path, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body) });
  return r.json();
}
async function postForm(path, form) {
  const r = await fetch(path, { method: "POST", body: form });
  return r.json();
}

/* ---------- grid rendering ---------- */
function renderGrid(container, items, kind) {
  container.innerHTML = "";
  items.forEach(it => {
    const cell = document.createElement("div");
    cell.className = "cell";
    const cap = kind === "browse" ? (it.question || "") : (it.caption || "");
    const score = (kind === "search" && it.score != null) ? `<span class="score">${it.score}</span>` : "";
    cell.innerHTML = `
      <div class="corner"></div>${score}
      <img loading="lazy" src="${it.image_url}" alt="" onerror="cocoFallback(this)" />
      <div class="meta"><span class="id">${it.image_id}</span><span class="cap">${escapeHtml(cap)}</span></div>`;
    cell.onclick = () => openDrawer(it, kind);
    container.appendChild(cell);
  });
}
const escapeHtml = s => String(s).replace(/[&<>"]/g, c => ({ "&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;" }[c]));

// COCO ids may live in val2014 or train2014 — try the other folder before giving up.
window.cocoFallback = function (img) {
  if (!img.dataset.tried && /val2014/.test(img.src)) {
    img.dataset.tried = "1";
    img.src = img.src.replace(/val2014/g, "train2014");
  } else {
    img.style.opacity = ".15";
  }
};

/* ---------- drawer ---------- */
const drawer = $("#drawer"), scrim = $("#scrim");
function openDrawer(it, kind) {
  $("#drawerImg").src = it.image_url;
  const rows = [["Image id", it.image_id]];
  if (it.question) rows.push(["Question", it.question]);
  if (it.answers?.length) rows.push(["Ground-truth answers", it.answers.join(" · ")]);
  if (it.caption) rows.push(["Caption", it.caption]);
  if (it.score != null) rows.push(["Score", it.score]);
  if (it.ocr_text) rows.push(["Extracted OCR text", it.ocr_text || "—"]);
  if (it.language) rows.push(["Language", it.language]);
  $("#drawerBody").innerHTML = rows.map(([k, v]) => `<dt>${k}</dt><dd>${escapeHtml(v)}</dd>`).join("");
  drawer.classList.add("on"); scrim.classList.add("on");
}
const closeDrawer = () => { drawer.classList.remove("on"); scrim.classList.remove("on"); };
$("#drawerX").onclick = closeDrawer; scrim.onclick = closeDrawer;

/* ---------- file dropzones (one image, shared by Visual QA + OCR Lab) ---------- */
let sharedImage = null;
function setSharedImage(f) {
  if (!f) return;
  sharedImage = f;
  const url = URL.createObjectURL(f);
  ["#vqaDrop", "#ocrDrop"].forEach(sel => {
    const z = $(sel); if (!z) return;
    z.classList.add("has-img");
    let img = $("img", z);
    if (!img) { img = document.createElement("img"); z.appendChild(img); }
    img.src = url;
  });
}
function wireDrop(zoneSel, inputSel) {
  const zone = $(zoneSel), input = $(inputSel);
  zone.onclick = () => input.click();
  input.onchange = () => setSharedImage(input.files[0]);
  ["dragover", "dragenter"].forEach(ev => zone.addEventListener(ev, e => { e.preventDefault(); zone.classList.add("drag"); }));
  ["dragleave", "drop"].forEach(ev => zone.addEventListener(ev, e => { e.preventDefault(); zone.classList.remove("drag"); }));
  zone.addEventListener("drop", e => { if (e.dataTransfer.files[0]) setSharedImage(e.dataTransfer.files[0]); });
}
wireDrop("#vqaDrop", "#vqaFile");
wireDrop("#ocrDrop", "#ocrFile");

/* ---------- RETRIEVE ---------- */
$("#searchK").oninput = e => $("#searchKv").textContent = e.target.value;
async function runSearch() {
  const q = $("#searchInput").value.trim();
  if (!q) { $("#searchStatus").textContent = "enter a query."; return; }
  busy(true, "RETRIEVE"); $("#searchGrid").classList.add("loading"); $("#searchStatus").textContent = "scanning index…";
  const res = await postJSON("/api/search", { query: q, mode: segVal($("#searchMode")), top_k: +$("#searchK").value });
  busy(false); $("#searchGrid").classList.remove("loading");
  if (res.error) { $("#searchStatus").textContent = "⚠ " + res.error; return; }
  renderGrid($("#searchGrid"), res.results || [], "search");
  $("#searchStatus").textContent = `${(res.results || []).length} frames · mode ${segVal($("#searchMode"))}`;
  readout.textContent = "RETRIEVE · " + (res.results || []).length + " hits";
}
$("#searchGo").onclick = runSearch;
$("#searchInput").addEventListener("keydown", e => { if (e.key === "Enter") runSearch(); });

/* ---------- VISUAL QA ---------- */
$("#vqaGo").onclick = async () => {
  const f = sharedImage; const q = $("#vqaQ").value.trim();
  if (!f) { $("#vqaAnswer").textContent = "drop an image first."; return; }
  if (!q) { $("#vqaAnswer").textContent = "ask a question first."; return; }
  busy(true, "VISUAL QA"); $("#vqaAnswer").classList.add("loading"); $("#vqaAnswer").textContent = "looking…";
  const fd = new FormData(); fd.append("image", f); fd.append("question", q); fd.append("system_prompt", $("#vqaSys").value);
  const res = await postForm("/api/vqa", fd);
  busy(false); $("#vqaAnswer").classList.remove("loading");
  $("#vqaAnswer").textContent = res.error ? "⚠ " + res.error : res.answer;
};

/* ---------- OCR LAB ---------- */
$("#ocrGo").onclick = async () => {
  const f = sharedImage; const q = $("#ocrQ").value.trim();
  if (!f) { $(".val", $("#ocrText")).textContent = "drop an image first."; return; }
  if (!q) { $(".val", $("#ocrText")).textContent = "ask a question first."; return; }
  busy(true, "OCR LAB");
  $("#ocrText").classList.add("loading");
  $(".val", $("#ocrText")).textContent = "reading…"; $("#ocrAnsOcr").textContent = "…"; $("#ocrAnsVis").textContent = "…";
  const fd = new FormData(); fd.append("image", f); fd.append("question", q); fd.append("lang", segVal($("#ocrLang")));
  const res = await postForm("/api/ocr_qa", fd);
  busy(false); $("#ocrText").classList.remove("loading");
  if (res.error) { $(".val", $("#ocrText")).textContent = "⚠ " + res.error; return; }
  $(".val", $("#ocrText")).textContent = res.ocr_text || "(no text detected)";
  $("#ocrAnsOcr").textContent = res.answer_ocr || "—";
  $("#ocrAnsVis").textContent = res.answer_vision || "—";
};

/* ---------- INDEX BROWSER ---------- */
async function browse(mode) {
  const ds = segVal($("#browseDs"));
  busy(true, "INDEX BROWSER"); $("#browseGrid").classList.add("loading");
  let res;
  if (mode === "search") {
    const q = $("#browseQ").value.trim();
    if (!q) { $("#browseStatus").textContent = "enter text to search."; busy(false); $("#browseGrid").classList.remove("loading"); return; }
    $("#browseStatus").textContent = "searching OCR text…";
    res = await postJSON("/api/dataset/search", { dataset: ds, query: q, k: 12 });
  } else {
    $("#browseStatus").textContent = "loading index…";
    res = await fetch(`/api/dataset?dataset=${ds}&limit=12`).then(r => r.json());
  }
  busy(false); $("#browseGrid").classList.remove("loading");
  if (res.error) { $("#browseStatus").textContent = "⚠ " + res.error; return; }
  const items = res.items || [];
  renderGrid($("#browseGrid"), items, "browse");
  $("#browseStatus").textContent = items.length ? `${items.length} records · ${ds}` : "no records (is the index built?)";
}
$("#browseLoad").onclick = () => browse("load");
$("#browseGo").onclick = () => browse("search");
$("#browseQ").addEventListener("keydown", e => { if (e.key === "Enter") browse("search"); });

/* ---------- AGENT ---------- */
function agentSession() {
  let s = localStorage.getItem("vlps_session");
  if (!s) { s = (crypto.randomUUID ? crypto.randomUUID() : "s" + Date.now() + Math.random()); localStorage.setItem("vlps_session", s); }
  return s;
}
$("#agentGo").onclick = runAgent;
$("#agentQ").addEventListener("keydown", e => { if (e.key === "Enter") runAgent(); });
$("#agentReset").onclick = async () => {
  const old = agentSession();
  await postJSON("/api/agent/reset", { session: old });
  localStorage.removeItem("vlps_session");
  $("#chatLog").innerHTML = "";
};
async function runAgent() {
  const q = $("#agentQ").value.trim(); if (!q) return;
  const log = $("#chatLog");
  log.insertAdjacentHTML("beforeend", `<div class="msg user">${escapeHtml(q)}</div>`);
  $("#agentQ").value = "";
  const think = document.createElement("div"); think.className = "msg think"; think.textContent = "orchestrating tools…";
  log.appendChild(think); log.scrollTop = log.scrollHeight;
  busy(true, "AGENT");
  const res = await postJSON("/api/agent", { message: q, session: agentSession() });
  busy(false); think.remove();
  let html = escapeHtml(res.error ? "⚠ " + res.error : res.answer);
  if (res.tools && res.tools.length) {
    html += `<div class="tool-trace">` + res.tools.map(t => {
      const a = t.args || {};
      const detail = a.query ? `"${a.query}"` : (a.image_id ? `#${a.image_id}` : "");
      return `<span class="tchip">${escapeHtml(t.tool)}${detail ? " · " + escapeHtml(detail) : ""}</span>`;
    }).join("") + `</div>`;
  }
  if (res.images && res.images.length) {
    html += `<div class="msg-imgs">` + res.images.map(im =>
      `<figure><img src="${im.image_url}" onerror="cocoFallback(this)" />` +
      `<figcaption>${escapeHtml(im.image_id)}${im.caption ? " · " + escapeHtml(im.caption) : ""}</figcaption></figure>`
    ).join("") + `</div>`;
  }
  log.insertAdjacentHTML("beforeend", `<div class="msg bot">${html}</div>`);
  log.scrollTop = log.scrollHeight;
}

/* ---------- boot ---------- */
fetch("/api/config").then(r => r.json()).then(c => {
  $("#modelBadge").textContent = "◍ " + (c.model || "online");
  $("#modelBadge").classList.add("live");
  $("#footModel").textContent = (c.model || "") + " · " + (c.user || "");
}).catch(() => { $("#modelBadge").textContent = "◍ offline"; });
selectTool("retrieve");