import type { Pipeline, PipelineNode } from "./types";

export type LayoutMode = "circular" | "rows" | "grid";

export type LayoutResult = {
  positions: Record<string, { x: number; y: number }>;
  flatChildren: Record<string, PipelineNode>;
};

const RING_1_RADIUS = 460;
const RING_2_RADIUS = 230;

function computeCircularLayout(pipeline: Pipeline, expandedIds: Set<string>): LayoutResult {
  const top = pipeline.nodes.filter((n) => n.category !== "orchestration");
  const hub = pipeline.nodes.find((n) => n.category === "orchestration");
  const count = top.length;
  const positions: Record<string, { x: number; y: number }> = {};
  const flatChildren: Record<string, PipelineNode> = {};

  if (hub) positions[hub.id] = { x: 0, y: 0 };

  top.forEach((n, i) => {
    const angle = -Math.PI / 2 + i * ((2 * Math.PI) / count);
    positions[n.id] = {
      x: RING_1_RADIUS * Math.cos(angle),
      y: RING_1_RADIUS * Math.sin(angle),
    };

    if (expandedIds.has(n.id) && n.children?.length) {
      const childCount = n.children.length;
      const spread = Math.min(0.9, 0.28 * childCount);
      n.children.forEach((child, ci) => {
        flatChildren[child.id] = child;
        const childAngle =
          angle + (ci - (childCount - 1) / 2) * (spread / Math.max(childCount - 1, 1));
        positions[child.id] = {
          x: positions[n.id].x + RING_2_RADIUS * Math.cos(childAngle),
          y: positions[n.id].y + RING_2_RADIUS * Math.sin(childAngle),
        };
      });
    }
  });

  return { positions, flatChildren };
}

// Rows: nodes grouped into horizontal bands by pipeline order (top-down DAG),
// matching a layered "structural" view rather than a circular loop.
const ROW_ORDER: Record<string, number> = {
  vision: 0,
  asr: 0,
  llm: 1,
  tts: 2,
  keyframer: 3,
  avatar: 3,
  video: 4,
  orchestration: 5,
};

const ROW_HEIGHT = 190;
const COLUMN_WIDTH = 260;
const CHILD_ROW_HEIGHT = 110;

function computeRowsLayout(pipeline: Pipeline, expandedIds: Set<string>): LayoutResult {
  const positions: Record<string, { x: number; y: number }> = {};
  const flatChildren: Record<string, PipelineNode> = {};

  const rows = new Map<number, PipelineNode[]>();
  pipeline.nodes.forEach((n) => {
    const row = ROW_ORDER[n.category] ?? 0;
    if (!rows.has(row)) rows.set(row, []);
    rows.get(row)!.push(n);
  });

  const sortedRowIndices = Array.from(rows.keys()).sort((a, b) => a - b);

  sortedRowIndices.forEach((rowIndex) => {
    const nodesInRow = rows.get(rowIndex)!;
    const rowWidth = (nodesInRow.length - 1) * COLUMN_WIDTH;
    nodesInRow.forEach((n, i) => {
      const x = i * COLUMN_WIDTH - rowWidth / 2;
      const y = rowIndex * ROW_HEIGHT;
      positions[n.id] = { x, y };

      if (expandedIds.has(n.id) && n.children?.length) {
        const childCount = n.children.length;
        const childRowWidth = (childCount - 1) * (COLUMN_WIDTH * 0.7);
        n.children.forEach((child, ci) => {
          flatChildren[child.id] = child;
          positions[child.id] = {
            x: x + ci * (COLUMN_WIDTH * 0.7) - childRowWidth / 2,
            y: y + CHILD_ROW_HEIGHT,
          };
        });
      }
    });
  });

  return { positions, flatChildren };
}

// Grid: an even N-column grid with no hub/ring bias — used by the Test Bench,
// where nodes are picked/inspected individually rather than read as a flow.
// The Test Bench renders compact node cards (code's already been read on the
// main frontend), so its grid cells are noticeably tighter than the default.
const GRID_COLS = 3;
const GRID_CELL_W = 300;
const GRID_CELL_H = 230;
const GRID_COMPACT_COLS = 4;
const GRID_COMPACT_CELL_W = 190;
const GRID_COMPACT_CELL_H = 130;

function computeGridLayout(
  pipeline: Pipeline,
  expandedIds: Set<string>,
  compact = false,
): LayoutResult {
  const positions: Record<string, { x: number; y: number }> = {};
  const flatChildren: Record<string, PipelineNode> = {};
  const cols = compact ? GRID_COMPACT_COLS : GRID_COLS;
  const cellW = compact ? GRID_COMPACT_CELL_W : GRID_CELL_W;
  const cellH = compact ? GRID_COMPACT_CELL_H : GRID_CELL_H;

  pipeline.nodes.forEach((n, i) => {
    const col = i % cols;
    const row = Math.floor(i / cols);
    const x = col * cellW;
    const y = row * cellH;
    positions[n.id] = { x, y };

    if (expandedIds.has(n.id) && n.children?.length) {
      n.children.forEach((child, ci) => {
        flatChildren[child.id] = child;
        positions[child.id] = { x: x + (ci + 1) * 40, y: y + cellH * 0.7 };
      });
    }
  });

  return { positions, flatChildren };
}

export function computeLayout(
  pipeline: Pipeline,
  expandedIds: Set<string>,
  mode: LayoutMode = "circular",
  compactGrid = false,
): LayoutResult {
  if (mode === "rows") return computeRowsLayout(pipeline, expandedIds);
  if (mode === "grid") return computeGridLayout(pipeline, expandedIds, compactGrid);
  return computeCircularLayout(pipeline, expandedIds);
}

// Body-layer grouping: cluster nodes by functional layer — Senses (input),
// Brain (cognitive + serving), Face (generative output) — matching the
// biological-metaphor architecture (senses -> brain -> face), not literal
// anatomy. A distinct grouping mode from circular/rows: it repositions nodes
// without remounting the graph, so the transition can glide instead of
// re-landing.
export type BodyGroup = "senses" | "brain" | "face";

export const BODY_GROUP_OF: Record<string, BodyGroup> = {
  vision: "senses",
  asr: "senses",
  llm: "brain",
  orchestration: "brain",
  tts: "face",
  avatar: "face",
  video: "face",
  keyframer: "face",
};

export const BODY_GROUP_LABEL: Record<BodyGroup, string> = {
  senses: "THE SENSES",
  brain: "THE BRAIN",
  face: "THE FACE",
};

export const BODY_GROUP_COLOR: Record<BodyGroup, string> = {
  senses: "#38bdf8",
  brain: "#2dd4bf",
  face: "#fbbf24",
};

const BODY_GROUP_X: Record<BodyGroup, number> = {
  senses: -480,
  brain: 0,
  face: 480,
};

const BODY_GROUP_ORDER: BodyGroup[] = ["senses", "brain", "face"];
const BODY_NODE_V_GAP = 200;
const BODY_HEADER_MARGIN = 110;

export type BodyLayoutResult = {
  positions: Record<string, { x: number; y: number }>;
  groupHeaders: { group: BodyGroup; x: number; y: number }[];
};

export function computeBodyGroupLayout(pipeline: Pipeline): BodyLayoutResult {
  const positions: Record<string, { x: number; y: number }> = {};
  const buckets: Record<BodyGroup, string[]> = { senses: [], brain: [], face: [] };

  pipeline.nodes.forEach((n) => {
    const group = BODY_GROUP_OF[n.category] ?? "brain";
    buckets[group].push(n.id);
  });

  const groupHeaders: BodyLayoutResult["groupHeaders"] = [];

  BODY_GROUP_ORDER.forEach((group) => {
    const x = BODY_GROUP_X[group];
    const ids = buckets[group];
    const totalHeight = (ids.length - 1) * BODY_NODE_V_GAP;
    const topY = -totalHeight / 2;
    ids.forEach((id, i) => {
      positions[id] = { x, y: topY + i * BODY_NODE_V_GAP };
    });
    groupHeaders.push({ group, x, y: topY - BODY_HEADER_MARGIN });
  });

  return { positions, groupHeaders };
}