ai-chatbot/lib/canvas/arrange.ts

import type { ELK, ElkNode } from "elkjs/lib/elk.bundled.js";

export const CANVAS_ARRANGE_NODE_GAP = 100;
export const CANVAS_ARRANGE_COMPONENT_GAP = 1.5 * CANVAS_ARRANGE_NODE_GAP;
export const CANVAS_ARRANGE_MAX_ROW_WIDTH = 3200;
export const CANVAS_ARRANGE_SNAP_GRID_SIZE = 12;
export const CANVAS_ARRANGE_VIEWPORT_OFFSET = 48;
export const CANVAS_ARRANGE_VIEWPORT_PADDING = 0.12;

export type CanvasArrangeNode = {
  height: number;
  id: string;
  parentId?: string;
  width: number;
  x: number;
  y: number;
};

export type CanvasArrangeEdge = {
  source: string;
  target: string;
};

export type CanvasArrangePositionUpdate = {
  id: string;
  position: { x: number; y: number };
};

type LayoutComponent = {
  height: number;
  positions: Map<string, { x: number; y: number }>;
  width: number;
};

let elkPromise: Promise<ELK> | null = null;

function getElk() {
  if (!elkPromise) {
    elkPromise = import("elkjs/lib/elk.bundled.js")
      .then(({ default: Elk }) => new Elk())
      .catch((error) => {
        elkPromise = null;
        throw error;
      });
  }
  return elkPromise;
}

function snapArrangeCoordinate(value: number) {
  return Math.round(value / CANVAS_ARRANGE_SNAP_GRID_SIZE) * CANVAS_ARRANGE_SNAP_GRID_SIZE;
}

function getRootNodeId(nodeId: string, nodesById: Map<string, CanvasArrangeNode>) {
  let currentId = nodeId;
  const visited = new Set<string>();

  for (let index = 0; index < 256; index += 1) {
    const node = nodesById.get(currentId);
    if (!node?.parentId || visited.has(currentId)) break;
    visited.add(currentId);
    currentId = node.parentId;
  }

  return currentId;
}

function groupConnectedRootNodes(
  nodes: CanvasArrangeNode[],
  edges: CanvasArrangeEdge[],
  nodesById: Map<string, CanvasArrangeNode>
) {
  const rootNodes = nodes.filter((node) => !node.parentId);
  if (rootNodes.length === 0) return [];

  const rootIds = new Set(rootNodes.map((node) => node.id));
  const parents = new Map(rootNodes.map((node) => [node.id, node.id]));

  const find = (nodeId: string) => {
    let rootId = nodeId;
    while (parents.get(rootId) !== rootId) {
      rootId = parents.get(rootId) || rootId;
    }

    let currentId = nodeId;
    while (parents.get(currentId) !== rootId) {
      const nextId = parents.get(currentId);
      parents.set(currentId, rootId);
      if (!nextId) break;
      currentId = nextId;
    }
    return rootId;
  };

  for (const edge of edges) {
    const sourceId = getRootNodeId(edge.source, nodesById);
    const targetId = getRootNodeId(edge.target, nodesById);
    if (sourceId === targetId || !rootIds.has(sourceId) || !rootIds.has(targetId)) continue;
    const sourceRoot = find(sourceId);
    const targetRoot = find(targetId);
    if (sourceRoot !== targetRoot) parents.set(sourceRoot, targetRoot);
  }

  const groups = new Map<string, CanvasArrangeNode[]>();
  for (const node of rootNodes) {
    const rootId = find(node.id);
    const group = groups.get(rootId) || [];
    group.push(node);
    groups.set(rootId, group);
  }

  const sortedGroups = [...groups.values()].map((group) =>
    group.sort((first, second) => {
      const yDifference = first.y - second.y;
      return Math.abs(yDifference) > 1 ? yDifference : first.x - second.x;
    })
  );

  return sortedGroups.sort((first, second) => {
    const firstY = Math.min(...first.map((node) => node.y));
    const secondY = Math.min(...second.map((node) => node.y));
    if (Math.abs(firstY - secondY) > 1) return firstY - secondY;
    return (
      Math.min(...first.map((node) => node.x)) -
      Math.min(...second.map((node) => node.x))
    );
  });
}

async function layoutComponent(
  componentNodes: CanvasArrangeNode[],
  edges: CanvasArrangeEdge[],
  nodesById: Map<string, CanvasArrangeNode>
): Promise<LayoutComponent | null> {
  const componentIds = new Set(componentNodes.map((node) => node.id));
  const seenEdges = new Set<string>();
  const layoutEdges: NonNullable<ElkNode["edges"]> = [];

  for (const edge of edges) {
    const sourceId = getRootNodeId(edge.source, nodesById);
    const targetId = getRootNodeId(edge.target, nodesById);
    if (
      sourceId === targetId ||
      !componentIds.has(sourceId) ||
      !componentIds.has(targetId)
    ) {
      continue;
    }
    const edgeKey = `${sourceId}->${targetId}`;
    if (seenEdges.has(edgeKey)) continue;
    seenEdges.add(edgeKey);
    layoutEdges.push({
      id: `e-${sourceId}-${targetId}`,
      sources: [sourceId],
      targets: [targetId],
    });
  }

  const graph: ElkNode = {
    children: componentNodes.map((node) => ({
      height: Math.max(1, node.height),
      id: node.id,
      width: Math.max(1, node.width),
    })),
    edges: layoutEdges,
    id: "root",
    layoutOptions: {
      "elk.algorithm": "layered",
      "elk.direction": "RIGHT",
      "elk.layered.spacing.nodeNodeBetweenLayers": String(2 * CANVAS_ARRANGE_NODE_GAP),
      "elk.padding": "[ top=0, left=0, bottom=0, right=0 ]",
      "elk.spacing.nodeNode": String(CANVAS_ARRANGE_NODE_GAP),
    },
  };

  let layout: ElkNode;
  try {
    layout = await (await getElk()).layout(graph);
  } catch {
    return null;
  }

  const positions = new Map<string, { x: number; y: number }>();
  for (const child of layout.children || []) {
    if (!child.id) continue;
    positions.set(child.id, { x: child.x || 0, y: child.y || 0 });
  }

  const minX = Math.min(...[...positions.values()].map((position) => position.x));
  const minY = Math.min(...[...positions.values()].map((position) => position.y));
  if (Number.isFinite(minX) && Number.isFinite(minY)) {
    for (const position of positions.values()) {
      position.x -= minX;
      position.y -= minY;
    }
  }

  let width = 0;
  let height = 0;
  for (const [nodeId, position] of positions) {
    const node = nodesById.get(nodeId);
    if (!node) continue;
    width = Math.max(width, position.x + node.width);
    height = Math.max(height, position.y + node.height);
  }

  return positions.size > 0 ? { height, positions, width } : null;
}

export async function arrangeCanvasNodes(
  nodes: CanvasArrangeNode[],
  edges: CanvasArrangeEdge[]
): Promise<CanvasArrangePositionUpdate[]> {
  const rootNodes = nodes.filter((node) => !node.parentId);
  if (rootNodes.length <= 1) return [];

  const nodesById = new Map(nodes.map((node) => [node.id, node]));
  const groups = groupConnectedRootNodes(nodes, edges, nodesById);
  const components: LayoutComponent[] = [];

  for (const group of groups) {
    const component = await layoutComponent(group, edges, nodesById);
    if (component) components.push(component);
  }

  const originX = snapArrangeCoordinate(Math.min(...rootNodes.map((node) => node.x)));
  const originY = snapArrangeCoordinate(Math.min(...rootNodes.map((node) => node.y)));
  const updates: CanvasArrangePositionUpdate[] = [];
  let cursorX = originX;
  let cursorY = originY;
  let rowHeight = 0;

  for (const component of components) {
    if (
      cursorX > originX &&
      cursorX + component.width > originX + CANVAS_ARRANGE_MAX_ROW_WIDTH
    ) {
      cursorY += rowHeight + CANVAS_ARRANGE_COMPONENT_GAP;
      cursorX = originX;
      rowHeight = 0;
    }

    for (const [nodeId, position] of component.positions) {
      const node = nodesById.get(nodeId);
      if (!node) continue;
      const x = snapArrangeCoordinate(cursorX + position.x);
      const y = snapArrangeCoordinate(cursorY + position.y);
      if (Math.abs(x - node.x) > 0.5 || Math.abs(y - node.y) > 0.5) {
        updates.push({ id: nodeId, position: { x, y } });
      }
    }

    cursorX += component.width + CANVAS_ARRANGE_COMPONENT_GAP;
    rowHeight = Math.max(rowHeight, component.height);
  }

  return updates;
}

export function applyCanvasArrangeUpdates<TNode extends CanvasArrangeNode>(
  nodes: TNode[],
  updates: CanvasArrangePositionUpdate[]
) {
  const positionsById = new Map(updates.map((update) => [update.id, update.position]));
  return nodes.map((node) => {
    const position = positionsById.get(node.id);
    return position ? { ...node, ...position } : node;
  });
}

export function getArrangeCanvasViewport({
  canvasHeight,
  canvasWidth,
  maxZoom,
  minZoom,
  nodes,
}: {
  canvasHeight: number;
  canvasWidth: number;
  maxZoom: number;
  minZoom: number;
  nodes: CanvasArrangeNode[];
}) {
  if (nodes.length === 0 || canvasWidth <= 0 || canvasHeight <= 0) return null;

  const minX = Math.min(...nodes.map((node) => node.x));
  const minY = Math.min(...nodes.map((node) => node.y));
  const maxX = Math.max(...nodes.map((node) => node.x + node.width));
  const maxY = Math.max(...nodes.map((node) => node.y + node.height));
  const width = maxX - minX;
  const height = maxY - minY;
  if (width <= 0 || height <= 0) return null;

  const horizontalPadding = Math.floor(
    ((canvasWidth - canvasWidth / (1 + CANVAS_ARRANGE_VIEWPORT_PADDING)) * 0.5)
  );
  const verticalPadding = Math.floor(
    ((canvasHeight - canvasHeight / (1 + CANVAS_ARRANGE_VIEWPORT_PADDING)) * 0.5)
  );
  const zoom = Math.min(
    maxZoom,
    Math.max(
      minZoom,
      Math.min(
        (canvasWidth - 2 * horizontalPadding) / width,
        (canvasHeight - 2 * verticalPadding) / height
      )
    )
  );

  return {
    x: CANVAS_ARRANGE_VIEWPORT_OFFSET - minX * zoom,
    y: CANVAS_ARRANGE_VIEWPORT_OFFSET - minY * zoom,
    zoom,
  };
}