import { useMemo } from 'react';
import { useDiagramByRoute } from '../../api/queries/diagrams';
import type { DiagramNode, DiagramEdge, DiagramLayout } from '../../api/queries/diagrams';
import type { DiagramSection } from './types';
import { isErrorCompoundType, isCompletionCompoundType } from './node-colors';

const SECTION_GAP = 80;

export function useDiagramData(
  application: string,
  routeId: string,
  direction: 'LR' | 'TB' = 'LR',
  preloadedLayout?: DiagramLayout,
) {
  // When a preloaded layout is provided, disable the internal fetch
  const fetchApp = preloadedLayout ? undefined : application;
  const fetchRoute = preloadedLayout ? undefined : routeId;
  const { data: fetchedLayout, isLoading, error } = useDiagramByRoute(fetchApp, fetchRoute, direction);

  const layout = preloadedLayout ?? fetchedLayout;

  const result = useMemo(() => {
    if (!layout?.nodes) {
      return { sections: [], totalWidth: 0, totalHeight: 0 };
    }

    const allEdges = layout.edges ?? [];

    // Separate main nodes from completion and error handler compound sections
    const mainNodes: DiagramNode[] = [];
    const completionSections: { label: string; nodes: DiagramNode[] }[] = [];
    const errorSections: { label: string; nodes: DiagramNode[] }[] = [];

    for (const node of layout.nodes) {
      if (isCompletionCompoundType(node.type) && node.children && node.children.length > 0) {
        completionSections.push({
          label: node.label || 'onCompletion',
          nodes: node.children,
        });
      } else if (isErrorCompoundType(node.type) && node.children && node.children.length > 0) {
        errorSections.push({
          label: node.label || 'Error Handler',
          nodes: node.children,
        });
      } else {
        mainNodes.push(node);
      }
    }

    // Collect node IDs for edge filtering
    const mainNodeIds = new Set<string>();
    collectNodeIds(mainNodes, mainNodeIds);

    const mainEdges = allEdges.filter(
      e => mainNodeIds.has(e.sourceId) && mainNodeIds.has(e.targetId),
    );

    // Normalize main section to start at (0, 0) — ELK can place nodes
    // at arbitrary positions within its root graph
    const mainBounds = computeBounds(mainNodes);
    const mainOffX = mainBounds.minX;
    const mainOffY = mainBounds.minY;
    const shiftedMainNodes = shiftNodes(mainNodes, mainOffX, mainOffY);
    const shiftedMainEdges = mainEdges.map(e => ({
      ...e,
      points: e.points.map(p => [p[0] - mainOffX, p[1] - mainOffY]),
    }));
    const mainWidth = mainBounds.maxX - mainBounds.minX;
    const mainHeight = mainBounds.maxY - mainBounds.minY;

    const sections: DiagramSection[] = [
      {
        label: 'Main Route',
        nodes: shiftedMainNodes,
        edges: shiftedMainEdges,
        offsetY: 0,
      },
    ];

    let currentY = mainHeight + SECTION_GAP;
    let maxWidth = mainWidth;

    const addHandlerSections = (
      handlers: { label: string; nodes: DiagramNode[] }[],
      variant: 'completion' | 'error',
    ) => {
      for (const hs of handlers) {
        const bounds = computeBounds(hs.nodes);
        const offX = bounds.minX;
        const offY = bounds.minY;
        const shiftedNodes = shiftNodes(hs.nodes, offX, offY);
        const nodeIds = new Set<string>();
        collectNodeIds(hs.nodes, nodeIds);
        const edges = allEdges
          .filter(e => nodeIds.has(e.sourceId) && nodeIds.has(e.targetId))
          .map(e => ({
            ...e,
            points: e.points.map(p => [p[0] - offX, p[1] - offY]),
          }));
        const sectionHeight = bounds.maxY - bounds.minY;
        const sectionWidth = bounds.maxX - bounds.minX;
        sections.push({
          label: hs.label,
          nodes: shiftedNodes,
          edges,
          offsetY: currentY,
          variant,
        });
        currentY += sectionHeight + SECTION_GAP;
        if (sectionWidth > maxWidth) maxWidth = sectionWidth;
      }
    };

    // Completion handlers first (above error handlers)
    addHandlerSections(completionSections, 'completion');
    // Then error handlers
    addHandlerSections(errorSections, 'error');

    const totalWidth = Math.max(layout.width ?? 0, mainWidth, maxWidth);
    const totalHeight = currentY;

    return { sections, totalWidth, totalHeight };
  }, [layout]);

  return {
    ...result,
    isLoading: preloadedLayout ? false : isLoading,
    error: preloadedLayout ? null : error,
  };
}

/** Shift all node coordinates by subtracting an offset, recursively. */
function shiftNodes(nodes: DiagramNode[], offX: number, offY: number): DiagramNode[] {
  return nodes.map(n => ({
    ...n,
    x: (n.x ?? 0) - offX,
    y: (n.y ?? 0) - offY,
    children: n.children ? shiftNodes(n.children, offX, offY) : undefined,
  }));
}

function collectNodeIds(nodes: DiagramNode[], set: Set<string>) {
  for (const n of nodes) {
    if (n.id) set.add(n.id);
    if (n.children) collectNodeIds(n.children, set);
  }
}

function computeBounds(nodes: DiagramNode[]): {
  minX: number; minY: number; maxX: number; maxY: number;
} {
  let minX = Infinity, minY = Infinity, maxX = 0, maxY = 0;
  for (const n of nodes) {
    const x = n.x ?? 0;
    const y = n.y ?? 0;
    const w = n.width ?? 80;
    const h = n.height ?? 40;
    if (x < minX) minX = x;
    if (y < minY) minY = y;
    if (x + w > maxX) maxX = x + w;
    if (y + h > maxY) maxY = y + h;
    if (n.children) {
      const childBounds = computeBounds(n.children);
      if (childBounds.maxX > maxX) maxX = childBounds.maxX;
      if (childBounds.maxY > maxY) maxY = childBounds.maxY;
    }
  }
  return { minX: minX === Infinity ? 0 : minX, minY: minY === Infinity ? 0 : minY, maxX, maxY };
}