import { absolutePosition, isGroupNode } from "./groups";
import type { CanvasEdge, CanvasNode } from "./types";

/** Canvas size of a node as rendered (falls back to defaults before React Flow has measured it). */
export const sizeOf = (n: CanvasNode) => ({
  w: Math.round(n.measured?.width ?? n.width ?? 260),
  h: Math.round(n.measured?.height ?? n.height ?? (isGroupNode(n) ? 160 : n.data.collapsed ? 32 : 140)),
});

const GAP_X = 80;
const GAP_Y = 40;
const snap = (v: number) => Math.round(v / 10) * 10;

type Rect = { x: number; y: number; w: number; h: number };
const hits = (a: Rect, b: Rect, pad = 16) => a.x < b.x + b.w + pad && a.x + a.w + pad > b.x && a.y < b.y + b.h + pad && a.y + a.h + pad > b.y;

/**
 * Left-to-right layered layout: each node goes in the column of its longest path from a source,
 * columns are ordered to follow their inputs, and nothing overlaps. Works per parent (top level or
 * inside one group), in that parent's coordinates. With `ids`, only those nodes move (the block is
 * shifted down if it would land on a sibling left in place). Returns new parent-relative positions.
 */
export function layeredLayout(nodes: CanvasNode[], edges: CanvasEdge[], ids?: string[]): Map<string, { x: number; y: number }> {
  const out = new Map<string, { x: number; y: number }>();
  const byId = new Map(nodes.map((n) => [n.id, n]));
  // The unit of a node at a given parent level: itself, or the ancestor group that sits at that level.
  const unitAt = (id: string, parent: string | undefined): string | undefined => {
    let cur = byId.get(id);
    while (cur && cur.parentId !== parent) cur = cur.parentId ? byId.get(cur.parentId) : undefined;
    return cur?.id;
  };
  const scope = ids?.length ? nodes.filter((n) => ids.includes(n.id)) : nodes.filter((n) => !n.parentId);
  const parents = [...new Set(scope.map((n) => n.parentId))];
  for (const parent of parents) {
    const set = scope.filter((n) => n.parentId === parent);
    if (!set.length) continue;
    const inSet = new Set(set.map((n) => n.id));
    // Edges between units at this level (edges into / out of groups count for the group).
    const links: [string, string][] = [];
    for (const e of edges) {
      const a = unitAt(e.source, parent);
      const b = unitAt(e.target, parent);
      if (a && b && a !== b && inSet.has(a) && inSet.has(b)) links.push([a, b]);
    }
    // Longest-path depth, bounded so cycles can't loop forever.
    const depth = new Map(set.map((n) => [n.id, 0]));
    for (let i = 0; i < set.length; i++) {
      let changed = false;
      for (const [a, b] of links) {
        const d = depth.get(a)! + 1;
        if (d > depth.get(b)! && d < set.length) {
          depth.set(b, d);
          changed = true;
        }
      }
      if (!changed) break;
    }
    const cols: CanvasNode[][] = [];
    for (const n of set) (cols[depth.get(n.id)!] ??= []).push(n);
    const filled = cols.filter(Boolean);
    // Order: first column keeps its current vertical order; later ones follow the mean slot of their inputs.
    const slot = new Map<string, number>();
    filled.forEach((col, ci) => {
      if (ci === 0) col.sort((a, b) => a.position.y - b.position.y);
      else {
        const key = (n: CanvasNode) => {
          const ins = links.filter(([, b]) => b === n.id).map(([a]) => slot.get(a)).filter((v): v is number => v !== undefined);
          return ins.length ? ins.reduce((s, v) => s + v, 0) / ins.length : Number.MAX_SAFE_INTEGER;
        };
        col.sort((a, b) => key(a) - key(b) || a.position.y - b.position.y);
      }
      col.forEach((n, i) => slot.set(n.id, i));
    });
    const heights = filled.map((col) => col.reduce((s, n) => s + sizeOf(n).h, 0) + GAP_Y * (col.length - 1));
    const tallest = Math.max(...heights);
    const x0 = Math.min(...set.map((n) => n.position.x));
    const y0 = Math.min(...set.map((n) => n.position.y));
    const placed: Rect[] = [];
    let x = x0;
    filled.forEach((col, ci) => {
      let y = y0 + (tallest - heights[ci]) / 2;
      for (const n of col) {
        const { w, h } = sizeOf(n);
        placed.push({ x: snap(x), y: snap(y), w, h });
        out.set(n.id, { x: snap(x), y: snap(y) });
        y += h + GAP_Y;
      }
      x += Math.max(...col.map((n) => sizeOf(n).w)) + GAP_X;
    });
    // Siblings that stay put: move the block down until it clears them.
    const others = nodes.filter((n) => n.parentId === parent && !inSet.has(n.id)).map((n) => ({ ...n.position, ...sizeOf(n) }));
    if (others.length) {
      let dy = 0;
      for (let i = 0; i < 100; i++) {
        const block = placed.map((r) => ({ ...r, y: r.y + dy }));
        const hit = others.find((o) => block.some((r) => hits(r, o)));
        if (!hit) break;
        dy = hit.y + hit.h + GAP_Y - Math.min(...placed.map((r) => r.y));
      }
      if (dy) for (const [id, p] of out) out.set(id, { x: p.x, y: snap(p.y + dy) });
    }
  }
  return out;
}

/** Absolute canvas rectangle of a node (children of groups are stored relative to their group). */
export function absRect(nodes: CanvasNode[], n: CanvasNode) {
  const p = absolutePosition(nodes, n);
  const s = sizeOf(n);
  return { x: Math.round(p.x), y: Math.round(p.y), width: s.w, height: s.h };
}