@flowuent-org/diagramming-core 1.1.0 → 1.1.2

package/packages/diagrams/src/lib/templates/systemFlow/components/SystemFlowDiagramFlow.tsx

@@ -0,0 +1,1179 @@
+import React, {
+    useCallback,
+    useState,
+    useRef,
+    useEffect,
+    useMemo,
+} from 'react';
+
+import {
+    ReactFlow,
+    Background,
+    useNodesState,
+    useEdgesState,
+    useReactFlow,
+    Panel,
+    BackgroundVariant,
+    Connection,
+    NodeChange,
+    applyNodeChanges,
+    ReactFlowProvider,
+    Node,
+    Edge,
+    NodeProps,
+} from '@xyflow/react';
+import '@xyflow/react/dist/style.css';
+import {
+    SystemFlowNode,
+    SystemFlowEdge,
+    HistoryState,
+    EditingState,
+    SystemFlowNodeSettingsState,
+    SystemFlowEdgeSettingsState,
+    NodeSelectionContextType,
+    LayoutDirection,
+    UpdateFunction,
+    CloseFunction,
+    DeleteFunction
+} from '../../../types/system-flow-types';
+
+import {
+    Box,
+    Button,
+    Paper,
+    ToggleButton,
+    ToggleButtonGroup,
+    Snackbar,
+    Alert,
+} from '@mui/material';
+import DownloadIcon from '@mui/icons-material/Download';
+import UploadIcon from '@mui/icons-material/Upload';
+import FullscreenIcon from '@mui/icons-material/Fullscreen';
+import FullscreenExitIcon from '@mui/icons-material/FullscreenExit';
+import CloseIcon from '@mui/icons-material/Close';
+import UndoIcon from '@mui/icons-material/Undo';
+import RedoIcon from '@mui/icons-material/Redo';
+import ViewWeekIcon from '@mui/icons-material/ViewWeek';
+import ViewStreamIcon from '@mui/icons-material/ViewStream';
+import {
+    systemFlowDefaultNodes,
+    systemFlowDefaultEdges,
+    treeStructureTemplate,
+    treeStructureTemplate1,
+    treeStructureTemplate2,
+    treeStructureTemplate3,
+} from '../data/template-data';
+import {
+    systemFlowEdgeTypes,
+    systemFlowNodeTypes as baseSystemFlowNodeTypes,
+} from './SystemFlowComponents';
+import DiagramControls from '../controls/SystemFlowControls';
+import { NodeSelectionContext } from '../controls/SystemFlowControls';
+import SystemFlowNodeSettings from './SystemFlowNodeSettings';
+import EdgeSettings from './EdgeSettings';
+import DiagramPanel from '../../../components/DiagramPanel';
+import Sidebar from '../../../components/SideBar';
+import AiButton from '../../../components/AiButton';
+import MiniMapCommon from '../../../components/MiniMapCommon';
+
+// Browser-specific fullscreen API types
+interface WebkitElement extends Element {
+    webkitRequestFullscreen(): Promise<void>;
+}
+
+interface WebkitDocument extends Document {
+    webkitExitFullscreen(): Promise<void>;
+}
+
+interface MSElement extends Element {
+    msRequestFullscreen(): Promise<void>;
+}
+
+interface MSDocument extends Document {
+    msExitFullscreen(): Promise<void>;
+}
+
+interface SystemFlowDiagramFlowProps {
+    nodes: SystemFlowNode[];
+    setNodes: (nodes: SystemFlowNode[] | ((prev: SystemFlowNode[]) => SystemFlowNode[])) => void;
+    edges: SystemFlowEdge[];
+    setEdges: (edges: SystemFlowEdge[] | ((prev: SystemFlowEdge[]) => SystemFlowEdge[])) => void;
+}
+
+// Utility: Deep compare node arrays by id, type, data, width, height
+const deepCompareNodes = (a: Node[], b: Node[]) => {
+    if (a.length !== b.length) return false;
+    for (let i = 0; i < a.length; i++) {
+        const nodeA = a[i];
+        const nodeB = b[i];
+        if (
+            nodeA.id !== nodeB.id ||
+            nodeA.type !== nodeB.type ||
+            JSON.stringify(nodeA.data) !== JSON.stringify(nodeB.data)
+        ) {
+            return false;
+        }
+    }
+    return true;
+};
+
+// Utility: Deep compare edge arrays by id, source, target, type, data
+const deepCompareEdges = (a: Edge[], b: Edge[]) => {
+    if (a.length !== b.length) return false;
+    for (let i = 0; i < a.length; i++) {
+        const edgeA = a[i];
+        const edgeB = b[i];
+        if (
+            edgeA.id !== edgeB.id ||
+            edgeA.source !== edgeB.source ||
+            edgeA.target !== edgeB.target ||
+            edgeA.type !== edgeB.type ||
+            JSON.stringify(edgeA.data) !== JSON.stringify(edgeB.data)
+        ) {
+            return false;
+        }
+    }
+    return true;
+};
+
+// Debounce utility
+const debounce = (fn: Function, delay: number) => {
+    let timer: NodeJS.Timeout | null = null;
+    return (...args: unknown[]) => {
+        if (timer) clearTimeout(timer);
+        timer = setTimeout(() => fn(...args), delay);
+    };
+};
+
+// Helper: get center position in diagram
+const getCenterPosition = (width: number, height: number) => {
+    return { x: width / 2 - 75, y: height / 2 - 40 };
+};
+
+const SystemFlowDiagramFlowContent: React.FC<
+    SystemFlowDiagramFlowProps
+> = React.memo(({
+    nodes: initialNodes,
+    setNodes: parentSetNodes,
+    edges: initialEdges,
+    setEdges: parentSetEdges,
+}) => {
+    // State declarations
+    const [nodes, setNodes, onNodesChange] = useNodesState(initialNodes as Node[]);
+    const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges as Edge[]);
+
+    // Custom onNodesChange handler for synchronized tree structure movement
+    const customOnNodesChange = useCallback((changes: NodeChange[]) => {
+        // First apply the default changes
+        onNodesChange(changes);
+
+        // Then handle synchronized movement for tree structure nodes
+        changes.forEach((change: NodeChange) => {
+            if (change.type === 'position' && 'position' in change && change.position) {
+                const position = change.position as { x: number; y: number };
+                const node = nodes.find(n => n.id === change.id);
+
+                if (node && node.data.isTreeStructure && node.data.structureId) {
+                    // This is a tree structure node being moved, move all nodes in the same structure
+                    const currentNode = nodes.find(n => n.id === change.id);
+                    if (!currentNode) return;
+
+                    // Calculate the delta movement
+                    const deltaX = position.x - currentNode.position.x;
+                    const deltaY = position.y - currentNode.position.y;
+
+                    // Move all nodes in the same structure by the same delta
+                    const updatedNodes = nodes.map(n => {
+                        if (n.data.isTreeStructure && n.data.structureId === node.data.structureId) {
+                            return {
+                                ...n,
+                                position: {
+                                    x: n.position.x + deltaX,
+                                    y: n.position.y + deltaY
+                                }
+                            };
+                        }
+                        return n;
+                    });
+
+                    setNodes(updatedNodes);
+                }
+            }
+        });
+    }, [nodes, setNodes, onNodesChange]);
+    const [history, setHistory] = useState<HistoryState[]>([]);
+    const [historyIndex, setHistoryIndex] = useState<number>(-1);
+    const isUndoRedoAction = useRef(false);
+    const [isFullscreen, setIsFullscreen] = useState<boolean>(false);
+    const diagramContainerRef = useRef<HTMLDivElement>(null);
+    const [editing, setEditing] = useState<{
+        nodeId?: string | null;
+        edgeId?: string | null;
+        label?: string;
+        type?: 'node' | 'edge';
+    }>({});
+    const [systemFlowNodeSettings, setSystemFlowNodeSettings] = useState<{
+        id: string | null;
+    }>({ id: null });
+    const [edgeSettings, setEdgeSettings] = useState<{ id: string | null }>({
+        id: null,
+    });
+    const [isSettingsSidebarOpen, setIsSettingsSidebarOpen] = useState(false);
+
+    // Track if a child node has been edited and needs layout recalculation
+    const [pendingLayoutRecalculation, setPendingLayoutRecalculation] = useState<{
+        structureId: string;
+        rootNodeId: string;
+    } | null>(null);
+
+    // Function declarations in correct order
+    const saveLabel = useCallback(() => {
+        if (editing.nodeId && editing.type === 'node') {
+            setNodes((nds: Node[]) =>
+                nds.map((node: Node) => {
+                    if (node.id === editing.nodeId) {
+                        return {
+                            ...node,
+                            data: {
+                                ...node.data,
+                                label: editing.label,
+                            },
+                        };
+                    }
+                    return node;
+                }),
+            );
+            setEditing({
+                nodeId: undefined,
+                edgeId: undefined,
+                label: undefined,
+                type: undefined,
+            });
+        } else if (editing.edgeId && editing.type === 'edge') {
+            setEdges((eds: Edge[]) =>
+                eds.map((edge: Edge) => {
+                    if (edge.id === editing.edgeId) {
+                        return {
+                            ...edge,
+                            data: {
+                                ...edge.data,
+                                label: editing.label,
+                            },
+                        };
+                    }
+                    return edge;
+                }),
+            );
+            setEditing({
+                nodeId: undefined,
+                edgeId: undefined,
+                label: undefined,
+                type: undefined,
+            });
+        }
+    }, [editing, setNodes, setEdges]);
+
+    const handleKeyDown = useCallback(
+        (e: React.KeyboardEvent) => {
+            if (e.key === 'Enter') {
+                saveLabel();
+            } else if (e.key === 'Escape') {
+                setEditing({
+                    nodeId: undefined,
+                    edgeId: undefined,
+                    label: undefined,
+                    type: undefined,
+                });
+            }
+        },
+        [saveLabel],
+    );
+
+    const loadTemplateData = useCallback(
+        (templateType: string) => {
+
+            if (templateType === 'tree-structure') {
+                setNodes(treeStructureTemplate.nodes as Node[]);
+                setEdges(treeStructureTemplate.edges as Edge[]);
+            } else if (templateType === 'structure1') {
+                setNodes(treeStructureTemplate1.nodes as Node[]);
+                setEdges(treeStructureTemplate1.edges as Edge[]);
+            } else if (templateType === 'structure2') {
+                setNodes(treeStructureTemplate2.nodes as Node[]);
+                setEdges(treeStructureTemplate2.edges as Edge[]);
+            } else if (templateType === 'structure3') {
+                setNodes(treeStructureTemplate3.nodes as Node[]);
+                setEdges(treeStructureTemplate3.edges as Edge[]);
+            } else {
+                // Default system flow nodes/edges
+                setNodes(systemFlowDefaultNodes as Node[]);
+                setEdges(systemFlowDefaultEdges as Edge[]);
+            }
+        },
+        [setNodes, setEdges],
+    );
+
+    // Initialize nodes and edges with our data
+    useEffect(() => {
+        // Only initialize if history is empty
+        if (history.length === 0) {
+            const initialState: HistoryState = {
+                nodes: nodes as SystemFlowNode[],
+                edges: edges as SystemFlowEdge[],
+            };
+            setHistory([initialState]);
+            setHistoryIndex(0);
+        }
+    }, []);
+
+    // Save current state to history when nodes or edges change (debounced)
+    const debouncedUpdateHistory = useRef(
+        debounce(
+            (
+                nodes: Node[],
+                edges: Edge[],
+                history: HistoryState[],
+                historyIndex: number,
+                setHistory: React.Dispatch<React.SetStateAction<HistoryState[]>>,
+                setHistoryIndex: React.Dispatch<React.SetStateAction<number>>,
+            ) => {
+                const newHistoryEntry: HistoryState = {
+                    nodes: [...nodes] as SystemFlowNode[],
+                    edges: [...edges] as SystemFlowEdge[],
+                };
+                const currentState = history[historyIndex];
+                if (
+                    currentState &&
+                    deepCompareNodes(currentState.nodes, newHistoryEntry.nodes) &&
+                    deepCompareEdges(currentState.edges, newHistoryEntry.edges)
+                ) {
+                    return; // No change, don't add to history
+                }
+                // Trim future history if we're in the middle of the history stack
+                let newHistory = history.slice(0, historyIndex + 1);
+                newHistory.push(newHistoryEntry);
+                // Limit to last 50 entries
+                if (newHistory.length > 50) {
+                    newHistory = newHistory.slice(newHistory.length - 50);
+                }
+                setHistory(newHistory);
+                setHistoryIndex(newHistory.length - 1);
+            },
+            300,
+        ),
+    ).current;
+
+    useEffect(() => {
+        if (isUndoRedoAction.current || history.length === 0) {
+            return;
+        }
+        debouncedUpdateHistory(
+            nodes,
+            edges,
+            history,
+            historyIndex,
+            setHistory,
+            setHistoryIndex,
+        );
+        // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, [nodes, edges]);
+
+    // Function to perform undo operation
+    const handleUndo = useCallback(() => {
+        if (historyIndex > 0) {
+            isUndoRedoAction.current = true;
+            const prevState = history[historyIndex - 1];
+            setNodes(prevState.nodes);
+            setEdges(prevState.edges);
+            setHistoryIndex(historyIndex - 1);
+            setTimeout(() => {
+                isUndoRedoAction.current = false;
+            }, 0);
+        }
+    }, [history, historyIndex, setNodes, setEdges]);
+
+    // Function to perform redo operation
+    const handleRedo = useCallback(() => {
+        if (historyIndex < history.length - 1) {
+            isUndoRedoAction.current = true;
+            const nextState = history[historyIndex + 1];
+            setNodes(nextState.nodes);
+            setEdges(nextState.edges);
+            setHistoryIndex(historyIndex + 1);
+            setTimeout(() => {
+                isUndoRedoAction.current = false;
+            }, 0);
+        }
+    }, [history, historyIndex, setNodes, setEdges]);
+
+    // Clear handler
+    const handleClear = useCallback(() => {
+        const cleared = { nodes: [], edges: [] };
+        const newHistory = history.slice(0, historyIndex + 1);
+        newHistory.push(cleared);
+        setNodes([]);
+        setEdges([]);
+        setHistory(newHistory);
+        setHistoryIndex(newHistory.length - 1);
+    }, [history, historyIndex, setNodes, setEdges]);
+
+    // Listen for keyboard shortcuts
+    useEffect(() => {
+        const handleKeyDown = (e: KeyboardEvent) => {
+            // Check for Ctrl+Z or Cmd+Z (undo)
+            if ((e.ctrlKey || e.metaKey) && e.key === 'z' && !e.shiftKey) {
+                e.preventDefault();
+                handleUndo();
+            }
+            // Check for Ctrl+Y or Cmd+Shift+Z (redo)
+            if (
+                ((e.ctrlKey || e.metaKey) && e.key === 'y') ||
+                ((e.ctrlKey || e.metaKey) && e.shiftKey && e.key === 'z')
+            ) {
+                e.preventDefault();
+                handleRedo();
+            }
+        };
+
+        window.addEventListener('keydown', handleKeyDown);
+        return () => {
+            window.removeEventListener('keydown', handleKeyDown);
+        };
+    }, [handleUndo, handleRedo]);
+
+    // Node type map for system flow nodes
+    const nodeTypeMap: Record<string, { type: string; data?: Record<string, unknown> }> = {
+        SystemNode: { type: 'SystemNode', data: { label: 'Node', color: '#2196F3' } },
+        bendpoint: { type: 'bendpoint', data: {} },
+        Structure1: {
+            type: 'TreeSystemNode',
+            data: {
+                label: 'Structure 1',
+                color: '#2196F3',
+                isTreeStructure: true,
+                structureId: `structure-1-${Date.now()}`,
+                nodeType: 'root',
+                structureType: 'structure1'
+            }
+        },
+        Structure2: {
+            type: 'TreeSystemNode',
+            data: {
+                label: 'Structure 2',
+                color: '#2196F3',
+                isTreeStructure: true,
+                structureId: `structure-2-${Date.now()}`,
+                nodeType: 'root',
+                structureType: 'structure2'
+            }
+        },
+        Structure3: {
+            type: 'TreeSystemNode',
+            data: {
+                label: 'Structure 3',
+                color: '#2196F3',
+                isTreeStructure: true,
+                structureId: `structure-3-${Date.now()}`,
+                nodeType: 'root',
+                structureType: 'structure3'
+            }
+        },
+    };
+
+    // Add node at center (internal)
+    const addNodeAtCenter = useCallback(
+        (nodeTypeId: string) => {
+            const nodeInfo = nodeTypeMap[nodeTypeId];
+            if (!nodeInfo) return;
+            const { type, data } = nodeInfo;
+            const width = diagramContainerRef.current?.offsetWidth || 1200;
+            const height = diagramContainerRef.current?.offsetHeight || 800;
+            const center = getCenterPosition(width, height);
+            const newId = `${nodeTypeId}-${nodes.length + 1}`;
+
+            // Add onResize handler to node data
+            const onResize = (id: string, width: number, height: number) => {
+                setNodes((nds) =>
+                    nds.map((node) =>
+                        node.id === id
+                            ? {
+                                ...node,
+                                measured: { width, height },
+                                width: width,
+                                height: height,
+                            }
+                            : node,
+                    ),
+                );
+            };
+
+            setNodes((nds: Node[]) => [
+                ...nds,
+                {
+                    id: newId,
+                    type,
+                    position: center,
+                    data: { ...data, onResize },
+                    draggable: true,
+                },
+            ]);
+        },
+        [nodes, setNodes],
+    );
+
+    // Update onAddEntity to use internal addNodeAtCenter
+    const handleAddEntity = useCallback(
+        (nodeTypeId = 'SystemNode') => {
+            addNodeAtCenter(nodeTypeId);
+        },
+        [addNodeAtCenter],
+    );
+
+    // Function to export diagram as JSON
+    const exportDiagram = useCallback(() => {
+        const diagramData = {
+            nodes: nodes,
+            edges: edges,
+        };
+
+        const dataStr = JSON.stringify(diagramData, null, 2);
+        const dataUri =
+            'data:application/json;charset=utf-8,' + encodeURIComponent(dataStr);
+
+        const exportFileDefaultName = 'system-flow-diagram.json';
+
+        const linkElement = document.createElement('a');
+        linkElement.setAttribute('href', dataUri);
+        linkElement.setAttribute('download', exportFileDefaultName);
+        linkElement.click();
+    }, [nodes, edges]);
+
+    // Function to import diagram from JSON
+    const importDiagram = useCallback(
+        (event: React.ChangeEvent<HTMLInputElement>) => {
+            if (!event.target.files || event.target.files.length === 0) return;
+
+            const fileReader = new FileReader();
+            fileReader.readAsText(event.target.files[0], 'UTF-8');
+            fileReader.onload = (e) => {
+                try {
+                    if (!e.target || typeof e.target.result !== 'string') return;
+
+                    const imported = JSON.parse(e.target.result);
+                    if (imported.nodes && imported.edges) {
+                        setNodes(imported.nodes);
+                        setEdges(imported.edges);
+                    }
+                } catch (error) {
+                    alert('Failed to import diagram. Invalid file format.');
+                }
+            };
+        },
+        [setNodes, setEdges],
+    );
+
+    // Toggle fullscreen mode
+    const toggleFullscreen = useCallback(() => {
+        if (!diagramContainerRef.current) return;
+
+        if (!isFullscreen) {
+            if (diagramContainerRef.current.requestFullscreen) {
+                diagramContainerRef.current.requestFullscreen();
+            } else if ((diagramContainerRef.current as unknown as WebkitElement).webkitRequestFullscreen) {
+                (diagramContainerRef.current as unknown as WebkitElement).webkitRequestFullscreen?.();
+            } else if ((diagramContainerRef.current as unknown as MSElement).msRequestFullscreen) {
+                (diagramContainerRef.current as unknown as MSElement).msRequestFullscreen?.();
+            }
+        } else {
+            if (document.exitFullscreen) {
+                document.exitFullscreen();
+            } else if ((document as unknown as WebkitDocument).webkitExitFullscreen) {
+                (document as unknown as WebkitDocument).webkitExitFullscreen?.();
+            } else if ((document as unknown as MSDocument).msExitFullscreen) {
+                (document as unknown as MSDocument).msExitFullscreen?.();
+            }
+        }
+
+        setIsFullscreen(!isFullscreen);
+    }, [isFullscreen]);
+
+    // Listen for fullscreen change events from browser
+    useEffect(() => {
+        const handleFullscreenChange = () => {
+            setIsFullscreen(!!document.fullscreenElement);
+        };
+
+        document.addEventListener('fullscreenchange', handleFullscreenChange);
+        document.addEventListener('webkitfullscreenchange', handleFullscreenChange);
+        document.addEventListener('mozfullscreenchange', handleFullscreenChange);
+        document.addEventListener('MSFullscreenChange', handleFullscreenChange);
+
+        return () => {
+            document.removeEventListener('fullscreenchange', handleFullscreenChange);
+            document.removeEventListener(
+                'webkitfullscreenchange',
+                handleFullscreenChange,
+            );
+            document.removeEventListener(
+                'mozfullscreenchange',
+                handleFullscreenChange,
+            );
+            document.removeEventListener(
+                'MSFullscreenChange',
+                handleFullscreenChange,
+            );
+        };
+    }, []);
+
+    // Modified onConnect to open EdgeSettings after edge is created
+    const onConnect = useCallback(
+        (params: Connection) => {
+            if (!params.source || !params.target) return;
+            const sourceNode = nodes.find((n) => n.id === params.source);
+            const targetNode = nodes.find((n) => n.id === params.target);
+            if (!sourceNode || !targetNode) return;
+
+            const newEdge = {
+                id: `edge-${edges.length + 1}`,
+                source: params.source,
+                target: params.target,
+                type: 'floatingEdge',
+                sourceHandle: params.sourceHandle,
+                targetHandle: params.targetHandle,
+                data: {
+                    label: `Flow ${edges.length + 1}`,
+                    type: 'controlFlow',
+                    color: '#666',
+                },
+                animated: false,
+            };
+            setEdges((eds: Edge[]) => [...eds, newEdge]);
+            setEdgeSettings({ id: newEdge.id });
+            setIsSettingsSidebarOpen(true);
+        },
+        [nodes, edges, setEdges],
+    );
+
+    // Update onNodeClick function
+    const onNodeClick = useCallback((event: React.MouseEvent, node: Node) => {
+        setSystemFlowNodeSettings({ id: node.id });
+        setIsSettingsSidebarOpen(true);
+    }, []);
+
+    // Update onNodeDoubleClick function
+    const onNodeDoubleClick = useCallback(
+        (event: React.MouseEvent, node: Node) => {
+            event.preventDefault();
+            event.stopPropagation();
+            setEditing({
+                nodeId: node.id,
+                edgeId: null,
+                label: (node.data?.label as string) || '',
+                type: 'node',
+            });
+        },
+        [setEditing],
+    );
+
+    // Open EdgeSettings on double click
+    const onEdgeDoubleClick = useCallback(
+        (event: React.MouseEvent, edge: Edge) => {
+            event.preventDefault();
+            event.stopPropagation();
+            if (edge.type === 'floatingEdge') {
+                setEdgeSettings({ id: edge.id });
+                setIsSettingsSidebarOpen(true);
+            }
+        },
+        [],
+    );
+
+    // Function to update selected system flow node
+    const updateSelectedSystemFlowNode = useCallback(
+        (updates: Record<string, unknown>) => {
+            if (!systemFlowNodeSettings.id) return;
+
+            const currentNode = nodes.find(n => n.id === systemFlowNodeSettings.id);
+            const isChildNode = currentNode &&
+                currentNode.data.isTreeStructure &&
+                currentNode.data.nodeType === 'child' &&
+                currentNode.data.structureId;
+
+            setNodes((nds: Node[]) =>
+                nds.map((node: Node) => {
+                    if (node.id === systemFlowNodeSettings.id) {
+                        return {
+                            ...node,
+                            data: {
+                                ...node.data,
+                                ...updates,
+                            },
+                        };
+                    }
+                    return node;
+                }),
+            );
+
+            // If editing a child node, mark for layout recalculation
+            if (isChildNode) {
+                const rootNode = nodes.find(n =>
+                    n.data.isTreeStructure &&
+                    n.data.structureId === currentNode.data.structureId &&
+                    n.data.nodeType === 'root'
+                );
+
+                if (rootNode) {
+                    setPendingLayoutRecalculation({
+                        structureId: currentNode.data.structureId as string,
+                        rootNodeId: rootNode.id
+                    });
+                }
+            }
+        },
+        [systemFlowNodeSettings.id, setNodes, nodes],
+    );
+
+    // Close system flow node settings
+    const closeSystemFlowNodeSettings = useCallback(() => {
+        const currentNodeId = systemFlowNodeSettings.id;
+        setSystemFlowNodeSettings({ id: null });
+        setIsSettingsSidebarOpen(false);
+
+        // If we're closing settings for a tree structure node, handle accordingly
+        if (currentNodeId) {
+            const currentNode = nodes.find(n => n.id === currentNodeId);
+            if (currentNode && currentNode.data.isTreeStructure) {
+
+                // If it's a root node being edited, auto-center it
+                if (currentNode.data.nodeType === 'root' && currentNode.data.isRootBeingEdited) {
+                    const autoCenterRootNode = currentNode.data?.autoCenterRootNode;
+                    if (typeof autoCenterRootNode === 'function') {
+                        autoCenterRootNode();
+                    }
+                }
+
+                // If it's a child node, trigger layout recalculation
+                else if (currentNode.data.nodeType === 'child' && currentNode.data.structureId) {
+
+                    // Find the root node and trigger its layout recalculation
+                    const rootNode = nodes.find(n =>
+                        n.data.isTreeStructure &&
+                        n.data.structureId === currentNode.data.structureId &&
+                        n.data.nodeType === 'root'
+                    );
+
+                    if (rootNode) {
+                        // Call the forceLayoutRecalculation function directly if available
+                        const forceLayoutRecalculation = rootNode.data?.forceLayoutRecalculation;
+                        if (typeof forceLayoutRecalculation === 'function') {
+                            forceLayoutRecalculation();
+                        } else {
+                            // Fallback: Force a layout recalculation by updating the root node
+                            setNodes((nds: Node[]) =>
+                                nds.map((node: Node) => {
+                                    if (node.id === rootNode.id) {
+                                        return {
+                                            ...node,
+                                            data: {
+                                                ...node.data,
+                                                // Add a timestamp to force re-render and layout recalculation
+                                                lastLayoutUpdate: Date.now()
+                                            }
+                                        };
+                                    }
+                                    return node;
+                                })
+                            );
+                        }
+                    }
+                }
+            }
+        }
+
+        // Clear pending layout recalculation when settings close
+        setPendingLayoutRecalculation(null);
+    }, [systemFlowNodeSettings.id, nodes, setNodes]);
+
+    // Update edge from EdgeSettings
+    const updateSelectedEdge = useCallback(
+        (updates: Record<string, unknown>) => {
+            if (!edgeSettings.id) return;
+
+            setEdges((eds: Edge[]) => {
+                const currentEdge = eds.find(e => e.id === edgeSettings.id);
+                if (!currentEdge) return eds;
+
+                const newEdges = [...eds];
+
+                // Update the current edge
+                const edgeIndex = newEdges.findIndex(e => e.id === edgeSettings.id);
+                if (edgeIndex !== -1) {
+                    newEdges[edgeIndex] = {
+                        ...newEdges[edgeIndex],
+                        data: {
+                            ...newEdges[edgeIndex].data,
+                            ...updates,
+                        },
+                    };
+                }
+
+                // Handle two-way edge creation/removal
+                const isTwoWay = updates.isTwoWay as boolean;
+                const wasTwoWay = currentEdge.data?.isTwoWay as boolean;
+
+                if (isTwoWay && !wasTwoWay) {
+                    // Create reverse edge for two-way connection
+                    const reverseEdgeId = `reverse-${edgeSettings.id}`;
+                    const reverseEdge: Edge = {
+                        id: reverseEdgeId,
+                        source: currentEdge.target,
+                        target: currentEdge.source,
+                        type: 'floatingEdge',
+                        data: {
+                            ...currentEdge.data,
+                            ...updates,
+                            isTwoWay: true,
+                            reverseEdgeId: edgeSettings.id,
+                        },
+                    };
+
+                    // Update current edge to reference reverse edge
+                    newEdges[edgeIndex] = {
+                        ...newEdges[edgeIndex],
+                        data: {
+                            ...newEdges[edgeIndex].data,
+                            reverseEdgeId,
+                        },
+                    };
+
+                    newEdges.push(reverseEdge);
+                } else if (!isTwoWay && wasTwoWay) {
+                    // Remove reverse edge
+                    const reverseEdgeId = currentEdge.data?.reverseEdgeId;
+                    if (reverseEdgeId) {
+                        // Remove reverse edge
+                        const reverseEdgeIndex = newEdges.findIndex(e => e.id === reverseEdgeId);
+                        if (reverseEdgeIndex !== -1) {
+                            newEdges.splice(reverseEdgeIndex, 1);
+                        }
+
+                        // Remove reverse edge reference from current edge
+                        newEdges[edgeIndex] = {
+                            ...newEdges[edgeIndex],
+                            data: {
+                                ...newEdges[edgeIndex].data,
+                                reverseEdgeId: undefined,
+                            },
+                        };
+                    }
+                } else if (isTwoWay && wasTwoWay) {
+                    // Update both edges when two-way is already enabled
+                    const reverseEdgeId = currentEdge.data?.reverseEdgeId;
+                    if (reverseEdgeId) {
+                        const reverseEdgeIndex = newEdges.findIndex(e => e.id === reverseEdgeId);
+                        if (reverseEdgeIndex !== -1) {
+                            newEdges[reverseEdgeIndex] = {
+                                ...newEdges[reverseEdgeIndex],
+                                data: {
+                                    ...newEdges[reverseEdgeIndex].data,
+                                    ...updates,
+                                },
+                            };
+                        }
+                    }
+                }
+
+                return newEdges;
+            });
+        },
+        [edgeSettings.id, setEdges],
+    );
+
+    // Close EdgeSettings
+    const closeEdgeSettings = useCallback(() => {
+        setEdgeSettings({ id: null });
+        setIsSettingsSidebarOpen(false);
+    }, []);
+
+    // Collapse all nodes on pane click
+    const handlePaneClick = useCallback(() => {
+        setSystemFlowNodeSettings({ id: null });
+
+        // Check if any root node is being edited and auto-center it
+        const rootNodesBeingEdited = nodes.filter(n =>
+            n.data.isTreeStructure &&
+            n.data.nodeType === 'root' &&
+            n.data.isRootBeingEdited
+        );
+
+        rootNodesBeingEdited.forEach(rootNode => {
+            const autoCenterRootNode = rootNode.data?.autoCenterRootNode;
+            if (typeof autoCenterRootNode === 'function') {
+                autoCenterRootNode();
+            }
+        });
+
+        // If there's a pending layout recalculation, trigger it now
+        if (pendingLayoutRecalculation) {
+
+            const rootNode = nodes.find(n => n.id === pendingLayoutRecalculation.rootNodeId);
+            if (rootNode) {
+                // Call the forceLayoutRecalculation function directly if available
+                const forceLayoutRecalculation = rootNode.data?.forceLayoutRecalculation;
+                if (typeof forceLayoutRecalculation === 'function') {
+                    forceLayoutRecalculation();
+                } else {
+                    // Fallback: Force a layout recalculation by updating the root node
+                    setNodes((nds: Node[]) =>
+                        nds.map((node: Node) => {
+                            if (node.id === rootNode.id) {
+                                return {
+                                    ...node,
+                                    data: {
+                                        ...node.data,
+                                        // Add a timestamp to force re-render and layout recalculation
+                                        lastLayoutUpdate: Date.now()
+                                    }
+                                };
+                            }
+                            return node;
+                        })
+                    );
+                }
+            }
+
+            // Clear the pending layout recalculation
+            setPendingLayoutRecalculation(null);
+        }
+    }, [pendingLayoutRecalculation, nodes, setNodes]);
+
+    // Memoize node and edge maps for O(1) lookup
+    const nodeMap = useMemo(() => {
+        const map = new Map<string, Node>();
+        nodes.forEach((node) => map.set(node.id, node));
+        return map;
+    }, [nodes]);
+    const edgeMap = useMemo(() => {
+        const map = new Map<string, Edge>();
+        edges.forEach((edge) => map.set(edge.id, edge));
+        return map;
+    }, [edges]);
+
+    // Memoize styles to prevent recreation on every render
+    const diagramStyles = useMemo(() => ({
+        background: '#1a1a1a',
+    }), []);
+
+    const fitViewOptions = useMemo(() => ({
+        padding: 0.2,
+        minZoom: 0.5,
+        maxZoom: 2,
+    }), []);
+
+    const defaultViewport = useMemo(() => ({
+        x: 0,
+        y: 0,
+        zoom: 1
+    }), []);
+
+    return (
+        <Box sx={{ display: 'flex', height: '100vh' }}>
+            <Sidebar>
+                <NodeSelectionContext.Provider
+                    value={{
+                        selectedNodeType: null,
+                        setSelectedNodeType: () => { },
+                        addNodeAtCenter,
+                    }}
+                >
+                    <DiagramControls
+                        onAddEntity={handleAddEntity}
+                        onUndo={handleUndo}
+                        onRedo={handleRedo}
+                        canUndo={historyIndex > 0}
+                        canRedo={historyIndex < history.length - 1}
+                        onLoadTemplate={loadTemplateData}
+                        edges={edges as SystemFlowEdge[]}
+                        nodes={nodes as SystemFlowNode[]}
+                        onEditEdge={(edge) => {
+                            setEdgeSettings({ id: edge.id });
+                            setIsSettingsSidebarOpen(true);
+                        }}
+                    />
+                </NodeSelectionContext.Provider>
+            </Sidebar>
+
+            <Box
+                sx={{
+                    flexGrow: 1,
+                    '& .react-flow__edge': {
+                        zIndex: '5000 !important',
+                    },
+                    '& .react-flow__edge-path': {
+                        pointerEvents: 'all !important',
+                    },
+                    '& .edge-path-selector': {
+                        strokeWidth: 20,
+                        fill: 'none',
+                        stroke: 'transparent',
+                        cursor: 'pointer',
+                    },
+                    '& .react-flow__edges': {
+                        zIndex: '5000 !important',
+                    },
+                    '& .react-flow__node': {
+                        pointerEvents: 'all !important',
+                        cursor: 'grab',
+                        '&:active': {
+                            cursor: 'grabbing',
+                        },
+                        '& .react-flow__edge': {
+                            zIndex: '5000 !important',
+                        },
+                    },
+                    '& .react-flow__node.dragging': {
+                        cursor: 'grabbing !important',
+                    },
+                    '& .react-flow.dragging': {
+                        cursor: 'grabbing',
+                    },
+                }}
+            >
+                <Box
+                    sx={{
+                        width: '100%',
+                        height: '100%',
+                        padding: '10px',
+                        bgcolor: '#1a1a1a',
+                    }}
+                >
+                    <ReactFlow
+                        nodes={nodes}
+                        edges={edges}
+                        onNodesChange={customOnNodesChange}
+                        onEdgesChange={onEdgesChange}
+                        onConnect={onConnect}
+                        nodeTypes={baseSystemFlowNodeTypes}
+                        edgeTypes={systemFlowEdgeTypes}
+                        fitView
+                        fitViewOptions={fitViewOptions}
+                        minZoom={0.1}
+                        maxZoom={2}
+                        defaultViewport={defaultViewport}
+                        style={diagramStyles}
+                        proOptions={{ hideAttribution: true }}
+                        selectNodesOnDrag={false}
+                        connectOnClick={false}
+                        onNodeClick={onNodeClick}
+                        onNodeDoubleClick={onNodeDoubleClick}
+                        onEdgeDoubleClick={onEdgeDoubleClick}
+                        onPaneClick={handlePaneClick}
+                    >
+                        <Background
+                            color="#1a1a1a"
+                            variant={BackgroundVariant.Dots}
+                            gap={16}
+                        />
+
+                        {/* MiniMap positioned at bottom-left */}
+                        <MiniMapCommon />
+
+                        {/* Diagram Panel with all controls */}
+                        <DiagramPanel
+                            historyIndex={historyIndex}
+                            historyLength={history.length}
+                            onUndo={handleUndo}
+                            onRedo={handleRedo}
+                            onClear={handleClear}
+                            onImport={(data) => {
+                                setNodes(data.nodes || []);
+                                setEdges(data.edges || []);
+                                // Add to history
+                                const newHistoryState = {
+                                    nodes: data.nodes || [],
+                                    edges: data.edges || [],
+                                };
+                                setHistory((prev) => [
+                                    ...prev.slice(0, historyIndex + 1),
+                                    newHistoryState,
+                                ]);
+                                setHistoryIndex((prev) => prev + 1);
+                            }}
+                            isSettingsOpen={isSettingsSidebarOpen}
+                            onToggleSettings={() => setIsSettingsSidebarOpen(true)}
+                            showCustomControls={true}
+                            onGridLayout={() => { }}
+                            onTreeLayout={() => { }}
+                        />
+
+                        {/* AI Button positioned at bottom-right */}
+                        <AiButton />
+
+                        {/* Fullscreen button positioned at top-left */}
+                    </ReactFlow>
+                </Box>
+            </Box>
+
+            {/* Settings panels */}
+            {systemFlowNodeSettings.id && (
+                <SystemFlowNodeSettings
+                    nodeId={systemFlowNodeSettings.id}
+                    data={nodeMap.get(systemFlowNodeSettings.id)?.data || {}}
+                    onUpdate={updateSelectedSystemFlowNode}
+                    onClose={closeSystemFlowNodeSettings}
+                    edges={edges as SystemFlowEdge[]}
+                    setEdges={(updater) => {
+                        if (typeof updater === 'function') {
+                            setEdges((prev) => updater(prev as SystemFlowEdge[]) as Edge[]);
+                        } else {
+                            setEdges(updater as Edge[]);
+                        }
+                    }}
+                    nodes={nodes as SystemFlowNode[]}
+                />
+            )}
+
+            {edgeSettings.id && (
+                <EdgeSettings
+                    edge={edgeMap.get(edgeSettings.id) as SystemFlowEdge}
+                    nodes={nodes as SystemFlowNode[]}
+                    onUpdate={updateSelectedEdge}
+                    onClose={closeEdgeSettings}
+                    onDelete={() => {
+                        setEdges((eds: Edge[]) => {
+                            const currentEdge = eds.find(e => e.id === edgeSettings.id);
+                            if (!currentEdge) return eds;
+
+                            // If it's a two-way edge, also remove the reverse edge
+                            const reverseEdgeId = currentEdge.data?.reverseEdgeId;
+                            if (reverseEdgeId) {
+                                return eds.filter((edge: Edge) =>
+                                    edge.id !== edgeSettings.id && edge.id !== reverseEdgeId
+                                );
+                            }
+
+                            // Regular edge deletion
+                            return eds.filter((edge: Edge) => edge.id !== edgeSettings.id);
+                        });
+                        closeEdgeSettings();
+                    }}
+                />
+            )}
+        </Box>
+    );
+});
+
+SystemFlowDiagramFlowContent.displayName = 'SystemFlowDiagramFlowContent';
+
+// Main wrapper component with ReactFlowProvider
+const SystemFlowDiagramFlow: React.FC<SystemFlowDiagramFlowProps> = (
+    props,
+) => {
+    return (
+        <ReactFlowProvider>
+            <SystemFlowDiagramFlowContent {...props} />
+        </ReactFlowProvider>
+    );
+};
+
+export default SystemFlowDiagramFlow;