import { useCallback, useMemo, useEffect, useState, useRef } from "react"; import ReactFlow, { Background, Controls, MiniMap, Connection, NodeTypes, EdgeTypes, BackgroundVariant, useNodesState, useEdgesState, addEdge, Node, Edge, NodeChange, EdgeChange, ConnectionMode, useReactFlow, Viewport, useOnSelectionChange, } from "reactflow"; import "reactflow/dist/style.css"; import { useGraphWithHistory } from "../../hooks/useGraphWithHistory"; import { useDocumentHistory } from "../../hooks/useDocumentHistory"; import { useEditorStore } from "../../stores/editorStore"; import { useActiveDocument } from "../../stores/workspace/useActiveDocument"; import { useWorkspaceStore } from "../../stores/workspaceStore"; import CustomNode from "../Nodes/CustomNode"; import CustomEdge from "../Edges/CustomEdge"; import ContextMenu from "./ContextMenu"; import EmptyState from "../Common/EmptyState"; import { createNode } from "../../utils/nodeUtils"; import DeleteIcon from "@mui/icons-material/Delete"; import { useConfirm } from "../../hooks/useConfirm"; import type { Actor, Relation } from "../../types"; interface GraphEditorProps { selectedNode: Actor | null; selectedEdge: Relation | null; onNodeSelect: (node: Actor | null) => void; onEdgeSelect: (edge: Relation | null) => void; onAddNodeRequest?: (nodeTypeId: string, position?: { x: number; y: number }) => void; } /** * GraphEditor - Main interactive graph visualization component * * Features: * - Interactive node dragging and positioning * - Edge creation by dragging from node handles * - Background grid * - MiniMap for navigation * - Zoom and pan controls * - Synchronized with workspace active document * * Usage: Core component that wraps React Flow with custom nodes and edges */ const GraphEditor = ({ onNodeSelect, onEdgeSelect, onAddNodeRequest }: GraphEditorProps) => { // Sync with workspace active document const { activeDocumentId } = useActiveDocument(); const { saveViewport, getViewport, createDocument } = useWorkspaceStore(); const { nodes: storeNodes, edges: storeEdges, nodeTypes: nodeTypeConfigs, edgeTypes: edgeTypeConfigs, setNodes, setEdges, addEdge: addEdgeWithHistory, addNode: addNodeWithHistory, deleteNode, deleteEdge, } = useGraphWithHistory(); const { pushToHistory } = useDocumentHistory(); const { showGrid, snapToGrid, gridSize, panOnDrag, zoomOnScroll, selectedRelationType, } = useEditorStore(); // React Flow instance for screen-to-flow coordinates and viewport control const { screenToFlowPosition, setViewport, getViewport: getCurrentViewport, } = useReactFlow(); // Track previous document ID to save viewport before switching const prevDocumentIdRef = useRef(null); // Confirmation dialog const { confirm, ConfirmDialogComponent } = useConfirm(); // React Flow state (synchronized with store) const [nodes, setNodesState, onNodesChange] = useNodesState( storeNodes as Node[], ); const [edges, setEdgesState, onEdgesChange] = useEdgesState( storeEdges as Edge[], ); // Track if a drag is in progress to capture state before drag const dragInProgressRef = useRef(false); // Track pending selection (ID of item to select after next sync) const pendingSelectionRef = useRef<{ type: 'node' | 'edge', id: string } | null>(null); // Context menu state const [contextMenu, setContextMenu] = useState<{ x: number; y: number; type: "pane" | "node" | "edge"; target?: Node | Edge; } | null>(null); // Sync store changes to React Flow state // IMPORTANT: Preserve selection state, unless we have a pending selection (new item added) useEffect(() => { const hasPendingSelection = pendingSelectionRef.current !== null; const pendingType = pendingSelectionRef.current?.type; const pendingId = pendingSelectionRef.current?.id; setNodesState((currentNodes) => { // If we have a pending selection, deselect all nodes (or select the new node) if (hasPendingSelection) { const pendingNodeId = pendingType === 'node' ? pendingId : null; return (storeNodes as Node[]).map((node) => ({ ...node, selected: node.id === pendingNodeId, })); } // Otherwise, preserve existing selection state const selectionMap = new Map( currentNodes.map((node) => [node.id, node.selected]) ); return (storeNodes as Node[]).map((node) => ({ ...node, selected: selectionMap.get(node.id) || false, })); }); setEdgesState((currentEdges) => { // If we have a pending selection, deselect all edges (or select the new edge) if (hasPendingSelection) { const pendingEdgeId = pendingType === 'edge' ? pendingId : null; const newEdges = (storeEdges as Edge[]).map((edge) => ({ ...edge, selected: edge.id === pendingEdgeId, })); // Clear pending selection after applying it to both nodes and edges pendingSelectionRef.current = null; return newEdges; } // Otherwise, preserve existing selection state const selectionMap = new Map( currentEdges.map((edge) => [edge.id, edge.selected]) ); return (storeEdges as Edge[]).map((edge) => ({ ...edge, selected: selectionMap.get(edge.id) || false, })); }); }, [storeNodes, storeEdges, setNodesState, setEdgesState]); // Save viewport when switching documents and restore viewport for new document useEffect(() => { if (!activeDocumentId) return; // Save viewport for the previous document if ( prevDocumentIdRef.current && prevDocumentIdRef.current !== activeDocumentId ) { const currentViewport = getCurrentViewport(); saveViewport(prevDocumentIdRef.current, currentViewport); console.log( `Saved viewport for document: ${prevDocumentIdRef.current}`, currentViewport, ); } // Restore viewport for the new document const savedViewport = getViewport(activeDocumentId); if (savedViewport) { console.log( `Restoring viewport for document: ${activeDocumentId}`, savedViewport, ); setViewport(savedViewport, { duration: 0 }); } // Update the ref to current document prevDocumentIdRef.current = activeDocumentId; }, [ activeDocumentId, saveViewport, getViewport, setViewport, getCurrentViewport, ]); // Save viewport periodically (debounced) const handleViewportChange = useCallback( (_event: MouseEvent | TouchEvent | null, viewport: Viewport) => { if (!activeDocumentId) return; // Debounce viewport saves const timeoutId = setTimeout(() => { saveViewport(activeDocumentId, viewport); }, 500); return () => clearTimeout(timeoutId); }, [activeDocumentId, saveViewport], ); // Handle selection changes using ReactFlow's dedicated hook const handleSelectionChange = useCallback( ({ nodes: selectedNodes, edges: selectedEdges, }: { nodes: Node[]; edges: Edge[]; }) => { // If a node is selected, notify parent if (selectedNodes.length == 1) { const selectedNode = selectedNodes[0] as Actor; onNodeSelect(selectedNode); // Don't call onEdgeSelect - parent will handle clearing edge selection } // If an edge is selected, notify parent else if (selectedEdges.length == 1) { const selectedEdge = selectedEdges[0] as Relation; onEdgeSelect(selectedEdge); // Don't call onNodeSelect - parent will handle clearing node selection } // Nothing selected else { onNodeSelect(null); onEdgeSelect(null); } }, [onNodeSelect, onEdgeSelect], ); // Register the selection change handler with ReactFlow useOnSelectionChange({ onChange: handleSelectionChange, }); // Sync React Flow state back to store when nodes/edges change // IMPORTANT: This handler tracks drag operations for undo/redo const handleNodesChange = useCallback( (changes: NodeChange[]) => { // Check if a drag operation just started (dragging: true) const dragStartChanges = changes.filter( (change) => change.type === "position" && "dragging" in change && change.dragging === true, ); // Capture state BEFORE the drag operation begins (for undo/redo) // This ensures we can restore to the position before dragging if (dragStartChanges.length > 0 && !dragInProgressRef.current) { dragInProgressRef.current = true; // Capture the state before any changes are applied pushToHistory("Move Actor"); } // Apply the changes onNodesChange(changes); // Check if any drag operation just completed (dragging: false) const dragEndChanges = changes.filter( (change) => change.type === "position" && "dragging" in change && change.dragging === false, ); // If a drag just ended, sync to store if (dragEndChanges.length > 0) { dragInProgressRef.current = false; // Debounce to allow React Flow state to settle setTimeout(() => { // Sync to store - use callback to get fresh state setNodesState((currentNodes) => { setNodes(currentNodes as Actor[]); return currentNodes; }); }, 0); } else { // For non-drag changes (dimension, etc), just sync to store const hasNonSelectionChanges = changes.some( (change) => change.type !== "select" && change.type !== "remove" && change.type !== "position", ); if (hasNonSelectionChanges) { setTimeout(() => { setNodesState((currentNodes) => { setNodes(currentNodes as Actor[]); return currentNodes; }); }, 0); } } }, [onNodesChange, setNodesState, setNodes, pushToHistory], ); const handleEdgesChange = useCallback( (changes: EdgeChange[]) => { onEdgesChange(changes); // Only sync to store for non-selection changes const hasNonSelectionChanges = changes.some( (change) => change.type !== "select" && change.type !== "remove", ); if (hasNonSelectionChanges) { // Debounce store updates to avoid loops setTimeout(() => { setEdgesState((currentEdges) => { setEdges(currentEdges as Relation[]); return currentEdges; }); }, 0); } }, [onEdgesChange, setEdgesState, setEdges], ); // Handle new edge connections const handleConnect = useCallback( (connection: Connection) => { if (!connection.source || !connection.target) return; // Use selected relation type or fall back to first available const edgeType = selectedRelationType || edgeTypeConfigs[0]?.id || "default"; // Create edge with custom data (no label - will use type default) const edgeWithData = { ...connection, type: "custom", data: { type: edgeType, // Don't set label - will use type's label as default }, }; // Use React Flow's addEdge helper to properly format the edge const updatedEdges = addEdge(edgeWithData, storeEdges as Edge[]); // Find the newly added edge (it will be the last one) const newEdge = updatedEdges[updatedEdges.length - 1] as Relation; // Set pending selection - will be applied after Zustand sync pendingSelectionRef.current = { type: 'edge', id: newEdge.id }; // Use the history-tracked addEdge function (triggers sync which will apply selection) addEdgeWithHistory(newEdge); }, [ storeEdges, edgeTypeConfigs, addEdgeWithHistory, selectedRelationType, ], ); // Handle node deletion const handleNodesDelete = useCallback( (nodesToDelete: Node[]) => { nodesToDelete.forEach((node) => { deleteNode(node.id); }); }, [deleteNode], ); // Handle edge deletion const handleEdgesDelete = useCallback( (edgesToDelete: Edge[]) => { edgesToDelete.forEach((edge) => { deleteEdge(edge.id); }); }, [deleteEdge], ); // Register custom node types const nodeTypes: NodeTypes = useMemo( () => ({ custom: CustomNode, }), [], ); // Register custom edge types const edgeTypes: EdgeTypes = useMemo( () => ({ custom: CustomEdge, }), [], ); // Handle node click - ReactFlow handles selection automatically const handleNodeClick = useCallback(() => { setContextMenu(null); // Close context menu if open }, []); // Handle edge click - ReactFlow handles selection automatically const handleEdgeClick = useCallback(() => { setContextMenu(null); // Close context menu if open }, []); // Handle right-click on pane (empty space) const handlePaneContextMenu = useCallback((event: React.MouseEvent) => { event.preventDefault(); setContextMenu({ x: event.clientX, y: event.clientY, type: "pane", }); }, []); // Handle right-click on node const handleNodeContextMenu = useCallback( (event: React.MouseEvent, node: Node) => { event.preventDefault(); setContextMenu({ x: event.clientX, y: event.clientY, type: "node", target: node, }); }, [], ); // Handle right-click on edge const handleEdgeContextMenu = useCallback( (event: React.MouseEvent, edge: Edge) => { event.preventDefault(); setContextMenu({ x: event.clientX, y: event.clientY, type: "edge", target: edge, }); }, [], ); // Handle left-click on pane to close context menu const handlePaneClick = useCallback(() => { if (contextMenu) { setContextMenu(null); } }, [contextMenu]); // Shared node creation logic (used by context menu and left panel) const handleAddNode = useCallback( (nodeTypeId: string, position?: { x: number; y: number }) => { // Use provided position or random position for toolbar/panel const nodePosition = position || { x: Math.random() * 400 + 100, y: Math.random() * 300 + 100, }; const nodeTypeConfig = nodeTypeConfigs.find((nt) => nt.id === nodeTypeId); const newNode = createNode(nodeTypeId, nodePosition, nodeTypeConfig); // Set pending selection - will be applied after Zustand sync pendingSelectionRef.current = { type: 'node', id: newNode.id }; // Use history-tracked addNode (triggers sync which will apply selection) addNodeWithHistory(newNode); }, [nodeTypeConfigs, addNodeWithHistory], ); // Call the onAddNodeRequest callback if provided useEffect(() => { if (onAddNodeRequest) { onAddNodeRequest(handleAddNode); } }, [onAddNodeRequest, handleAddNode]); // Add new actor at context menu position const handleAddActorFromContextMenu = useCallback( (nodeTypeId: string) => { if (!contextMenu) return; const position = screenToFlowPosition({ x: contextMenu.x, y: contextMenu.y, }); handleAddNode(nodeTypeId, position); setContextMenu(null); }, [contextMenu, screenToFlowPosition, handleAddNode], ); // Show empty state when no document is active if (!activeDocumentId) { return ( createDocument()} onOpenDocumentManager={() => { // This will be handled by the parent component // We'll trigger it via a custom event window.dispatchEvent(new CustomEvent("openDocumentManager")); }} /> ); } return (
{/* Background grid */} {showGrid && ( )} {/* Zoom and pan controls */} {/* MiniMap for navigation */} { const actor = node as Actor; const nodeType = nodeTypeConfigs.find( (nt) => nt.id === actor.data?.type, ); return nodeType?.color || "#6b7280"; }} pannable zoomable /> {/* Context Menu - Pane */} {contextMenu && contextMenu.type === "pane" && ( ({ label: nodeType.label, color: nodeType.color, onClick: () => handleAddActorFromContextMenu(nodeType.id), })), }, ]} onClose={() => setContextMenu(null)} /> )} {/* Context Menu - Node */} {contextMenu && contextMenu.type === "node" && contextMenu.target && ( , onClick: async () => { const confirmed = await confirm({ title: "Delete Actor", message: "Are you sure you want to delete this actor? All connected relations will also be deleted.", confirmLabel: "Delete", severity: "danger", }); if (confirmed) { deleteNode(contextMenu.target!.id); setContextMenu(null); } }, }, ], }, ]} onClose={() => setContextMenu(null)} /> )} {/* Context Menu - Edge */} {contextMenu && contextMenu.type === "edge" && contextMenu.target && ( , onClick: async () => { const confirmed = await confirm({ title: "Delete Relation", message: "Are you sure you want to delete this relation?", confirmLabel: "Delete", severity: "danger", }); if (confirmed) { deleteEdge(contextMenu.target!.id); setContextMenu(null); } }, }, ], }, ]} onClose={() => setContextMenu(null)} /> )} {/* Confirmation Dialog */} {ConfirmDialogComponent}
); }; export default GraphEditor;