import type { XYPosition } from "@xyflow/react";
import type {
  CanvasEdge,
  CanvasNode,
  DataType,
  FunctionDef,
  GroupPort,
  NodeDef,
} from "./types";
import { canConnect, FN_PREFIX, getDef, GROUP_TYPE, isGroup } from "./registry";
import { makeEdge, makeNode, uid } from "./workflows";

export const GROUP_PAD = 20;
export const GROUP_HEAD = 36;
export const GROUP_MIN_W = 300;
export const GROUP_COLLAPSED_W = 240;

export const isGroupNode = (n: CanvasNode) => isGroup(n.data.defType);
export const titleOf = (n: CanvasNode) => n.data.title ?? getDef(n.data.defType).title;

/* Tree helpers ---------------------------------------------------------- */

export const childrenOf = (nodes: CanvasNode[], parentId: string | undefined) =>
  nodes.filter((n) => (n.parentId ?? undefined) === parentId);

export function descendants(nodes: CanvasNode[], id: string): CanvasNode[] {
  const out: CanvasNode[] = [];
  const walk = (pid: string) => {
    for (const c of childrenOf(nodes, pid)) {
      out.push(c);
      if (isGroupNode(c)) walk(c.id);
    }
  };
  walk(id);
  return out;
}

export function ancestors(nodes: CanvasNode[], id: string): string[] {
  const out: string[] = [];
  let cur = nodes.find((n) => n.id === id);
  while (cur?.parentId) {
    out.push(cur.parentId);
    cur = nodes.find((n) => n.id === cur!.parentId);
  }
  return out;
}

export function absolutePosition(nodes: CanvasNode[], n: CanvasNode): XYPosition {
  let x = n.position.x;
  let y = n.position.y;
  let cur = n;
  while (cur.parentId) {
    const p = nodes.find((x) => x.id === cur.parentId);
    if (!p) break;
    x += p.position.x;
    y += p.position.y;
    cur = p;
  }
  return { x, y };
}

const heightOf = (n: CanvasNode) =>
  n.measured?.height ?? n.height ?? (isGroupNode(n) ? 160 : n.data.collapsed ? 32 : 140);
const widthOf = (n: CanvasNode) => n.measured?.width ?? n.width ?? 260;

/** React Flow needs parents listed before their children. */
export function sortForFlow(nodes: CanvasNode[]): CanvasNode[] {
  const out: CanvasNode[] = [];
  const seen = new Set<string>();
  const place = (n: CanvasNode) => {
    if (seen.has(n.id)) return;
    if (n.parentId) {
      const p = nodes.find((x) => x.id === n.parentId);
      if (p) place(p);
    }
    seen.add(n.id);
    out.push(n);
  };
  nodes.forEach(place);
  return out;
}

/**
 * Sizes every expanded group to hug its children. Children that drifted
 * above/left of the padding shift the group instead, so nothing clips.
 */
export function refitGroups(nodes: CanvasNode[]): CanvasNode[] {
  let out = nodes;
  // Deepest first so nested groups are sized before their parents.
  const groups = out.filter(isGroupNode).sort((a, b) => ancestors(out, b.id).length - ancestors(out, a.id).length);
  for (const g of groups) {
    const kids = childrenOf(out, g.id);
    if (g.data.collapsed) continue;
    if (!kids.length) {
      if (g.width !== GROUP_MIN_W || g.height !== GROUP_HEAD + 80) {
        out = out.map((n) => (n.id === g.id ? { ...n, width: GROUP_MIN_W, height: GROUP_HEAD + 80 } : n));
      }
      continue;
    }
    const minX = Math.min(...kids.map((k) => k.position.x));
    const minY = Math.min(...kids.map((k) => k.position.y));
    const maxX = Math.max(...kids.map((k) => k.position.x + widthOf(k)));
    const maxY = Math.max(...kids.map((k) => k.position.y + heightOf(k)));
    const dx = minX - GROUP_PAD;
    const dy = minY - (GROUP_HEAD + GROUP_PAD);
    const width = Math.max(GROUP_MIN_W, maxX - dx + GROUP_PAD);
    const height = maxY - dy + GROUP_PAD;
    if (!dx && !dy && g.width === width && g.height === height) continue;
    out = out.map((n) => {
      if (n.id === g.id) {
        return { ...n, position: { x: n.position.x + dx, y: n.position.y + dy }, width, height };
      }
      if (n.parentId === g.id && (dx || dy)) {
        return { ...n, position: { x: n.position.x - dx, y: n.position.y - dy } };
      }
      return n;
    });
  }
  return out;
}

/* Boundary ports -------------------------------------------------------- */

/**
 * Ports a group exposes while collapsed: inputs not fed from inside, and
 * outputs used outside or not used at all.
 */
export function boundaryPorts(nodes: CanvasNode[], edges: CanvasEdge[], groupId: string) {
  const inner = descendants(nodes, groupId).filter((n) => !isGroupNode(n));
  const ids = new Set(inner.map((n) => n.id));
  const inputs: GroupPort[] = [];
  const outputs: GroupPort[] = [];
  for (const n of inner) {
    const def = getDef(n.data.defType);
    for (const p of def.inputs) {
      const fromInside = edges.some((e) => e.target === n.id && e.targetHandle === p.id && ids.has(e.source));
      if (!fromInside) inputs.push({ ...p, id: `${n.id}.${p.id}`, label: p.label || p.id, nodeId: n.id, portId: p.id });
    }
    for (const p of def.outputs) {
      const toInside = edges.some((e) => e.source === n.id && e.sourceHandle === p.id && ids.has(e.target));
      const toOutside = edges.some((e) => e.source === n.id && e.sourceHandle === p.id && !ids.has(e.target));
      if (!toInside || toOutside) outputs.push({ ...p, id: `${n.id}.${p.id}`, label: p.label || p.id, nodeId: n.id, portId: p.id });
    }
  }
  const dedupe = (ports: GroupPort[]) => {
    const count = new Map<string, number>();
    for (const p of ports) count.set(p.label, (count.get(p.label) ?? 0) + 1);
    return ports.map((p) =>
      (count.get(p.label) ?? 0) > 1
        ? { ...p, label: `${titleOf(inner.find((n) => n.id === p.nodeId)!)} · ${p.label}` }
        : p,
    );
  };
  return { inputs: dedupe(inputs), outputs: dedupe(outputs) };
}

/** Splits a collapsed-group handle id back into its inner node and port. */
export function splitHandle(handle: string | null | undefined) {
  if (!handle) return null;
  const i = handle.indexOf(".");
  if (i < 0) return null;
  return { nodeId: handle.slice(0, i), portId: handle.slice(i + 1) };
}

/* Display graph for the canvas ------------------------------------------ */

/** Hides children of collapsed groups and re-routes their edges to the group. */
export function displayGraph(nodes: CanvasNode[], edges: CanvasEdge[]) {
  const collapsedAncestor = new Map<string, string | null>();
  const topCollapsed = (id: string): string | null => {
    if (collapsedAncestor.has(id)) return collapsedAncestor.get(id)!;
    let result: string | null = null;
    for (const a of ancestors(nodes, id).reverse()) {
      const g = nodes.find((n) => n.id === a);
      if (g?.data.collapsed) {
        result = a;
        break;
      }
    }
    collapsedAncestor.set(id, result);
    return result;
  };
  const dNodes = nodes.map((n) => {
    const hiddenBy = topCollapsed(n.id);
    if (hiddenBy) return n.hidden ? n : { ...n, hidden: true };
    if (isGroupNode(n) && n.data.collapsed) {
      return {
        ...n,
        width: GROUP_COLLAPSED_W,
        height: undefined,
        data: { ...n.data, groupPorts: boundaryPorts(nodes, edges, n.id) },
      };
    }
    const draggable = !n.data.locked;
    if (n.hidden || (n.draggable ?? true) !== draggable) return { ...n, hidden: false, draggable };
    return n;
  });
  const dEdges: CanvasEdge[] = [];
  for (const e of edges) {
    const sTop = topCollapsed(e.source);
    const tTop = topCollapsed(e.target);
    if (sTop && sTop === tTop) continue;
    if (!sTop && !tTop) {
      dEdges.push(e);
      continue;
    }
    dEdges.push({
      ...e,
      source: sTop ?? e.source,
      sourceHandle: sTop ? `${e.source}.${e.sourceHandle}` : e.sourceHandle,
      target: tTop ?? e.target,
      targetHandle: tTop ? `${e.target}.${e.targetHandle}` : e.targetHandle,
    });
  }
  return { nodes: dNodes, edges: dEdges };
}

/* Pipeline structure ---------------------------------------------------- */

export interface TreeItem {
  id: string;
  children?: TreeItem[];
}

/**
 * Nested pipeline order. Each level is a topological order of its units
 * (a group counts as one unit carrying its descendants' dependencies),
 * tie-broken by the stored rank. Units in a cross-level cycle fall back
 * to rank order.
 */
export function structure(nodes: CanvasNode[], edges: CanvasEdge[], order: string[]): TreeItem[] {
  const rank = new Map(order.map((id, i) => [id, i]));
  const byId = new Map(nodes.map((n) => [n.id, n]));
  const unitOf = new Map<string, Map<string, string | null>>(); // parent -> nodeId -> unit id at that level
  const unitAt = (parent: string | undefined, nodeId: string): string | null => {
    const key = parent ?? "";
    let m = unitOf.get(key);
    if (!m) {
      m = new Map();
      unitOf.set(key, m);
    }
    if (m.has(nodeId)) return m.get(nodeId)!;
    let cur = byId.get(nodeId);
    let result: string | null = null;
    while (cur) {
      if ((cur.parentId ?? undefined) === parent) {
        result = cur.id;
        break;
      }
      cur = cur.parentId ? byId.get(cur.parentId) : undefined;
    }
    m.set(nodeId, result);
    return result;
  };
  const cmp = (a: string, b: string) => {
    const ra = rank.get(a) ?? Infinity;
    const rb = rank.get(b) ?? Infinity;
    if (ra !== rb) return ra - rb;
    const na = byId.get(a)!;
    const nb = byId.get(b)!;
    return na.position.y - nb.position.y || na.position.x - nb.position.x;
  };
  const level = (parent: string | undefined): TreeItem[] => {
    const units = childrenOf(nodes, parent).map((n) => n.id);
    const set = new Set(units);
    const indeg = new Map(units.map((u) => [u, 0]));
    const out = new Map<string, Set<string>>();
    for (const e of edges) {
      const s = unitAt(parent, e.source);
      const t = unitAt(parent, e.target);
      if (!s || !t || s === t || !set.has(s) || !set.has(t)) continue;
      if (out.get(s)?.has(t)) continue;
      out.set(s, new Set([...(out.get(s) ?? []), t]));
      indeg.set(t, (indeg.get(t) ?? 0) + 1);
    }
    const ready = units.filter((u) => indeg.get(u) === 0).sort(cmp);
    const seq: string[] = [];
    while (ready.length) {
      const u = ready.shift()!;
      seq.push(u);
      for (const t of out.get(u) ?? []) {
        indeg.set(t, indeg.get(t)! - 1);
        if (indeg.get(t) === 0) {
          ready.push(t);
          ready.sort(cmp);
        }
      }
    }
    for (const u of [...units].sort(cmp)) if (!seq.includes(u)) seq.push(u);
    return seq.map((id) => {
      const n = byId.get(id)!;
      return isGroupNode(n) ? { id, children: level(id) } : { id };
    });
  };
  return level(undefined);
}

export function flatten(tree: TreeItem[]): string[] {
  const out: string[] = [];
  const walk = (items: TreeItem[]) => {
    for (const it of items) {
      out.push(it.id);
      if (it.children) walk(it.children);
    }
  };
  walk(tree);
  return out;
}

/** Leaf (runnable) ids in pipeline order. */
export const leafSequence = (nodes: CanvasNode[], edges: CanvasEdge[], order: string[]) => {
  const groups = new Set(nodes.filter(isGroupNode).map((n) => n.id));
  return flatten(structure(nodes, edges, order)).filter((id) => !groups.has(id));
};

/** Execution order: a real topological sort of leaves, tie-broken by pipeline rank. */
export function executionOrder(nodes: CanvasNode[], edges: CanvasEdge[], order: string[]): string[] {
  const seq = leafSequence(nodes, edges, order);
  const rank = new Map(seq.map((id, i) => [id, i]));
  const runnable = nodes.filter(
    (n) => rank.has(n.id) && n.data.mode !== "never" && getDef(n.data.defType).cost !== 0,
  );
  const ids = new Set(runnable.map((n) => n.id));
  const indeg = new Map<string, number>();
  const out = new Map<string, string[]>();
  for (const id of ids) indeg.set(id, 0);
  for (const e of edges) {
    if (!ids.has(e.source) || !ids.has(e.target)) continue;
    indeg.set(e.target, (indeg.get(e.target) ?? 0) + 1);
    out.set(e.source, [...(out.get(e.source) ?? []), e.target]);
  }
  const cmp = (a: string, b: string) => rank.get(a)! - rank.get(b)!;
  const ready = [...ids].filter((id) => indeg.get(id) === 0).sort(cmp);
  const result: string[] = [];
  while (ready.length) {
    const id = ready.shift()!;
    result.push(id);
    for (const t of out.get(id) ?? []) {
      indeg.set(t, indeg.get(t)! - 1);
      if (indeg.get(t) === 0) {
        ready.push(t);
        ready.sort(cmp);
      }
    }
  }
  return result;
}

/* Auto wiring ----------------------------------------------------------- */

const edgeType = (a: DataType, b: DataType): DataType => (a === "any" ? b : a);

/**
 * Wires the free inputs of `nodeId` to the nearest preceding leaf in `seq`
 * with a compatible output, spreading same-typed inputs across sources.
 */
export function autoWire(nodes: CanvasNode[], edges: CanvasEdge[], seq: string[], nodeId: string): CanvasEdge[] {
  const node = nodes.find((n) => n.id === nodeId);
  if (!node || isGroupNode(node)) return [];
  const idx = seq.indexOf(nodeId);
  if (idx < 0) return [];
  const def = getDef(node.data.defType);
  const added: CanvasEdge[] = [];
  const used = new Set(edges.filter((e) => e.target === nodeId).map((e) => `${e.source}.${e.sourceHandle}`));
  for (const p of def.inputs) {
    if (edges.some((e) => e.target === nodeId && e.targetHandle === p.id)) continue;
    let pick: { id: string; port: string; type: DataType } | null = null;
    let fallback: { id: string; port: string; type: DataType } | null = null;
    for (let j = idx - 1; j >= 0 && !pick; j--) {
      const src = nodes.find((n) => n.id === seq[j]);
      if (!src || src.data.mode === "never") continue;
      for (const o of getDef(src.data.defType).outputs) {
        if (!canConnect(o.type, p.type)) continue;
        const key = `${src.id}.${o.id}`;
        if (used.has(key)) {
          fallback ??= { id: src.id, port: o.id, type: o.type };
          continue;
        }
        pick = { id: src.id, port: o.id, type: o.type };
        break;
      }
    }
    pick ??= fallback;
    if (!pick) continue;
    used.add(`${pick.id}.${pick.port}`);
    added.push({ ...makeEdge(pick.id, pick.port, nodeId, p.id), data: { type: edgeType(pick.type, p.type) } });
  }
  return added;
}

/** Points the next leaf's compatible inputs at a freshly inserted leaf. */
export function rewireNext(nodes: CanvasNode[], edges: CanvasEdge[], fromId: string, nextId: string): CanvasEdge[] {
  const from = nodes.find((n) => n.id === fromId);
  const next = nodes.find((n) => n.id === nextId);
  if (!from || !next || isGroupNode(from) || isGroupNode(next)) return edges;
  const outs = getDef(from.data.defType).outputs;
  let result = edges;
  const usedOut = new Set<string>();
  for (const p of getDef(next.data.defType).inputs) {
    const o =
      outs.find((x) => canConnect(x.type, p.type) && !usedOut.has(x.id)) ??
      outs.find((x) => canConnect(x.type, p.type));
    if (!o) continue;
    usedOut.add(o.id);
    result = result.filter((e) => !(e.target === nextId && e.targetHandle === p.id));
    result = [...result, { ...makeEdge(fromId, o.id, nextId, p.id), data: { type: edgeType(o.type, p.type) } }];
  }
  return result;
}

/**
 * Enforces "each step feeds the next": drops edges that point backwards in
 * the leaf sequence and re-wires the inputs that freed up.
 */
export function repairEdges(nodes: CanvasNode[], edges: CanvasEdge[], order: string[]): CanvasEdge[] {
  const seq = leafSequence(nodes, edges, order);
  const rank = new Map(seq.map((id, i) => [id, i]));
  let next = edges.filter((e) => (rank.get(e.source) ?? -1) < (rank.get(e.target) ?? Infinity));
  for (const id of seq) next = [...next, ...autoWire(nodes, next, seq, id)];
  return next;
}

/* Group nodes ----------------------------------------------------------- */

export function makeGroupNode(id: string, name: string, x: number, y: number, parentId?: string): CanvasNode {
  const n = makeNode(id, GROUP_TYPE, x, y, {}, { title: name });
  n.type = "group";
  n.width = GROUP_MIN_W;
  n.height = GROUP_HEAD + 80;
  if (parentId) n.parentId = parentId;
  return n;
}

/* Reusable functions ---------------------------------------------------- */

/** Snapshot of a group's contents, with local ids. */
export function buildFunction(nodes: CanvasNode[], edges: CanvasEdge[], groupId: string, name: string): FunctionDef {
  const inner = descendants(nodes, groupId);
  const map = new Map(inner.map((n, i) => [n.id, `l${i}`]));
  const ids = new Set(inner.map((n) => n.id));
  const fnNodes = inner.map((n) => ({
    ...n,
    id: map.get(n.id)!,
    parentId: n.parentId === groupId ? undefined : map.get(n.parentId!),
    selected: false,
    dragging: false,
    hidden: false,
    data: { ...n.data, status: "idle" as const, progress: 0, values: { ...n.data.values }, groupPorts: undefined },
  }));
  const fnEdges = edges
    .filter((e) => ids.has(e.source) && ids.has(e.target))
    .map((e) => ({
      ...makeEdge(map.get(e.source)!, e.sourceHandle!, map.get(e.target)!, e.targetHandle!),
      data: e.data ? { ...e.data, flowing: false } : e.data,
    }));
  const leaves = inner.filter((n) => !isGroupNode(n));
  return {
    id: uid("fn"),
    name,
    description: `${leaves.length} step${leaves.length === 1 ? "" : "s"}: ${leaves.map(titleOf).join(", ")}`,
    nodes: fnNodes,
    edges: fnEdges,
  };
}

/** Fresh group node + copied children for a function, at an absolute position. */
export function instantiateFunction(fn: FunctionDef, x: number, y: number, parentId?: string) {
  const gid = uid("g");
  const group = makeGroupNode(gid, fn.name, x, y, parentId);
  group.data.fnId = fn.id;
  const map = new Map(fn.nodes.map((n) => [n.id, uid()]));
  const nodes = fn.nodes.map((n) => ({
    ...n,
    id: map.get(n.id)!,
    parentId: n.parentId ? map.get(n.parentId)! : gid,
    data: { ...n.data, values: { ...n.data.values } },
  }));
  const edges = fn.edges.map((e) => ({
    ...makeEdge(map.get(e.source)!, e.sourceHandle!, map.get(e.target)!, e.targetHandle!),
    data: e.data,
  }));
  return { group, nodes, edges };
}

/** Library/palette entry for a function. */
export function fnToDef(fn: FunctionDef): NodeDef {
  const nodes = fn.nodes.map((n) => (n.parentId ? n : { ...n, parentId: "__root__" }));
  const ports = boundaryPorts(nodes, fn.edges, "__root__");
  return {
    type: FN_PREFIX + fn.id,
    title: fn.name,
    category: "Functions",
    description: fn.description,
    inputs: ports.inputs,
    outputs: ports.outputs,
    widgets: [],
    width: GROUP_MIN_W,
    cost: 0,
  };
}