// compare.js — architecture comparison engine.
//
// Consumes two IRModels and produces a structured diff. The viewer owns all
// alignment + delta math (there is no shipped pairwise-diff artifact); it works
// purely from the exports: lineage (extends/modularity), the architecture facts
// block, components/templates + stable ids + attributes, repeats + config, and
// edges.
//
// Regime decision (Step 0):
//   - parent/child via extends / modularity.parent_model  → lineage (align by id)
//   - otherwise measure stable-id overlap:
//       high overlap → shared ids (align by id)
//       low  overlap → cross-lineage (align by semantic kind "lanes")

const HEADLINE_FIELDS = [
  "family",
  "view",
  "attention_variant",
  "positional",
  "is_moe",
  "moe",
  "sliding_window",
  "tie_word_embeddings",
];

// Resolve a possibly-symbolic scalar ("config.hidden_size") against a config.
function resolveVal(v, fields) {
  if (typeof v === "string" && v.startsWith("config.")) {
    const rv = fields ? fields[v.slice("config.".length)] : undefined;
    return rv != null ? rv : v.slice("config.".length);
  }
  return v;
}

function fmt(v) {
  if (v == null) return "—";
  if (typeof v === "object") return JSON.stringify(v);
  return String(v);
}

// Real (non-pseudo) nodes: components + templates + repeats.
function realNodes(ir) {
  const out = new Map();
  for (const [id, n] of ir.nodes) out.set(id, n);
  return out;
}

function diffAttrs(aAttrs, bAttrs, aFields, bFields) {
  aAttrs = aAttrs || {};
  bAttrs = bAttrs || {};
  const keys = [...new Set([...Object.keys(aAttrs), ...Object.keys(bAttrs)])].sort();
  return keys.map((k) => {
    const av = resolveVal(aAttrs[k], aFields);
    const bv = resolveVal(bAttrs[k], bFields);
    const same = k in aAttrs && k in bAttrs && String(av) === String(bv);
    return { key: k, a: fmt(av), b: fmt(bv), same };
  });
}

function decideRegime(a, b) {
  const aChild = a.extends === b.modelType;
  const bChild = b.extends === a.modelType;
  if (aChild || bChild) {
    const parent = aChild ? b : a;
    return { kind: "lineage", basis: "id", note: `Parent/child — ${parent.modelType} is the base; diff comes from modular patches.` };
  }
  // Measure stable-id overlap.
  const idsA = new Set(a.nodes.keys());
  const idsB = new Set(b.nodes.keys());
  let shared = 0;
  for (const id of idsA) if (idsB.has(id)) shared++;
  const overlap = shared / Math.max(1, Math.min(idsA.size, idsB.size));
  if (overlap >= 0.5) {
    return { kind: "shared_ids", basis: "id", note: `Stable ids align (${Math.round(overlap * 100)}% overlap).`, overlap };
  }
  return {
    kind: "cross_lineage",
    basis: "kind",
    note: `Different lineages (${Math.round(overlap * 100)}% id overlap) — aligned by semantic kind.`,
    overlap,
  };
}

function diffHeadline(a, b) {
  const fa = a.architecture || {};
  const fb = b.architecture || {};
  const fields = HEADLINE_FIELDS.filter((k) => fa[k] !== undefined || fb[k] !== undefined);
  return fields.map((k) => ({
    field: k,
    a: fmt(fa[k]),
    b: fmt(fb[k]),
    same: fmt(fa[k]) === fmt(fb[k]),
  }));
}

function pick(fields, keys) {
  for (const k of keys) if (fields[k] != null) return { key: k, val: fields[k] };
  return null;
}

function totalDepth(ir) {
  let n = 0;
  for (const r of ir.repeatsById.values()) n += ir.resolveCount(r, ir.configFields);
  return n;
}

function diffScale(a, b) {
  const rows = [];
  const add = (label, av, bv) => {
    if (av == null && bv == null) return;
    rows.push({ field: label, a: fmt(av), b: fmt(bv), same: fmt(av) === fmt(bv) });
  };
  const hidA = pick(a.configFields, ["hidden_size", "d_model"]);
  const hidB = pick(b.configFields, ["hidden_size", "d_model"]);
  add("hidden size", hidA && hidA.val, hidB && hidB.val);
  const ffA = pick(a.configFields, ["intermediate_size", "d_ff"]);
  const ffB = pick(b.configFields, ["intermediate_size", "d_ff"]);
  add("intermediate", ffA && ffA.val, ffB && ffB.val);
  const hA = pick(a.configFields, ["num_attention_heads", "num_heads"]);
  const hB = pick(b.configFields, ["num_attention_heads", "num_heads"]);
  add("attention heads", hA && hA.val, hB && hB.val);
  add("total layers (resolved)", totalDepth(a), totalDepth(b));
  add("vocab size", a.configFields.vocab_size, b.configFields.vocab_size);
  return rows;
}

function diffTopology(a, b) {
  const setA = new Set(a.edges.map((e) => e.kind));
  const setB = new Set(b.edges.map((e) => e.kind));
  const kinds = [...new Set([...setA, ...setB])].sort();
  return kinds.map((k) => ({ kind: k, inA: setA.has(k), inB: setB.has(k), same: setA.has(k) === setB.has(k) }));
}

function kindCounts(ir) {
  const m = new Map();
  for (const n of ir.nodes.values()) m.set(n.kind, (m.get(n.kind) || 0) + 1);
  return m;
}

function diffKindCounts(a, b) {
  const ca = kindCounts(a);
  const cb = kindCounts(b);
  const kinds = [...new Set([...ca.keys(), ...cb.keys()])].sort();
  return kinds.map((k) => {
    const av = ca.get(k) || 0;
    const bv = cb.get(k) || 0;
    return { kind: k, a: av, b: bv, same: av === bv };
  });
}

// Align matched nodes by stable id (lineage / shared-id regimes).
function alignById(a, b) {
  const na = realNodes(a);
  const nb = realNodes(b);
  const matched = [];
  const removed = [];
  const added = [];
  for (const [id, node] of na) {
    if (nb.has(id)) {
      const other = nb.get(id);
      const attrDeltas = diffAttrs(node.attributes, other.attributes, a.configFields, b.configFields);
      const kindChanged = node.kind !== other.kind;
      const changed = kindChanged || attrDeltas.some((d) => !d.same);
      matched.push({
        id,
        label: a.label(id, a.configFields),
        kindA: node.kind,
        kindB: other.kind,
        kindChanged,
        changed,
        attrDeltas: attrDeltas.filter((d) => !d.same),
      });
    } else {
      removed.push({ id, label: a.label(id, a.configFields), kind: node.kind });
    }
  }
  for (const [id, node] of nb) {
    if (!na.has(id)) added.push({ id, label: b.label(id, b.configFields), kind: node.kind });
  }
  // Most-changed first.
  matched.sort((x, y) => Number(y.changed) - Number(x.changed) || y.attrDeltas.length - x.attrDeltas.length);
  return { basis: "id", matched, added, removed };
}

// Align by semantic kind "lanes" (cross-lineage regime).
function alignByKind(a, b) {
  const byKind = (ir) => {
    const m = new Map();
    for (const [id, n] of ir.nodes) {
      if (!m.has(n.kind)) m.set(n.kind, []);
      m.get(n.kind).push({ id, node: n });
    }
    return m;
  };
  const ka = byKind(a);
  const kb = byKind(b);
  const kinds = [...new Set([...ka.keys(), ...kb.keys()])].sort();
  const sample = (list) => (list || []).find((x) => x.node.attributes) || (list || [])[0];
  const lanes = kinds.map((k) => {
    const la = ka.get(k) || [];
    const lb = kb.get(k) || [];
    const sa = sample(la);
    const sb = sample(lb);
    const attrDeltas = diffAttrs(
      sa && sa.node.attributes,
      sb && sb.node.attributes,
      a.configFields,
      b.configFields
    ).filter((d) => !d.same);
    return {
      kind: k,
      countA: la.length,
      countB: lb.length,
      onlyA: la.length > 0 && lb.length === 0,
      onlyB: lb.length > 0 && la.length === 0,
      exampleA: sa ? a.label(sa.id, a.configFields) : null,
      exampleB: sb ? b.label(sb.id, b.configFields) : null,
      attrDeltas,
    };
  });
  return { basis: "kind", lanes };
}

// Presence-set diff (like topology) for two string lists.
function diffSet(listA, listB) {
  const sa = new Set(listA || []);
  const sb = new Set(listB || []);
  const items = [...new Set([...sa, ...sb])].sort();
  return items.map((item) => ({ item, inA: sa.has(item), inB: sb.has(item), same: sa.has(item) === sb.has(item) }));
}

function diffCapabilities(a, b) {
  const ca = a.capabilities;
  const cb = b.capabilities;
  if (!ca && !cb) return null;
  const schedStr = (ir, c) => {
    const s = c && c.attention_schedule;
    return Array.isArray(s) && s.length ? ir.scheduleSummary(s) : "uniform";
  };
  return {
    backends: diffSet(ca && ca.attention_backends, cb && cb.attention_backends),
    taskHeads: diffSet(ca && ca.task_heads, cb && cb.task_heads),
    rows: [
      {
        field: "tensor parallel",
        a: ca ? String(!!ca.tensor_parallel) : "—",
        b: cb ? String(!!cb.tensor_parallel) : "—",
        same: !!(ca && ca.tensor_parallel) === !!(cb && cb.tensor_parallel),
      },
      {
        field: "attention patterns",
        a: (ca && ca.attention_patterns || []).join(", ") || "—",
        b: (cb && cb.attention_patterns || []).join(", ") || "—",
        same: JSON.stringify(ca && ca.attention_patterns) === JSON.stringify(cb && cb.attention_patterns),
      },
      {
        field: "attention schedule",
        a: schedStr(a, ca),
        b: schedStr(b, cb),
        same: schedStr(a, ca) === schedStr(b, cb),
      },
    ],
  };
}

function compareIR(a, b) {
  const regime = decideRegime(a, b);
  return {
    a: { modelType: a.modelType, modelClass: (a.provenance && a.provenance.model_class) || a.modelType },
    b: { modelType: b.modelType, modelClass: (b.provenance && b.provenance.model_class) || b.modelType },
    regime,
    headline: diffHeadline(a, b),
    scale: diffScale(a, b),
    capabilities: diffCapabilities(a, b),
    topology: diffTopology(a, b),
    kinds: diffKindCounts(a, b),
    nodes: regime.basis === "id" ? alignById(a, b) : alignByKind(a, b),
  };
}

// Build a single union artifact from two models (for the merged diff view):
// nodes/edges present in both are kept once; the rest are tagged by origin so
// the renderer can colour them (both = neutral, a-only, b-only). Meaningful when
// the two share stable ids (same lineage); for cross-lineage it degenerates to
// two disjoint colourings, so callers should only use it in the id regime.
function mergeIR(a, b) {
  const clone = (o) => JSON.parse(JSON.stringify(o));
  const unionArr = (x, y) => [...(x || []), ...(y || []).filter((v) => !(x || []).includes(v))];
  const A = a.raw;
  const B = b.raw;
  const raw = { ...A };

  const mergeList = (key) => {
    const by = new Map((A[key] || []).map((n) => [n.id, clone(n)]));
    for (const n of B[key] || []) {
      if (by.has(n.id)) by.get(n.id).children = unionArr(by.get(n.id).children, n.children);
      else by.set(n.id, clone(n));
    }
    return [...by.values()];
  };
  raw.components = mergeList("components");
  raw.templates = mergeList("templates");

  const repById = new Map((A.repeats || []).map((r) => [r.id, r]));
  for (const r of B.repeats || []) if (!repById.has(r.id)) repById.set(r.id, r);
  raw.repeats = [...repById.values()];

  const eKey = (e) => `${e.source}|${e.target}|${e.kind}`;
  const edgeMap = new Map();
  for (const e of A.edges || []) edgeMap.set(eKey(e), e);
  for (const e of B.edges || []) if (!edgeMap.has(eKey(e))) edgeMap.set(eKey(e), e);
  raw.edges = [...edgeMap.values()];

  // Node origin.
  const has = (ir, id) => ir.nodes.has(id) || ir.pseudoIds.has(id);
  const ids = new Set([...a.nodes.keys(), ...b.nodes.keys(), ...a.pseudoIds, ...b.pseudoIds]);
  const origin = new Map();
  for (const id of ids) {
    const inA = has(a, id);
    const inB = has(b, id);
    origin.set(id, inA && inB ? "both" : inA ? "a" : "b");
  }
  // Edge origin.
  const aE = new Set((A.edges || []).map(eKey));
  const bE = new Set((B.edges || []).map(eKey));
  const edgeOrigin = new Map();
  for (const k of edgeMap.keys()) edgeOrigin.set(k, aE.has(k) && bE.has(k) ? "both" : aE.has(k) ? "a" : "b");

  return { raw, origin, edgeOrigin };
}

export { compareIR, mergeIR };