aryx/src/renderer/components/pattern-graph/PatternGraphCanvas.tsx

import { useCallback, useEffect, useMemo, useRef } from 'react';
import {
  ReactFlow,
  ReactFlowProvider,
  Background,
  BackgroundVariant,
  Panel,
  MarkerType,
  useNodesState,
  useEdgesState,
  useReactFlow,
  type Node,
  type Edge,
  type OnConnect,
  type OnEdgesChange,
  type OnNodesChange,
} from '@xyflow/react';
import '@xyflow/react/dist/style.css';
import { LayoutGrid } from 'lucide-react';

import type { OrchestrationMode, PatternDefinition, PatternGraph } from '@shared/domain/pattern';
import { resolvePatternGraph } from '@shared/domain/pattern';
import type { ModelDefinition } from '@shared/domain/models';
import {
  addEdge,
  autoLayoutGraph,
  fromCanvasPositions,
  isConnectionAllowed,
  isEdgeDeletionAllowed,
  removeEdge,
  toCanvasEdges,
  toCanvasNodes,
  type GraphNodeData,
} from '@renderer/lib/patternGraph';

import { graphNodeTypes } from './GraphNodes';

interface PatternGraphCanvasProps {
  pattern: PatternDefinition;
  availableModels?: ReadonlyArray<ModelDefinition>;
  onGraphChange: (graph: PatternGraph) => void;
  onAgentRemove: (agentId: string) => void;
  onNodeSelect: (nodeId: string | null) => void;
  selectedNodeId: string | null;
}

function PatternGraphCanvasInner({
  pattern,
  availableModels,
  onGraphChange,
  onAgentRemove,
  onNodeSelect,
  selectedNodeId,
}: PatternGraphCanvasProps) {
  const { fitView } = useReactFlow();
  const graph = useMemo(() => resolvePatternGraph(pattern), [pattern]);
  const draggingRef = useRef(false);

  const [nodes, setNodes, onNodesChange] = useNodesState(
    toCanvasNodes(graph, pattern.agents, availableModels),
  );
  const [edges, setEdges, onEdgesChange] = useEdgesState(
    toCanvasEdges(graph, pattern.mode),
  );

  // Sync canvas when pattern changes externally
  useEffect(() => {
    setNodes(toCanvasNodes(graph, pattern.agents, availableModels));
    setEdges(toCanvasEdges(graph, pattern.mode));
  }, [graph, pattern.agents, pattern.mode, availableModels, setNodes, setEdges]);

  const handleNodesChange: OnNodesChange<Node<GraphNodeData>> = useCallback(
    (changes) => {
      // Intercept node removals and route agent deletions through the
      // authoritative removal path (which also removes from pattern.agents).
      const removals = changes.filter((c) => c.type === 'remove');
      const nonRemovals = changes.filter((c) => c.type !== 'remove');

      for (const removal of removals) {
        if (removal.type === 'remove') {
          const graphNode = graph.nodes.find((n) => n.id === removal.id);
          if (graphNode?.kind === 'agent' && graphNode.agentId) {
            onAgentRemove(graphNode.agentId);
          }
        }
      }

      if (nonRemovals.length > 0) {
        onNodesChange(nonRemovals);
      }

      const hasDragStart = nonRemovals.some(
        (c) => c.type === 'position' && 'dragging' in c && c.dragging,
      );
      const hasDragStop = nonRemovals.some(
        (c) => c.type === 'position' && !('dragging' in c && c.dragging),
      );

      if (hasDragStart) {
        draggingRef.current = true;
      }

      if (hasDragStop && draggingRef.current) {
        draggingRef.current = false;
        setNodes((currentNodes) => {
          const updatedGraph = fromCanvasPositions(pattern, currentNodes);
          onGraphChange(updatedGraph);
          return currentNodes;
        });
      }
    },
    [onNodesChange, pattern, graph, onGraphChange, onAgentRemove, setNodes],
  );

  const handleEdgesChange: OnEdgesChange = useCallback(
    (changes) => {
      // Route edge removals through the authoritative graph.
      // Non-deletable edges are already protected by the per-edge `deletable`
      // flag, so React Flow will not emit removal changes for them.
      const removals = changes.filter((c) => c.type === 'remove');
      if (removals.length > 0 && isEdgeDeletionAllowed(pattern.mode)) {
        let updatedGraph = graph;
        for (const removal of removals) {
          if (removal.type === 'remove') {
            updatedGraph = removeEdge(updatedGraph, removal.id);
          }
        }
        onGraphChange(updatedGraph);
      }

      const nonRemovals = changes.filter((c) => c.type !== 'remove');
      if (nonRemovals.length > 0) {
        onEdgesChange(nonRemovals);
      }
    },
    [onEdgesChange, pattern.mode, graph, onGraphChange],
  );

  const handleConnect: OnConnect = useCallback(
    (connection) => {
      if (!isConnectionAllowed(connection, pattern.mode, graph)) {
        return;
      }

      if (connection.source && connection.target) {
        const updatedGraph = addEdge(graph, connection.source, connection.target);
        onGraphChange(updatedGraph);
      }
    },
    [graph, pattern.mode, onGraphChange],
  );

  const handleNodeClick = useCallback(
    (_event: React.MouseEvent, node: Node<GraphNodeData>) => {
      onNodeSelect(node.id);
    },
    [onNodeSelect],
  );

  const handlePaneClick = useCallback(() => {
    onNodeSelect(null);
  }, [onNodeSelect]);

  const handleAutoLayout = useCallback(() => {
    const layouted = autoLayoutGraph(graph);
    onGraphChange(layouted);
    // Allow React to render the new positions before fitting
    requestAnimationFrame(() => fitView({ padding: 0.3 }));
  }, [graph, onGraphChange, fitView]);

  return (
    <div className="h-full w-full rounded-xl border border-[var(--color-border)] bg-[var(--color-surface-0)]/50">
      <ReactFlow
        nodes={nodes.map((n) => ({
          ...n,
          selected: n.id === selectedNodeId,
        }))}
        edges={edges}
        onNodesChange={handleNodesChange}
        onEdgesChange={handleEdgesChange}
        onConnect={handleConnect}
        onNodeClick={handleNodeClick}
        onPaneClick={handlePaneClick}
        nodeTypes={graphNodeTypes}
        fitView
        fitViewOptions={{ padding: 0.3 }}
        minZoom={0.3}
        maxZoom={2}
        proOptions={{ hideAttribution: true }}
        defaultEdgeOptions={{
          type: 'default',
          style: { stroke: '#245CF9', strokeWidth: 1.5 },
          markerEnd: { type: MarkerType.ArrowClosed, width: 16, height: 16, color: '#245CF9' },
        }}
        connectionLineStyle={{ stroke: '#245CF9', strokeWidth: 1.5 }}
        deleteKeyCode="Delete"
      >
        <Background variant={BackgroundVariant.Dots} gap={20} size={1} color="#1a1e2e" />
        <Panel position="top-right">
          <button
            type="button"
            onClick={handleAutoLayout}
            className="flex items-center gap-1.5 rounded-lg border border-[var(--color-border)] bg-[var(--color-surface-2)]/90 px-2.5 py-1.5 text-[11px] font-medium text-[var(--color-text-secondary)] shadow-sm backdrop-blur transition hover:border-[var(--color-border-glow)] hover:bg-[var(--color-surface-3)]/90 hover:text-[var(--color-text-primary)]"
            title="Auto-layout nodes"
          >
            <LayoutGrid className="size-3.5" />
            Auto layout
          </button>
        </Panel>
      </ReactFlow>
    </div>
  );
}

export function PatternGraphCanvas(props: PatternGraphCanvasProps) {
  return (
    <ReactFlowProvider>
      <PatternGraphCanvasInner {...props} />
    </ReactFlowProvider>
  );
}