// app.js — wires the IR model, layout, graph renderer, sidebar and inspector.

import { IRModel, EDGE_KINDS, patternColor } from "./ir.js";
import { Layout } from "./layout.js";
import { Graph, styleForKind } from "./graph.js";
import { compareIR, mergeIR } from "./compare.js";

const state = {
  ir: null,
  rawIR: null, // unmasked artifact JSON (spec masking derives state.ir from this)
  specLevel: 3, // 1 structure · 2 +capabilities · 3 +modularity (full)
  showTP: true, // show tensor-parallel glyphs
  showKernels: true, // show kernel badges
  expanded: new Set(),
  fields: {},
  showInfo: true, // shape/attribute captions on nodes (always on)
  edgeKinds: new Set(EDGE_KINDS), // all edge kinds shown
  showProvenance: true, // provenance in inspector (always on)
  selectedId: null,
  manifest: [], // architecture entries from the manifest
  irCache: new Map(), // artifact path -> IRModel (for comparison)
  compare: false,
  compareMerged: true, // merged union view (when models share a lineage)
};

let graph;

// --- Tiny DOM helper --------------------------------------------------------
function h(tag, attrs = {}, ...children) {
  const node = document.createElement(tag);
  for (const [k, v] of Object.entries(attrs)) {
    if (v == null) continue;
    if (k === "class") node.className = v;
    else if (k === "html") node.innerHTML = v;
    else if (k.startsWith("on") && typeof v === "function") node.addEventListener(k.slice(2), v);
    else node.setAttribute(k, v);
  }
  for (const c of children.flat()) {
    if (c == null) continue;
    node.appendChild(typeof c === "string" ? document.createTextNode(c) : c);
  }
  return node;
}

const $ = (sel) => document.querySelector(sel);

// --- Loading ----------------------------------------------------------------

// The generator writes straight here — the app reads it directly, no copying.
const IR_BASE = "ir/";

// Always refetch IR from disk so regenerated artifacts show on a plain reload
// (no hard-refresh / cache reset needed).
function fetchIR(url) {
  return fetch(url, { cache: "no-store" });
}

// Mask a raw artifact to a spec tier: 1 structure · 2 +capabilities ·
// 3 +modularity · 4 full. Lets the viewer preview conformance to a single tier.
function applySpec(raw, level) {
  const r = { ...raw };
  if (level < 2) delete r.capabilities; // tier 2
  if (level < 3) {
    delete r.extends; // tier 3
    delete r.patches;
  }
  return r;
}

// --- Hub org avatars (for kernel badges) ------------------------------------
const avatarCache = new Map(); // org -> avatarUrl | null

// The distinct orgs owning any kernel repo in this artifact.
function kernelOrgs(ir) {
  const set = new Set();
  const map = ir && ir.capabilities && ir.capabilities.kernels;
  if (map) for (const repos of Object.values(map)) (repos || []).forEach((r) => set.add(String(r).split("/")[0]));
  return [...set];
}

// Fetch missing org avatar URLs, then re-render so the graph picks them up.
async function ensureAvatars(orgs, onReady) {
  const missing = orgs.filter((o) => o && !avatarCache.has(o));
  if (!missing.length) return;
  await Promise.all(
    missing.map(async (org) => {
      for (const kind of ["organizations", "users"]) {
        try {
          const res = await fetch(`https://huggingface.co/api/${kind}/${org}/avatar`);
          if (res.ok) {
            const j = await res.json();
            if (j && j.avatarUrl) return void avatarCache.set(org, j.avatarUrl);
          }
        } catch (_) {}
      }
      avatarCache.set(org, null);
    })
  );
  onReady && onReady();
}

async function loadManifest() {
  const sel = $("#arch-select");
  sel.replaceChildren();
  try {
    const res = await fetchIR(IR_BASE + "manifest.json");
    if (!res.ok) throw new Error(res.status);
    const manifest = await res.json();
    state.manifest = manifest.architectures || [];
    state.manifest.forEach((a) => {
      const label = `${a.model_type}  —  ${a.model_class || ""}`.trim();
      sel.appendChild(h("option", { value: a.artifact }, label));
    });
    populateCompareSelectors();
    $("#manifest-note").textContent = `${state.manifest.length} architectures · schema ${manifest.artifact_schema_version || "?"}`;
    if (sel.options.length) {
      await loadArtifact(sel.value);
    }
  } catch (err) {
    $("#manifest-note").textContent = "No manifest found — upload or paste an IR artifact to begin.";
  }
}

// Manifest artifact paths (e.g. "artifacts/llama.json") are resolved relative
// to the manifest's directory (IR_BASE).
async function loadArtifact(path) {
  const url = IR_BASE + path;
  try {
    const res = await fetchIR(url);
    if (!res.ok) throw new Error(res.status);
    const raw = await res.json();
    setIR(raw);
  } catch (err) {
    setStatus(`Failed to load ${url}: ${err.message}`, true);
  }
}

// Fetch (and cache) an IRModel for a manifest artifact path — used by compare.
// Cached per spec tier so the comparison respects the same masking as the main view.
async function getIR(path) {
  const key = `${path}@${state.specLevel}`;
  if (state.irCache.has(key)) return state.irCache.get(key);
  const res = await fetchIR(IR_BASE + path);
  if (!res.ok) throw new Error(`HTTP ${res.status}`);
  const ir = new IRModel(applySpec(await res.json(), state.specLevel));
  state.irCache.set(key, ir);
  return ir;
}

function setIR(raw) {
  state.rawIR = raw;
  try {
    state.ir = new IRModel(applySpec(raw, state.specLevel));
  } catch (err) {
    setStatus(`Invalid IR: ${err.message}`, true);
    return;
  }
  graph.clearOverrides();
  state.edgeKinds = new Set(state.ir.edgeKinds); // all edge kinds present, shown
  renderArchFacts();
  renderCapabilityChips();
  renderModularity();
  state.expanded = new Set(state.ir.repeatsById.keys()); // repeats expanded by default
  state.fields = { ...state.ir.configFields };
  state.selectedId = state.ir.rootId;
  $("#config-editor").value = JSON.stringify(state.fields, null, 2);
  if (!state.ir.rootId || state.ir.nodes.size === 0) {
    setStatus(`${state.ir.modelType}: empty artifact — no components emitted`, true);
  } else {
    setStatus(`Loaded ${state.ir.modelType} · ${state.ir.nodes.size} components · ${state.ir.edges.length} edges`);
  }
  rebuild(true);
  loadHubModels(state.ir.modelType);
  // Fetch kernel-org avatars, then redraw so the badges show them.
  ensureAvatars(kernelOrgs(state.ir), () => {
    rebuild(false);
    if (state.compare) renderCompare();
  });
}

// Show/hide tier-gated UI (Compare is a modularity-tier feature → tier ≥ 3).
function updateSpecUI() {
  const canCompare = state.specLevel >= 3;
  const btn = $("#compare-toggle");
  if (btn) btn.style.display = canCompare ? "" : "none";
  if (!canCompare && state.compare) toggleCompare(); // leave compare mode
}

// Re-derive the model at a new spec tier, preserving the current view.
function reapplySpec() {
  if (!state.rawIR) return;
  updateSpecUI();
  state.ir = new IRModel(applySpec(state.rawIR, state.specLevel));
  state.edgeKinds = new Set(state.ir.edgeKinds);
  renderArchFacts();
  renderCapabilityChips();
  renderModularity();
  state.expanded = new Set([...state.expanded].filter((id) => state.ir.repeatsById.has(id)));
  if (!state.ir.node(state.selectedId)) state.selectedId = state.ir.rootId;
  rebuild(false);
  ensureAvatars(kernelOrgs(state.ir), () => rebuild(false));
  if (state.compare) renderCompare();
}

// --- Hugging Face Hub integration (fully client-side) -----------------------

function fmtDownloads(n) {
  if (n == null) return "";
  if (n >= 1e6) return `${(n / 1e6).toFixed(1)}M`;
  if (n >= 1e3) return `${(n / 1e3).toFixed(1)}k`;
  return String(n);
}

// Populate the dropdown with top Hub checkpoints for this architecture.
async function loadHubModels(modelType) {
  const sel = $("#hub-select");
  const note = $("#hub-note");
  sel.replaceChildren(h("option", { value: "" }, "— loading checkpoints… —"));
  try {
    const url = `https://huggingface.co/api/models?filter=${encodeURIComponent(
      modelType
    )}&sort=downloads&direction=-1&limit=30`;
    const res = await fetch(url);
    if (!res.ok) throw new Error(`HTTP ${res.status}`);
    const models = await res.json();
    sel.replaceChildren(
      h("option", { value: "" }, `— pick a ${modelType} checkpoint —`)
    );
    models.forEach((m) => {
      const dl = fmtDownloads(m.downloads);
      sel.appendChild(h("option", { value: m.id }, dl ? `${m.id}  ·  ${dl} ↓` : m.id));
    });
    note.innerHTML = `${models.length} public <code>${modelType}</code> checkpoints · loads its <code>config.json</code>`;
  } catch (err) {
    sel.replaceChildren(h("option", { value: "" }, "— Hub unavailable —"));
    note.textContent = `Could not reach the Hub (${err.message}). Edit the config manually below.`;
  }
}

// Fetch a checkpoint's config.json and apply it to resolve the template.
async function loadHubConfig(modelId) {
  const note = $("#hub-note");
  note.textContent = `Fetching config.json for ${modelId}…`;
  try {
    const res = await fetch(`https://huggingface.co/${modelId}/resolve/main/config.json`);
    if (!res.ok) throw new Error(res.status === 401 ? "gated/private" : `HTTP ${res.status}`);
    const cfg = await res.json();
    state.fields = cfg;
    $("#config-editor").value = JSON.stringify(cfg, null, 2);
    note.innerHTML = `Resolved against <code>${modelId}</code>`;
    setStatus(`Resolved template against ${modelId}`);
    renderCapabilityChips(); // checkpoint declares its own architecture class
    rebuild(false);
  } catch (err) {
    note.textContent = `Could not load config for ${modelId} (${err.message}).`;
    setStatus(`Failed to load config for ${modelId}: ${err.message}`, true);
  }
}

function setStatus(msg, isError = false) {
  const el = $("#status");
  el.textContent = msg;
  el.classList.toggle("error", isError);
}

// --- Rebuild / render -------------------------------------------------------

function rebuild(fit = false) {
  if (!state.ir) return;
  const tied = state.ir.tiedEmbeddings(state.fields);
  const layout = new Layout(state.ir, {
    expanded: state.expanded,
    fields: state.fields,
    showInfo: state.showInfo,
    tied,
  });
  graph.render(layout, {
    ir: state.ir,
    fields: state.fields,
    edgeKinds: state.edgeKinds,
    avatars: avatarCache,
    tied,
    // TP / kernels are capability-tier (spec §02): hidden below tier 2.
    showTP: state.showTP && state.specLevel >= 2,
    showKernels: state.showKernels && state.specLevel >= 2,
  });
  graph.setSelected(state.selectedId);
  if (fit) graph.fit();
  renderInspector();
}

// --- Inspector --------------------------------------------------------------

function row(label, value, opts = {}) {
  if (value == null || value === "") return null;
  return h(
    "div",
    { class: "insp-row" },
    h("div", { class: "insp-key" }, label),
    h("div", { class: `insp-val ${opts.mono ? "mono" : ""}` }, String(value))
  );
}

function section(title, ...rows) {
  const kept = rows.flat().filter(Boolean);
  if (!kept.length) return null;
  return h("div", { class: "insp-section" }, h("div", { class: "insp-h" }, title), ...kept);
}

// A key + a wrapped row of pill chips.
function chipsRow(label, items) {
  if (!items || !items.length) return null;
  return h(
    "div",
    { class: "insp-row" },
    h("div", { class: "insp-key" }, label),
    h("div", { class: "cap-chips" }, ...items.map((t) => h("span", { class: "cap-chip" }, String(t))))
  );
}

// Inspector "Attention schedule" section: colored per-layer grid + summary + legend.
function scheduleSection(sched, summary) {
  const grid = h(
    "div",
    { class: "sched-grid" },
    ...sched.map((p, i) =>
      h("span", {
        class: "pat-cell",
        title: `layer ${i}: ${p}`,
        style: `background:${patternColor(p)}`,
      })
    )
  );
  const distinct = [...new Set(sched)];
  const legend = h(
    "div",
    { class: "sched-legend" },
    ...distinct.map((p) =>
      h(
        "span",
        { class: "sched-leg-item" },
        h("span", { class: "pat-cell", style: `background:${patternColor(p)}` }),
        p.replace(/_attention$/, "")
      )
    )
  );
  return h(
    "div",
    { class: "insp-section" },
    h("div", { class: "insp-h" }, `Attention schedule · ${sched.length} layers`),
    h("div", { class: "sched-summary" }, summary || ""),
    grid,
    legend
  );
}

function edgePill(edge, endpointKey) {
  const otherId = endpointKey === "source" ? edge.target : edge.source;
  const style = styleForKind(edge.kind);
  const label = state.ir.label(otherId, state.fields);
  const dir = endpointKey === "source" ? "→" : "←";
  return h(
    "div",
    {
      class: "edge-item",
      title: `${edge.kind}: ${edge.source} → ${edge.target}`,
      onclick: () => select(otherId),
    },
    h("span", { class: "edge-dot", style: `background:${style.color}` }),
    h("span", { class: "edge-dir" }, dir),
    h("span", { class: "edge-name" }, label),
    h("span", { class: "edge-kind" }, edge.kind)
  );
}

function renderInspector() {
  const panel = $("#inspector-body");
  panel.replaceChildren();
  const id = state.selectedId;
  if (!id || !state.ir) {
    panel.appendChild(h("div", { class: "insp-empty" }, "Select a component to inspect it."));
    return;
  }

  const isInput = typeof id === "string" && id.startsWith("input:");
  const n = state.ir.node(id);
  const label = state.ir.label(id, state.fields);

  const add = (elm) => elm && panel.appendChild(elm); // sections may be empty → null

  panel.appendChild(
    h(
      "div",
      { class: "insp-title" },
      h("div", { class: "insp-label" }, label),
      h("div", { class: "insp-id mono" }, id)
    )
  );

  if (isInput) {
    add(section("Overview", row("Kind", "graph input"), row("ID", id, { mono: true })));
  } else if (n) {
    const kind = state.ir.kindLabel(n.kind);
    add(
      section(
        "Overview",
        row("Kind", `${kind}  (${n.kind})`),
        row("Class", n.className, { mono: true }),
        row("Component ID", n.id, { mono: true }),
        row("Node type", n.nodeType)
      )
    );

    if (n.nodeType === "repeat") {
      const r = n.repeat;
      const resolved = state.ir.resolveCount(n, state.fields);
      add(
        section(
          "Symbolic repeat",
          row("Repeated class", r.repeated_class_name, { mono: true }),
          row("Count expr", r.count_expr, { mono: true }),
          row("Resolved count", resolved),
          row("Baked count", r.count),
          row("Index symbol", r.index_symbol, { mono: true }),
          row("Container path", r.container_path_pattern, { mono: true }),
          row("Item path", r.item_path_pattern, { mono: true })
        )
      );
    }

    add(section("Source", row("Path pattern", n.pathPattern, { mono: true })));

    // Per-component semantic facts (current spec: attributes drive cache/route).
    if (n.attributes && typeof n.attributes === "object") {
      const attrRows = Object.entries(n.attributes).map(([k, v]) =>
        row(k, typeof v === "object" ? JSON.stringify(v) : String(v), { mono: true })
      );
      add(section("Attributes", ...attrRows));
    }

    // Tensor-parallel insight: shard style + meaning on a projection; the block's
    // sharding plan + communication cost on an attention/MLP/MoE container.
    if (n.kind === "projection" && n.attributes && n.attributes.tp) {
      const col = n.attributes.tp === "colwise";
      add(
        section(
          "Tensor parallel",
          row("Shard", col ? "column-parallel" : "row-parallel"),
          row(
            "Meaning",
            col
              ? "Output features split across TP ranks — computed locally, no collective."
              : "Input features split across TP ranks — all-reduce combines the partial sums."
          )
        )
      );
    } else {
      const plan = state.ir.tpPlan(id);
      if (plan) {
        add(
          section(
            "Tensor-parallel plan",
            row("Column-parallel", plan.colwise.join(", "), { mono: true }),
            row("Row-parallel", plan.rowwise.join(", "), { mono: true }),
            row("Collective", `${plan.rowwise.length} all-reduce${plan.rowwise.length === 1 ? "" : "s"} per block`)
          )
        );
      }
    }

    // Hub kernel: this node can be augmented by a kernel from the Hub.
    const kern = state.ir.nodeKernel(n);
    if (kern) {
      const orgs = [...new Set(kern.repos.map((r) => r.split("/")[0]))];
      const avatarRow = orgs.length
        ? h(
            "div",
            { class: "insp-row" },
            h("div", { class: "insp-key" }, "Orgs"),
            h(
              "div",
              { class: "org-avatars" },
              ...orgs.map((org) => {
                const url = avatarCache.get(org);
                const a = h("span", {
                  class: "org-avatar-chip",
                  title: org,
                  onclick: () => window.open(`https://huggingface.co/${org}`, "_blank", "noopener"),
                });
                if (url) a.appendChild(h("img", { src: url, alt: org }));
                a.appendChild(h("span", {}, org));
                return a;
              })
            )
          )
        : null;
      const repoRows = kern.repos.map((repo) =>
        h(
          "div",
          {
            class: "edge-item",
            title: `open ${repo} on the Hub`,
            onclick: () => window.open(`https://huggingface.co/${repo}`, "_blank", "noopener"),
          },
          h("span", { class: "kernel-bolt" }, "⚡"),
          h("span", { class: "mono" }, repo)
        )
      );
      add(section("Hub kernel", row("Kernel", kern.name, { mono: true }), avatarRow, ...repoRows));
    }

    // Model-level architecture facts, shown on the root/model node.
    if (n.kind === "model" && state.ir.architecture && typeof state.ir.architecture === "object") {
      const archRows = Object.entries(state.ir.architecture).map(([k, v]) =>
        row(k, typeof v === "object" ? JSON.stringify(v) : String(v), { mono: true })
      );
      add(section("Architecture facts", ...archRows));
    }

    // Capabilities (model node): architecture class, task heads, backends, TP.
    if (n.kind === "model" && state.ir.capabilities) {
      const c = state.ir.capabilities;
      const capRows = [];
      const classes = state.ir.architectureClasses(state.fields);
      if (classes.length) capRows.push(chipsRow("Architecture", classes));
      if (c.task_heads && c.task_heads.length) capRows.push(chipsRow("Task heads", c.task_heads));
      if (c.attention_backends && c.attention_backends.length)
        capRows.push(chipsRow("Attn backends", c.attention_backends.map((b) => b.replace(/_attention$/, ""))));
      if (c.attention_patterns && c.attention_patterns.length)
        capRows.push(chipsRow("Attn patterns", c.attention_patterns.map((p) => p.replace(/_attention$/, ""))));
      if (c.kernels && Object.keys(c.kernels).length)
        capRows.push(chipsRow("Hub kernels", Object.keys(c.kernels)));
      capRows.push(row("Tensor parallel", c.tensor_parallel ? "yes" : "no"));
      add(section("Capabilities", ...capRows));
    }

    // Attention schedule: on the model node (whole schedule) or on a repeat
    // whose layer count matches the schedule length.
    const sched =
      n.kind === "model"
        ? state.ir.capabilities && state.ir.capabilities.attention_schedule
        : n.nodeType === "repeat"
        ? state.ir.scheduleForRepeat(n, state.fields)
        : null;
    if (Array.isArray(sched) && sched.length) {
      add(scheduleSection(sched, state.ir.scheduleSummary(sched)));
    } else if (n.kind === "model" && state.ir.capabilities && state.ir.capabilities.attention_patterns) {
      add(section("Attention schedule", row("Uniform", state.ir.capabilities.attention_patterns.join(", "))));
    }

    // Observed dataflow: whole-model input/output on the model node, and the
    // symbolized in/out shape for any component captured as a forward stage.
    if (n.kind === "model" && state.ir.dataflow) {
      const df = state.ir.dataflow;
      add(
        section(
          "Observed dataflow",
          row("Input", df.input ? `${df.input.name || ""} ${fmtShape(df.input.shape) || ""}`.trim() : null, { mono: true }),
          row("Output", df.output ? fmtShape(df.output.shape) : null, { mono: true }),
          row("Source", df.source, { mono: true })
        )
      );
    }
    const shapes = state.ir.nodeShapes(id);
    if (shapes) {
      add(
        section(
          "Observed shape",
          row("In", fmtShape(shapes.in), { mono: true }),
          row("Out", fmtShape(shapes.out), { mono: true })
        )
      );
    }

    // Config-derived values: surface config fields whose name echoes the kind.
    const cfgRows = configHints(n).map(([k, v]) => row(k, JSON.stringify(v), { mono: true }));
    add(section("Config-derived", ...cfgRows));

    if (state.showProvenance) {
      const p = state.ir.provenance || {};
      add(
        section(
          "Provenance",
          row("Class name", n.className, { mono: true }),
          row("Path pattern", n.pathPattern, { mono: true }),
          row("Model class", p.model_class, { mono: true }),
          row("Model module", p.model_module, { mono: true })
        )
      );
    }
  }

  // Edges (verbatim IR edges incident to this id).
  const { incoming, outgoing } = state.ir.edgesFor(id);
  const edgesBox = [];
  if (incoming.length)
    edgesBox.push(h("div", { class: "edge-group-h" }, `Incoming (${incoming.length})`));
  incoming.forEach((e) => edgesBox.push(edgePill(e, "target")));
  if (outgoing.length)
    edgesBox.push(h("div", { class: "edge-group-h" }, `Outgoing (${outgoing.length})`));
  outgoing.forEach((e) => edgesBox.push(edgePill(e, "source")));
  if (edgesBox.length) {
    panel.appendChild(h("div", { class: "insp-section" }, h("div", { class: "insp-h" }, "Edges"), ...edgesBox));
  }
}

// Heuristic: config fields likely relevant to a component kind.
function configHints(n) {
  const f = state.fields || {};
  const wanted = {
    attention: ["num_attention_heads", "num_heads", "num_key_value_heads", "head_dim", "attention_dropout"],
    cross_attention: ["num_attention_heads", "num_heads", "d_kv"],
    feed_forward: ["intermediate_size", "d_ff", "hidden_act", "dense_act_fn", "is_gated_act"],
    embedding: ["vocab_size", "hidden_size", "d_model", "max_position_embeddings", "type_vocab_size"],
    position: ["max_position_embeddings", "rope_parameters", "relative_attention_num_buckets"],
    normalization: ["rms_norm_eps", "layer_norm_eps", "layer_norm_epsilon"],
    model: ["hidden_size", "d_model", "num_hidden_layers", "num_layers", "vocab_size"],
  };
  const keys = wanted[n.kind] || [];
  return keys.filter((k) => f[k] !== undefined).map((k) => [k, f[k]]);
}

// --- Interactions -----------------------------------------------------------

function select(id) {
  state.selectedId = id;
  graph.setSelected(id);
  renderInspector();
}

function toggle(id) {
  if (state.expanded.has(id)) state.expanded.delete(id);
  else state.expanded.add(id);
  rebuild(false);
}

function setAllRepeats(expand) {
  state.expanded = new Set();
  if (expand) {
    for (const r of state.ir.repeatsById.keys()) state.expanded.add(r);
  }
  rebuild(false);
}

function toggleTheme() {
  const light = document.body.classList.toggle("light");
  try {
    localStorage.setItem("tv-theme", light ? "light" : "dark");
  } catch (_) {}
}

// Compact architecture-facts summary in the sidebar (family/view/positional…).
function renderArchFacts() {
  const note = $("#arch-facts");
  const a = state.ir && state.ir.architecture;
  if (!a || typeof a !== "object") {
    note.textContent = "";
    return;
  }
  const parts = ["family", "view", "positional", "attention_variant"]
    .map((k) => a[k])
    .filter((v) => v != null);
  if (a.is_moe) parts.push("MoE");
  if (a.tie_word_embeddings !== undefined) parts.push(a.tie_word_embeddings ? "tied emb" : "untied emb");
  note.textContent = parts.join(" · ");
}

// Compact capability chips in the sidebar: the model class (checkpoint's own
// architecture when a config is loaded, else base class), TP and MoE.
function renderCapabilityChips() {
  const box = $("#capability-chips");
  box.replaceChildren();
  if (!state.ir) return;
  const c = state.ir.capabilities;
  const chip = (t, cls) => h("span", { class: `cap-chip ${cls || ""}` }, t);
  // Attention backends are shown as tags on the attention node itself.
  state.ir.architectureClasses(state.fields).forEach((cn) => box.appendChild(chip(cn, "arch")));
  if (c && c.tensor_parallel) box.appendChild(chip("TP", "on"));
  if (state.ir.architecture && state.ir.architecture.is_moe) box.appendChild(chip("MoE", "on"));
}

// Surface modular-inheritance info when the model extends another (extends !=
// null); diff magnitude is derived from patches.
function renderModularity() {
  const note = $("#modularity-note");
  const ir = state.ir;
  if (!ir || !ir.extends) {
    note.textContent = "";
    return;
  }
  const bits = [`extends <code>${ir.extends}</code>`];
  if (ir.patches.length) bits.push(`${ir.patches.length} patches`);
  const diff = ir.diffSize();
  if (diff) bits.push(`diff ${diff}`);
  note.innerHTML = `Modular: ${bits.join(" · ")}`;
}

function fmtShape(s) {
  return Array.isArray(s) ? `[${s.join(", ")}]` : s == null ? null : String(s);
}

// --- Architecture comparison ------------------------------------------------

function modelTypeOf(path) {
  const a = state.manifest.find((m) => m.artifact === path);
  return a ? a.model_type : path;
}

// The left pane is always the currently-viewed architecture; the dropdown picks
// the single "…with X" comparison target on the right.
function populateCompareSelectors() {
  const sel = $("#cmp-b");
  if (!sel) return;
  sel.replaceChildren();
  state.manifest.forEach((a) =>
    sel.appendChild(h("option", { value: a.artifact }, a.model_type))
  );
  const cur = $("#arch-select").value;
  const other = state.manifest.find((a) => a.artifact !== cur) || state.manifest[0];
  if (other) sel.value = other.artifact;
}

// Two side-by-side graph panels + one merged union panel, each with its own state.
let cmpPaneA = null;
let cmpPaneB = null;
let cmpMerge = null;

function makePane(svgSelector) {
  const pane = { ir: null, expanded: new Set(), shared: new Set(), graph: null };
  pane.graph = new Graph($(svgSelector), {
    onSelect: (id) => pane.graph.setSelected(id),
    onToggle: (id) => {
      if (pane.expanded.has(id)) pane.expanded.delete(id);
      else pane.expanded.add(id);
      renderPane(pane);
    },
  });
  return pane;
}

function renderPane(pane, fit = false) {
  const tied = pane.ir.tiedEmbeddings(pane.ir.configFields);
  const layout = new Layout(pane.ir, {
    expanded: pane.expanded,
    fields: pane.ir.configFields,
    showInfo: state.showInfo,
    tied,
  });
  pane.graph.render(layout, {
    ir: pane.ir,
    fields: pane.ir.configFields,
    edgeKinds: new Set(pane.ir.edgeKinds),
    sharedIds: pane.shared,
    avatars: avatarCache,
    tied,
    showTP: state.showTP && state.specLevel >= 2,
    showKernels: state.showKernels && state.specLevel >= 2,
  });
  if (fit) pane.graph.fit();
}

// Components present in both models, matched by stable id.
function sharedIdSet(a, b) {
  const s = new Set();
  for (const id of a.nodes.keys()) if (b.nodes.has(id)) s.add(id);
  return s;
}

// The merged union panel: one graph of A ∪ B, coloured by origin.
function makeMergePane() {
  const pane = { ir: null, expanded: new Set(), origin: null, edgeOrigin: null, graph: null };
  pane.graph = new Graph($("#cmp-graph-merged"), {
    onSelect: (id) => pane.graph.setSelected(id),
    onToggle: (id) => {
      if (pane.expanded.has(id)) pane.expanded.delete(id);
      else pane.expanded.add(id);
      renderMergePane(pane);
    },
  });
  return pane;
}

function renderMergePane(pane, fit = false) {
  const tied = pane.ir.tiedEmbeddings(pane.ir.configFields);
  const layout = new Layout(pane.ir, {
    expanded: pane.expanded,
    fields: pane.ir.configFields,
    showInfo: state.showInfo,
    tied,
  });
  pane.graph.render(layout, {
    ir: pane.ir,
    fields: pane.ir.configFields,
    edgeKinds: new Set(pane.ir.edgeKinds),
    avatars: avatarCache,
    origin: pane.origin,
    edgeOrigin: pane.edgeOrigin,
    tied,
    showTP: state.showTP && state.specLevel >= 2,
    showKernels: state.showKernels && state.specLevel >= 2,
  });
  if (fit) pane.graph.fit();
}

function toggleCompare() {
  state.compare = !state.compare;
  document.body.classList.toggle("compare", state.compare);
  $("#compare-toggle").classList.toggle("active", state.compare);
  if (state.compare) renderCompare();
  else graph.fit(); // main graph is visible again — restore its view
}

async function renderCompare() {
  const body = $("#compare-body");
  const av = $("#arch-select").value; // left pane = the architecture being viewed
  // Keep the target distinct from the current model when possible.
  if ($("#cmp-b").value === av) {
    const other = state.manifest.find((a) => a.artifact !== av);
    if (other) $("#cmp-b").value = other.artifact;
  }
  const bv = $("#cmp-b").value;
  $("#cmp-a-label").textContent = modelTypeOf(av);
  body.replaceChildren(h("div", { class: "cmp-msg" }, "Comparing…"));
  try {
    const [a, b] = await Promise.all([getIR(av), getIR(bv)]);
    const diff = compareIR(a, b);
    // Merged union only makes sense when the two share stable ids (same lineage).
    const useMerged = state.compareMerged && diff.regime.basis === "id";
    document.body.classList.toggle("cmp-merged", useMerged);
    $("#cmp-mode").classList.toggle("active", useMerged);
    $("#cmp-mode").disabled = diff.regime.basis !== "id";

    if (useMerged) {
      if (!cmpMerge) cmpMerge = makeMergePane();
      const m = mergeIR(a, b);
      cmpMerge.ir = new IRModel(m.raw);
      cmpMerge.origin = m.origin;
      cmpMerge.edgeOrigin = m.edgeOrigin;
      cmpMerge.expanded = new Set(cmpMerge.ir.repeatsById.keys());
      $("#cmp-merge-legend").innerHTML =
        `<span class="mrg-both">■ shared</span><span class="mrg-a">■ only ${a.modelType}</span><span class="mrg-b">■ only ${b.modelType}</span>`;
      const onlyA = [...m.origin.values()].filter((o) => o === "a").length;
      const onlyB = [...m.origin.values()].filter((o) => o === "b").length;
      $("#cmp-shared-note").innerHTML = `<span class="cmp-shared-dot"></span>${onlyA} only ${a.modelType} · ${onlyB} only ${b.modelType}`;
      renderMergePane(cmpMerge, true);
      ensureAvatars([...kernelOrgs(a), ...kernelOrgs(b)], () => renderMergePane(cmpMerge));
    } else {
      // Side-by-side: lazily create the two panels and sync their pan/zoom.
      if (!cmpPaneA) {
        cmpPaneA = makePane("#cmp-graph-a");
        cmpPaneB = makePane("#cmp-graph-b");
        cmpPaneA.graph.cb.onView = (v) => cmpPaneB.graph.setView(v);
        cmpPaneB.graph.cb.onView = (v) => cmpPaneA.graph.setView(v);
      }
      const shared = sharedIdSet(a, b);
      cmpPaneA.ir = a;
      cmpPaneA.expanded = new Set(a.repeatsById.keys());
      cmpPaneA.shared = shared;
      cmpPaneB.ir = b;
      cmpPaneB.expanded = new Set(b.repeatsById.keys());
      cmpPaneB.shared = shared;
      $("#cmp-a-title").textContent = a.modelType;
      $("#cmp-b-title").textContent = b.modelType;
      $("#cmp-shared-note").innerHTML = shared.size
        ? `<span class="cmp-shared-dot"></span>${shared.size} shared component${shared.size === 1 ? "" : "s"} (by stable id)`
        : "no shared components — different lineages";
      renderPane(cmpPaneA, true);
      renderPane(cmpPaneB, true);
      ensureAvatars([...kernelOrgs(a), ...kernelOrgs(b)], () => {
        renderPane(cmpPaneA);
        renderPane(cmpPaneB);
      });
    }

    body.replaceChildren(...renderDiff(diff));
  } catch (err) {
    body.replaceChildren(h("div", { class: "cmp-msg error" }, `Compare failed: ${err.message}`));
  }
}

// A 3-column diff card: key | A | B, differing rows highlighted. Cell renderers
// return a string or a DOM node (h() accepts both).
function diffCard(title, rows, cellA, cellB) {
  const ca = cellA || ((r) => String(r.a));
  const cb = cellB || ((r) => String(r.b));
  const trs = rows.map((r) =>
    h(
      "div",
      { class: `cmp-row ${r.same ? "" : "cmp-diff"}` },
      h("div", { class: "cmp-k" }, r.field || r.kind || r.label || r.item),
      h("div", { class: "cmp-v cmp-va" }, ca(r)),
      h("div", { class: "cmp-v cmp-vb" }, cb(r))
    )
  );
  return h("section", { class: "cmp-card" }, h("h3", {}, title), ...trs);
}

function renderDiff(d) {
  const out = [];

  // Banner
  const regimeLabel = { lineage: "same lineage", shared_ids: "shared ids", cross_lineage: "cross-lineage" }[d.regime.kind] || d.regime.kind;
  out.push(
    h(
      "div",
      { class: "cmp-banner" },
      h("div", { class: "cmp-title" }, `${d.a.modelType}  vs  ${d.b.modelType}`),
      h("div", { class: "cmp-regime" }, h("span", { class: `cmp-badge regime-${d.regime.kind}` }, regimeLabel), d.regime.note)
    )
  );

  // Column headers
  out.push(
    h(
      "div",
      { class: "cmp-row cmp-head" },
      h("div", { class: "cmp-k" }, ""),
      h("div", { class: "cmp-v cmp-va" }, d.a.modelType),
      h("div", { class: "cmp-v cmp-vb" }, d.b.modelType)
    )
  );

  // ✓ / · cell
  const chk = (on) => h("span", { class: on ? "cmp-yes" : "cmp-no" }, on ? "✓" : "·");

  if (d.headline.length) out.push(diffCard("Architecture", d.headline));
  if (d.scale.length) out.push(diffCard("Scale", d.scale));

  // Capabilities: backend + task-head presence matrices, plus TP/patterns/schedule rows.
  if (d.capabilities) {
    const cap = d.capabilities;
    out.push(diffCard("Capabilities · attention backends", cap.backends, (r) => chk(r.inA), (r) => chk(r.inB)));
    out.push(diffCard("Capabilities · task heads", cap.taskHeads, (r) => chk(r.inA), (r) => chk(r.inB)));
    out.push(diffCard("Capabilities · attention", cap.rows));
  }

  out.push(
    diffCard(
      "Topology (edge kinds)",
      d.topology,
      (r) => chk(r.inA),
      (r) => chk(r.inB)
    )
  );

  // Structure
  if (d.nodes.basis === "kind") {
    out.push(renderKindLanes(d.nodes.lanes));
  } else {
    out.push(renderIdAlign(d.nodes));
  }

  return out;
}

function renderKindLanes(lanes) {
  const rows = lanes.map((l) => {
    const tag = l.onlyA ? h("span", { class: "cmp-tag only-a" }, "A only") : l.onlyB ? h("span", { class: "cmp-tag only-b" }, "B only") : null;
    const deltas = l.attrDeltas.length
      ? h("span", { class: "cmp-deltas" }, "Δ " + l.attrDeltas.map((x) => x.key).join(", "))
      : null;
    return h(
      "div",
      { class: `cmp-row ${l.onlyA || l.onlyB || l.attrDeltas.length || l.countA !== l.countB ? "cmp-diff" : ""}` },
      h("div", { class: "cmp-k" }, l.kind),
      h("div", { class: "cmp-v cmp-va" }, `×${l.countA}`),
      h("div", { class: "cmp-v cmp-vb" }, `×${l.countB}`, tag, deltas)
    );
  });
  return h("section", { class: "cmp-card" }, h("h3", {}, "Components by kind (aligned into lanes)"), ...rows);
}

function renderIdAlign(nodes) {
  const kids = [];
  const changed = nodes.matched.filter((m) => m.changed);
  changed.forEach((m) => {
    const deltas = m.attrDeltas.map((dd) =>
      h("div", { class: "cmp-attr" }, h("span", { class: "cmp-attr-k mono" }, dd.key), h("span", {}, `${dd.a} → ${dd.b}`))
    );
    kids.push(
      h(
        "div",
        { class: "cmp-node cmp-changed" },
        h("div", { class: "cmp-node-h" }, h("span", { class: "mono" }, m.id), m.kindChanged ? h("span", { class: "cmp-tag" }, `${m.kindA}→${m.kindB}`) : null),
        ...deltas
      )
    );
  });
  nodes.added.forEach((n) => kids.push(h("div", { class: "cmp-node cmp-added" }, h("span", { class: "cmp-tag add" }, "added"), h("span", { class: "mono" }, n.id))));
  nodes.removed.forEach((n) => kids.push(h("div", { class: "cmp-node cmp-removed" }, h("span", { class: "cmp-tag del" }, "removed"), h("span", { class: "mono" }, n.id))));
  if (!kids.length) kids.push(h("div", { class: "cmp-msg" }, "No component-level differences."));
  return h("section", { class: "cmp-card" }, h("h3", {}, `Component deltas (${changed.length} changed · ${nodes.added.length} added · ${nodes.removed.length} removed)`), ...kids);
}

// --- Sidebar controls -------------------------------------------------------

function bindControls() {
  $("#arch-select").addEventListener("change", (e) => {
    loadArtifact(e.target.value);
    if (state.compare) renderCompare(); // left pane tracks the current model
  });

  $("#fit-btn").addEventListener("click", () => graph.fit());
  $("#expand-all").addEventListener("click", () => setAllRepeats(true));
  $("#collapse-all").addEventListener("click", () => setAllRepeats(false));
  $("#reset-layout").addEventListener("click", () => graph.resetPositions());
  $("#theme-toggle").addEventListener("click", toggleTheme);
  $("#compare-toggle").addEventListener("click", toggleCompare);
  $("#cmp-b").addEventListener("change", renderCompare);
  $("#cmp-mode").addEventListener("click", () => {
    state.compareMerged = !state.compareMerged;
    renderCompare();
  });
  $("#cmp-fit").addEventListener("click", () => {
    if (state.compareMerged && cmpMerge && document.body.classList.contains("cmp-merged")) cmpMerge.graph.fit();
    if (cmpPaneA) cmpPaneA.graph.fit();
    if (cmpPaneB) cmpPaneB.graph.fit();
  });

  $("#hub-select").addEventListener("change", (e) => {
    if (e.target.value) loadHubConfig(e.target.value);
  });

  // Spec-tier selector.
  $("#spec-seg").addEventListener("click", (e) => {
    const btn = e.target.closest("[data-spec]");
    if (!btn) return;
    state.specLevel = Number(btn.getAttribute("data-spec"));
    [...$("#spec-seg").children].forEach((b) => b.classList.toggle("active", b === btn));
    state.irCache.clear(); // compare cache is keyed by tier; drop stale full-tier entries
    reapplySpec();
  });
  // Feature toggles.
  $("#opt-tp").addEventListener("change", (e) => {
    state.showTP = e.target.checked;
    rebuild(false);
    if (state.compare) renderCompare();
  });
  $("#opt-kernels").addEventListener("change", (e) => {
    state.showKernels = e.target.checked;
    rebuild(false);
    if (state.compare) renderCompare();
  });

  // Config editor.
  $("#config-apply").addEventListener("click", () => {
    try {
      const parsed = JSON.parse($("#config-editor").value);
      state.fields = parsed;
      setStatus("Config applied · repeat counts re-resolved");
      renderCapabilityChips();
      rebuild(false);
    } catch (err) {
      setStatus(`Invalid config JSON: ${err.message}`, true);
    }
  });
  $("#config-reset").addEventListener("click", () => {
    if (!state.ir) return;
    state.fields = { ...state.ir.configFields };
    $("#config-editor").value = JSON.stringify(state.fields, null, 2);
    setStatus("Config reset to artifact defaults");
    renderCapabilityChips();
    rebuild(false);
  });
}

// --- Boot -------------------------------------------------------------------

function boot() {
  // Theme: default dark (HF-style); honour a saved preference.
  try {
    if (localStorage.getItem("tv-theme") === "light") document.body.classList.add("light");
  } catch (_) {}
  graph = new Graph($("#graph"), { onSelect: select, onToggle: toggle });
  bindControls();
  updateSpecUI();
  window.addEventListener("resize", () => graph.fit());
  loadManifest();
}

boot();

export { renderDiff, chipsRow, scheduleSection };