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

const SECTION_GAP = 40;

export function useDiagramData(
  application: string,
  routeId: string,
  direction: 'LR' | 'TB' = 'LR',
) {
  const { data: layout, isLoading, error } = useDiagramByRoute(application, routeId, direction);

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

    const allEdges = layout.edges ?? [];

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

    for (const node of layout.nodes) {
      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),
    );

    // Compute main section bounding box
    const mainBounds = computeBounds(mainNodes);

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

    let currentY = mainBounds.maxY + SECTION_GAP;
    let maxWidth = mainBounds.maxX;

    for (const es of errorSections) {
      const errorBounds = computeBounds(es.nodes);
      const offX = errorBounds.minX;
      const offY = errorBounds.minY;

      // Normalize node coordinates relative to the section's own origin
      const shiftedNodes = shiftNodes(es.nodes, offX, offY);

      const errorNodeIds = new Set<string>();
      collectNodeIds(es.nodes, errorNodeIds);

      // Shift edge points too
      const errorEdges = allEdges
        .filter(e => errorNodeIds.has(e.sourceId) && errorNodeIds.has(e.targetId))
        .map(e => ({
          ...e,
          points: e.points.map(p => [p[0] - offX, p[1] - offY]),
        }));

      const sectionHeight = errorBounds.maxY - errorBounds.minY;
      const sectionWidth = errorBounds.maxX - errorBounds.minX;

      sections.push({
        label: es.label,
        nodes: shiftedNodes,
        edges: errorEdges,
        offsetY: currentY,
        variant: 'error',
      });

      currentY += sectionHeight + SECTION_GAP;
      if (sectionWidth > maxWidth) maxWidth = sectionWidth;
    }

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

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

  return { ...result, isLoading, 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 };
}