"use strict";

const $ = (id) => document.getElementById(id);
const state = {index:null, episode:null, activeStep:null, query:null, tab:"trajectory", cache:new Map(), loadId:0};
const number = (x, digits=3) => typeof x === "number" ? x.toFixed(digits) : "—";
const text = (tag, value, cls) => {
  const node = document.createElement(tag);
  node.textContent = value;
  if (cls) node.className = cls;
  return node;
};
function node(tag, cls, children=[]) {
  const element = document.createElement(tag);
  if (cls) element.className = cls;
  for (const child of children) if (child) element.append(child);
  return element;
}
function button(label, cls, click) {
  const element = text("button", label, cls);
  element.type = "button";
  element.addEventListener("click", click);
  return element;
}
function badge(label, color="") { return text("span", label, "badge " + color); }
function stat(label, value, detail, color="") {
  return node("div", "stat" + (color === "green" ? " highlight" : ""), [
    text("div", label, "stat-label"), text("div", value, "stat-value " + color),
    text("div", detail, "stat-detail")
  ]);
}
function heading(title, subtitle) {
  return node("div", "panel-heading", [text("h3", title), text("small", subtitle)]);
}
function showError(error) {
  $("error").hidden = false;
  $("error").textContent = "Could not load this view. " + error.message;
}
function filteredEpisodes() {
  const search = $("search").value.toLowerCase().trim();
  return state.index.episodes.filter((e) =>
    (!search || (e.source + " " + e.target).toLowerCase().includes(search)) &&
    (!$("split-filter").value || e.split === $("split-filter").value) &&
    (!$("outcome-filter").value || e.success === ($("outcome-filter").value === "success"))
  );
}
function renderList() {
  const episodes = filteredEpisodes();
  $("list-count").textContent = episodes.length;
  $("episode-list").replaceChildren();
  for (const episode of episodes) {
    const item = button("", "episode-button" + (state.episode?.id === episode.id ? " selected" : ""), () => loadEpisode(episode.id));
    item.dataset.episode = episode.id;
    item.append(node("span", "episode-route", [
      text("strong", episode.source), text("small", "↳ " + episode.target)
    ]));
    item.append(node("span", "episode-meta", [
      node("span", "", [node("i", "status-dot" + (episode.success ? " success" : "")), text("span", episode.success ? "Success" : episode.termination)]),
      text("span", episode.labels + " labels · " + episode.queries + " Q calls")
    ]));
    $("episode-list").append(item);
  }
  if (!episodes.length) $("episode-list").append(text("div", "No matching trajectories.", "empty"));
}
async function loadEpisode(id, queryIndex=null) {
  const generation = ++state.loadId;
  const snapshot = state.index.summary.created_at;
  const cacheKey = snapshot + "/" + id;
  $("error").hidden = true;
  try {
    if (!state.cache.has(cacheKey)) {
      $("trajectory-view").replaceChildren(text("div", "Loading trajectory…", "loading"));
      const response = await fetch("data/episodes/" + encodeURIComponent(id) + ".json.gz?snapshot=" + encodeURIComponent(snapshot));
      if (!response.ok) throw new Error("Trajectory request returned " + response.status);
      let bytes = new Uint8Array(await response.arrayBuffer());
      if (bytes[0] === 31 && bytes[1] === 139) {
        const stream = new Blob([bytes]).stream().pipeThrough(new DecompressionStream("gzip"));
        bytes = new Uint8Array(await new Response(stream).arrayBuffer());
      }
      state.cache.set(cacheKey, JSON.parse(new TextDecoder().decode(bytes)));
    }
    if (generation !== state.loadId) return;
    state.episode = state.cache.get(cacheKey);
    const queries = Object.values(state.episode.queries);
    const selected = queryIndex !== null ? state.episode.queries[String(queryIndex)] :
      queries.find((q) => q.target_value !== null) || queries[0];
    state.activeStep = selected ? state.episode.steps.find((s) => s.query_index === selected.query_index) : state.episode.steps[0];
    state.query = selected || null;
    setTab("trajectory");
    renderTrajectory();
    renderList();
  } catch (error) { showError(error); }
}
function actionTitle(step) {
  const action = step.action;
  if (step.forced) return "Forced " + action.type + (action.link_id ? " · L" + action.link_id : "");
  if (action.type === "estimate_q") return "Q call" + (step.query_index ? " #" + step.query_index : "") + " · " +
    (action.action?.type === "back" ? "back" : "L" + (action.action?.link_id ?? "?"));
  if (action.type === "click") return "Click · L" + action.link_id;
  if (action.type === "read_chunk") return "Read chunk " + action.index;
  return action.type === "back" ? "Go back" : action.type;
}
function icon(step) { return step.action.type === "estimate_q" ? "Q" : step.action.type === "read_chunk" ? "≡" : step.action.type === "back" ? "↶" : "↗"; }
function updateHash() {
  const params = new URLSearchParams({episode:state.episode.id});
  if (state.query) params.set("query", state.query.query_index);
  history.replaceState(null, "", "#" + params);
}
function selectStep(step) {
  state.activeStep = step;
  state.query = step.query_index ? state.episode.queries[String(step.query_index)] : null;
  for (const item of document.querySelectorAll(".step-button")) {
    item.classList.toggle("active", Number(item.dataset.decision) === step.decision);
  }
  renderStats();
  renderDetail();
  updateHash();
}
function renderTrajectory() {
  const episode = state.episode;
  const title = node("h2", "", [text("span", episode.source), text("span", "→", "route-arrow"), text("span", episode.target)]);
  const header = node("div", "episode-header", [
    node("div", "", [title, node("div", "episode-subtitle", [
      text("span", episode.split === "train" ? "Training split" : "Validation split"),
      text("span", "Seed " + episode.seed), text("span", episode.navigation_steps + " / " + (state.index.summary.horizon ?? 10) + " navigation moves"),
      text("span", episode.decisions + " decisions")
    ])]), badge(episode.success ? "✓ Goal reached" : episode.termination, episode.success ? "green" : "")
  ]);
  const stats = node("div", "stats"); stats.id = "episode-stats";
  const timeline = node("div", "timeline");
  for (const step of episode.steps) {
    const query = step.query_index ? episode.queries[String(step.query_index)] : null;
    const item = button("", "step-button" + (state.activeStep === step ? " active" : ""), () => selectStep(step));
    item.dataset.decision = step.decision;
    const copy = node("div", "step-copy", [
      text("div", actionTitle(step), "step-title"), text("div", step.after?.page || "", "step-page")
    ]);
    if (query) copy.append(node("div", "step-values", [
      badge("Pred " + number(query.predicted_value), "purple"),
      badge(query.target_value !== null ? "MC " + number(query.target_value) : "No MC target", query.target_value !== null ? "green" : "")
    ]));
    item.append(text("span", icon(step), "step-icon" + (query ? " query" : "")), copy, text("span", String(step.decision + 1).padStart(2, "0"), "step-number"));
    timeline.append(item);
  }
  const left = node("section", "panel", [heading("Parent trajectory", episode.steps.length + " actions"), timeline]);
  const right = node("section", "panel"); right.id = "detail-panel";
  $("trajectory-view").replaceChildren(header, stats, node("div", "explorer-grid", [left, right]));
  renderStats(); renderDetail(); updateHash();
}
function renderStats() {
  const q = state.query, e = state.episode;
  $("episode-stats").replaceChildren(
    stat("Prompted prediction", q ? number(q.predicted_value) : "—", q ? "Q call #" + q.query_index : "Select a Q call", "purple"),
    stat("Monte Carlo target", q ? number(q.target_value) : "—", q?.target_value !== null && q ? q.successes + " / " + q.samples + " successful branches" : "Selected calls only", "green"),
    stat("Absolute difference", q?.predicted_value != null && q.target_value != null ? number(Math.abs(q.predicted_value - q.target_value)) : "—", "Prediction vs. MC fraction"),
    stat("Labeled Q calls", Object.values(e.queries).filter((q) => q.target_value !== null).length + " / " + Object.keys(e.queries).length, "At most two labels per trajectory")
  );
}
function bar(label, value, mc=false) {
  const fill = node("div", "fill" + (mc ? " mc" : ""));
  fill.style.width = (value === null ? 0 : value * 100) + "%";
  return node("div", "value-bar", [text("span", label), node("div", "track", [fill]), text("strong", number(value))]);
}
function renderDetail() {
  const panel = $("detail-panel"), q = state.query, step = state.activeStep;
  if (!step) { panel.replaceChildren(text("div", "No recorded actions.", "empty")); return; }
  const selector = node("select", "query-selector");
  selector.setAttribute("aria-label", "Select Q call");
  selector.append(new Option("Select a Q call", ""));
  for (const query of Object.values(state.episode.queries)) {
    selector.append(new Option("Q #" + query.query_index + (query.target_value !== null ? " · MC " + number(query.target_value) : " · no target"), String(query.query_index)));
  }
  selector.value = q ? String(q.query_index) : "";
  selector.addEventListener("change", () => {
    const selected = state.episode.steps.find((s) => String(s.query_index) === selector.value);
    if (selected) selectStep(selected);
  });
  const top = node("div", "panel-heading", [text("h3", q ? "Value call #" + q.query_index : "Decision " + (step.decision + 1)), selector]);
  const body = node("div", "query-detail", [
    text("div", q ? "PROPOSED ACTION · BEFORE NAVIGATION" : "RECORDED ACTION", "section-kicker"),
    text("h3", q ? (q.action?.type === "back" ? "Back" : "Click L" + (q.action?.link_id ?? "?")) +
      (q.candidate ? " → " + q.candidate.destination : "") : actionTitle(step), "query-title"),
    text("p", "At " + (step.before?.page || step.after.page) + " · " + (step.before?.remaining_steps ?? "—") + " navigation moves remaining", "description")
  ]);
  if (q) {
    body.append(node("div", "value-bars", [bar("Prompted Q", q.predicted_value), bar("MC target", q.target_value, true)]));
    body.append(text("p", q.target_value !== null ?
      "The parent received the prompted estimate. Independent MC branches supplied the training target afterward." :
      "This call was not selected for MC labeling. Its prompted estimate is shown for context; it is not a training target.", "description"));
  }
  panel.replaceChildren(top, body);
  if (q?.target_value !== null && q) renderBranches(panel, q);
  const inspector = node("section", "inspector"); inspector.id = "inspector";
  panel.append(inspector);
  inspectStep(step, "Parent decision " + (step.decision + 1));
}
function renderBranches(panel, query) {
  const header = heading("Continuation branches", query.successes + " of " + query.samples + " reached the goal");
  header.classList.add("branch-header");
  const area = node("div", "branch-area", [
    text("div", state.activeStep.before.page + " · Q #" + query.query_index, "fork-root"),
    text("div", "Force proposed action, then follow the prompted-Q policy", "fork-label")
  ]);
  const lanes = node("div", "branch-lanes");
  for (const branch of query.branches) {
    const outcome = node("div", "branch-outcome", [
      text("strong", "Branch " + (branch.index + 1)),
      badge(branch.success ? "✓ Success" : "× " + branch.termination, branch.success ? "green" : "red"),
      text("small", "Seed " + branch.seed)
    ]);
    const flow = node("div", "branch-flow");
    for (const [i, step] of branch.steps.entries()) {
      if (i) flow.append(text("span", "→", "flow-arrow"));
      const item = button("", "branch-node" + (step.forced ? " forced" : step.action.type === "estimate_q" ? " tool" : ""), () => {
        document.querySelectorAll(".branch-node.selected").forEach((n) => n.classList.remove("selected"));
        item.classList.add("selected");
        inspectStep(step, "Branch " + (branch.index + 1) + " · decision " + (step.decision + 1));
      });
      item.append(text("small", actionTitle(step)), text("strong", step.after?.page || ""));
      if (step.action.type === "estimate_q") item.append(badge("Pred " + number(step.predicted_value), "purple"));
      flow.append(item);
    }
    if (!branch.steps.length) flow.append(text("span", "Terminal at query state: " + branch.termination, "description"));
    lanes.append(node("div", "branch-row", [outcome, flow]));
  }
  area.append(lanes, text("div", "Purple nodes are prompted Q calls. Green-bordered nodes are forced actions. Click any node to inspect its observation and tool result. Branch Q calls never launch more MC.", "branch-legend"));
  panel.append(header, area);
}
function inspectStep(step, label) {
  const inspector = $("inspector");
  const observation = state.episode.observations[step.observation];
  inspector.replaceChildren(
    text("h4", label),
    node("div", "inspector-grid", [
      node("div", "", [text("div", "Action", "section-kicker"), text("pre", JSON.stringify(step.action, null, 2))]),
      node("div", "", [text("div", "Result", "section-kicker"), text("pre", JSON.stringify(step.result, null, 2))])
    ])
  );
  const states = node("details", "", [text("summary", "Game state and remaining budgets"), text("pre", JSON.stringify({before:step.before, after:step.after}, null, 2))]);
  inspector.append(states);
  if (observation) {
    const detail = node("details", "", [text("summary", "Visible article · " + observation.page + (observation.chunk !== undefined ? " · chunk " + observation.chunk : ""))]);
    if (observation.url?.startsWith("https://en.wikipedia.org/wiki/")) {
      const link = text("a", "View source article on Wikipedia ↗");
      link.href = observation.url; link.target = "_blank"; link.rel = "noopener";
      detail.append(link);
    }
    detail.append(text("pre", observation.text || "", "article"));
    inspector.append(detail);
  }
}
function svgNode(tag, attributes={}) {
  const element = document.createElementNS("http://www.w3.org/2000/svg", tag);
  for (const [key, value] of Object.entries(attributes)) element.setAttribute(key, value);
  return element;
}
function renderCalibration() {
  const ids = new Set(filteredEpisodes().map((e) => e.id));
  const points = state.index.comparisons.filter((p) => ids.has(p.episode_id));
  const mae = points.length ? points.reduce((a, p) => a + Math.abs(p.predicted - p.target), 0) / points.length : null;
  const mse = points.length ? points.reduce((a, p) => a + (p.predicted - p.target) ** 2, 0) / points.length : null;
  const stats = node("div", "stats", [
    stat("Compared calls", points.length, "Complete MC labels"),
    stat("Mean absolute error", number(mae, 3), "Against MC fractions"),
    stat("Mean squared error", number(mse, 3), "Against noisy MC targets"),
    stat("Continuations", state.index.summary.samples, "Per selected call", "green")
  ]);
  const chart = svgNode("svg", {viewBox:"0 0 600 450", class:"scatter", role:"img", "aria-label":"Prompted predictions versus Monte Carlo targets"});
  for (const t of [0, .25, .5, .75, 1]) {
    const x = 65 + t * 475, y = 365 - t * 315;
    chart.append(svgNode("line", {x1:65, x2:540, y1:y, y2:y, stroke:"#e5eae0"}));
    chart.append(svgNode("line", {x1:x, x2:x, y1:50, y2:365, stroke:"#edf0e9"}));
    const xt = svgNode("text", {x:x, y:388, "text-anchor":"middle"}); xt.textContent = t.toFixed(2); chart.append(xt);
    const yt = svgNode("text", {x:50, y:y+4, "text-anchor":"end"}); yt.textContent = t.toFixed(2); chart.append(yt);
  }
  chart.append(svgNode("line", {x1:65, y1:365, x2:540, y2:50, stroke:"#9bb397", "stroke-dasharray":"5 5"}));
  for (const p of points) {
    const circle = svgNode("circle", {cx:65+p.predicted*475, cy:365-p.target*315, r:5.5, class:"point", fill:p.split === "train" ? "#4a987e" : "#9270b8", tabindex:0, role:"button", "aria-label":"Open Q call " + p.query_index + ", predicted " + p.predicted + ", target " + p.target});
    const title = svgNode("title"); title.textContent = "Predicted " + number(p.predicted) + " · target " + number(p.target) + " · " + p.split;
    circle.append(title);
    circle.addEventListener("click", () => loadEpisode(p.episode_id, p.query_index));
    circle.addEventListener("keydown", (e) => { if (e.key === "Enter") loadEpisode(p.episode_id, p.query_index); });
    chart.append(circle);
  }
  const xlabel = svgNode("text", {x:302, y:427, "text-anchor":"middle"}); xlabel.textContent = "Prompted Q prediction";
  const ylabel = svgNode("text", {transform:"translate(17 210) rotate(-90)", "text-anchor":"middle"}); ylabel.textContent = "MC target · observed success fraction";
  chart.append(xlabel, ylabel);
  const plot = node("section", "panel", [heading("Predictions vs. MC targets", "Click a point to inspect its call"), node("div", "plot-wrap", [chart]), text("div", "Green: training · Purple: validation · Dashed line: agreement. Points can overlap at identical values.", "chart-caption")]);
  const samples = state.index.summary.samples;
  const bins = Array.from({length:samples + 1}, (_, i) => ({value:i / samples, count:points.filter((p) => p.target === i / samples).length}));
  const max = Math.max(1, ...bins.map((b) => b.count));
  const histogram = node("div", "histogram");
  for (const bin of bins) {
    const fill = node("div", "fill"); fill.style.width = bin.count / max * 100 + "%";
    histogram.append(node("div", "histogram-row", [text("span", number(bin.value)), node("div", "track", [fill]), text("strong", bin.count)]));
  }
  const notes = node("div", "method-note", [
    text("strong", "Read these as sample outcomes."),
    text("p", "Each target averages " + samples + " terminal success outcomes. A target of 0 or 1 is not proof that the true probability is exactly 0 or 1."),
    text("p", "Predictions come from the prompted Q tool. No trained scalar value model is represented here. Only sampled calls contribute training targets."),
    text("p", "The MC continuation policy can use prompted Q tools; the unchanged historical estimator prompt describes tool-free continuations.")
  ]);
  $("calibration-view").replaceChildren(stats, node("div", "metrics-grid", [plot, node("section", "panel", [heading("MC target distribution", points.length + " calls"), histogram, notes])]));
}
function setTab(tab) {
  state.tab = tab;
  $("trajectory-view").hidden = tab !== "trajectory";
  $("calibration-view").hidden = tab !== "calibration";
  $("trajectory-tab").classList.toggle("active", tab === "trajectory");
  $("calibration-tab").classList.toggle("active", tab === "calibration");
  if (tab === "calibration") renderCalibration();
}
async function initialize() {
  try {
    $("error").hidden = true;
    const response = await fetch("data/index.json", {cache:"no-store"});
    if (!response.ok) throw new Error("Snapshot request returned " + response.status);
    state.index = await response.json();
    state.cache.clear();
    const s = state.index.summary;
    $("collection-info").replaceChildren(text("div", s.episodes + " trajectories · " + s.labels + " labels"), text("small", "Qwen3.5-9B · " + s.samples + " MC continuations"));
    $("snapshot-badge").textContent = s.collection_complete ? "Collection complete" : "Snapshot · " + s.episodes + " / " + s.planned_episodes + " trajectories";
    $("snapshot-badge").classList.toggle("live", !s.collection_complete);
    $("dataset-link").href = "https://huggingface.co/datasets/" + s.dataset_repo;
    $("footer").textContent = "Snapshot " + new Date(s.created_at).toLocaleString() + " · Predictions are prompted Q estimates. MC targets are empirical success fractions. Wikipedia excerpts are CC BY-SA 4.0.";
    renderList();
    const hash = new URLSearchParams(location.hash.slice(1));
    const id = state.index.episodes.some((e) => e.id === hash.get("episode")) ? hash.get("episode") :
      state.index.episodes.find((e) => e.labels)?.id || state.index.episodes[0]?.id;
    if (id) await loadEpisode(id, hash.has("query") ? Number(hash.get("query")) : null);
    else $("trajectory-view").replaceChildren(text("div", "No completed trajectories in this snapshot.", "empty"));
  } catch (error) { showError(error); }
}
$("trajectory-tab").addEventListener("click", () => setTab("trajectory"));
$("calibration-tab").addEventListener("click", () => setTab("calibration"));
$("refresh").addEventListener("click", initialize);
for (const id of ["search", "split-filter", "outcome-filter"]) {
  $(id).addEventListener(id === "search" ? "input" : "change", () => {
    if (!state.index) return;
    renderList(); if (state.tab === "calibration") renderCalibration();
  });
}
initialize();