@flowuent-org/diagramming-core 1.1.0 → 1.1.2

package/packages/diagrams/src/lib/templates/stateMachine/components/StateMachineDiagramFlow.tsx

@@ -0,0 +1,1238 @@
+import React, {
+  useCallback,
+  useState,
+  useRef,
+  useEffect,
+  useMemo,
+} from 'react';
+
+// Fullscreen API type extensions
+interface FullscreenElement extends HTMLElement {
+  webkitRequestFullscreen?: () => Promise<void>;
+  msRequestFullscreen?: () => Promise<void>;
+}
+
+interface FullscreenDocument extends Document {
+  webkitExitFullscreen?: () => Promise<void>;
+  msExitFullscreen?: () => Promise<void>;
+}
+import {
+  ReactFlow,
+  Background,
+  useNodesState,
+  useEdgesState,
+  Panel,
+  BackgroundVariant,
+  Connection,
+  NodeChange,
+  applyNodeChanges,
+  ReactFlowProvider,
+  Node,
+  Edge,
+  NodeProps,
+} from '@xyflow/react';
+import '@xyflow/react/dist/style.css';
+import {
+  StateMachineNode,
+  StateMachineEdge,
+  HistoryState,
+  EditingState,
+  StateNodeSettingsState,
+  EdgeSettingsState,
+  NodeSelectionContextType,
+  LayoutDirection,
+  Position,
+  UpdateFunction,
+  CloseFunction,
+  DeleteFunction
+} from '../../../types/state-machine-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 {
+  stateMachineDefaultNodes,
+  stateMachineDefaultEdges,
+} from '../data/template-data';
+import {
+  stateMachineEdgeTypes,
+  stateMachineNodeTypes as baseStateMachineNodeTypes,
+  CompositeStateNode as CompositeStateNodeType,
+} from './StateMachineComponents';
+import DiagramControls from '../controls/StateDiagramControls';
+import { NodeSelectionContext } from '../controls/StateDiagramControls';
+import StateNodeSettings from './StateNodeSettings';
+import EdgeSettings from './EdgeSettings';
+import CompositeStateNodeSettings from './CompositeStateNodeSettings';
+import { useStateMachineLayout } from './StateMachineLayout';
+import DiagramPanel from '../../../components/DiagramPanel';
+import Sidebar from '../../../components/SideBar';
+import AiButton from '../../../components/AiButton';
+import MiniMapCommon from '../../../components/MiniMapCommon';
+
+interface StateMachineDiagramFlowProps {
+  nodes: StateMachineNode[];
+  setNodes: (nodes: StateMachineNode[] | ((prev: StateMachineNode[]) => StateMachineNode[])) => void;
+  edges: StateMachineEdge[];
+  setEdges: (edges: StateMachineEdge[] | ((prev: StateMachineEdge[]) => StateMachineEdge[])) => void;
+}
+
+interface CompositeStateNodeProps extends NodeProps {
+  onToggleExpand?: (id: string, expanded: boolean) => void;
+  onSettingsClick?: () => 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 to recursively collect all child node ids of collapsed composite nodes
+function getHiddenNodeIds(allNodes: Node[]): Set<string> {
+  const hidden = new Set<string>();
+  const collapsedComposites = allNodes.filter(
+    (n) => n.type === 'CompositeStateNode' && !n.data.expanded,
+  );
+  const collectChildren = (parentId: string) => {
+    allNodes.forEach((child) => {
+      if (child.parentId === parentId) {
+        hidden.add(child.id);
+        if (child.type === 'CompositeStateNode') {
+          collectChildren(child.id);
+        }
+      }
+    });
+  };
+  collapsedComposites.forEach((comp) => collectChildren(comp.id));
+  return hidden;
+}
+
+// Helper: get center position in diagram
+const getCenterPosition = (width: number, height: number) => {
+  return { x: width / 2 - 75, y: height / 2 - 40 };
+};
+
+const StateMachineDiagramFlowContent: React.FC<
+  StateMachineDiagramFlowProps
+> = ({
+  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[]);
+    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 [stateNodeSettings, setStateNodeSettings] = useState<{
+      id: string | null;
+    }>({ id: null });
+    const [edgeSettings, setEdgeSettings] = useState<{ id: string | null }>({
+      id: null,
+    });
+    const [selectedCompositeStateNodeId, setSelectedCompositeStateNodeId] =
+      useState<string | null>(null);
+    const [isSettingsSidebarOpen, setIsSettingsSidebarOpen] = useState(false);
+
+    // Add layout integration
+    const { getLayoutedElements, layoutDirection, setLayoutDirection } =
+      useStateMachineLayout(nodes as Node[], edges as StateMachineEdge[], setNodes);
+
+    // Function declarations in correct order
+    const handleToggleCompositeExpand = useCallback(
+      (id: string, expanded: boolean) => {
+        setNodes((nds: Node[]) =>
+          nds.map((node: Node) =>
+            node.id === id && node.type === 'CompositeStateNode'
+              ? { ...node, data: { ...node.data, expanded } }
+              : node,
+          ),
+        );
+      },
+      [setNodes],
+    );
+
+    const nodeTypes = useMemo(() => {
+      return {
+        ...baseStateMachineNodeTypes,
+        CompositeStateNode: (props: NodeProps) => {
+          const compositeProps: CompositeStateNodeProps = {
+            ...props,
+            onToggleExpand: handleToggleCompositeExpand,
+            onSettingsClick: undefined,
+          };
+          return <CompositeStateNodeType {...compositeProps} />;
+        },
+      };
+    }, [handleToggleCompositeExpand]);
+
+    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) => {
+        console.log(`Loading template: ${templateType}`);
+        // For now, just load the default state machine nodes/edges
+        setNodes(stateMachineDefaultNodes as Node[]);
+        setEdges(stateMachineDefaultEdges as Edge[]);
+      },
+      [setNodes, setEdges],
+    );
+
+    // Initialize nodes and edges with our data
+    useEffect(() => {
+      // Only initialize if history is empty
+      if (history.length === 0) {
+        const initialState = {
+          nodes: [...nodes],
+          edges: [...edges],
+        };
+        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 = {
+            nodes: [...nodes],
+            edges: [...edges],
+          };
+          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 state machine nodes
+    const nodeTypeMap: Record<string, { type: string; data?: Record<string, unknown> }> = {
+      StateNode: { type: 'StateNode', data: { label: 'State' } },
+      InitialStateNode: { type: 'InitialStateNode', data: {} },
+      FinalStateNode: { type: 'FinalStateNode', data: {} },
+      CompositeStateNode: {
+        type: 'CompositeStateNode',
+        data: { label: 'Composite State', expanded: false },
+      },
+      HistoryStateNode: { type: 'HistoryStateNode', data: {} },
+      JunctionPointNode: { type: 'JunctionPointNode', data: {} },
+      ChoicePointNode: { type: 'ChoicePointNode', data: {} },
+      EntryPointNode: { type: 'EntryPointNode', data: {} },
+      ExitPointNode: { type: 'ExitPointNode', data: {} },
+      ForkJoinNode: { type: 'ForkJoinNode', data: { orientation: 'horizontal' } },
+      BendpointNode: { type: 'bendpoint', data: {} },
+    };
+
+    // 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,
+            ),
+          );
+          // Update all selfconnecting edges for this node with new size/position
+          setEdges((eds) =>
+            eds.map((edge) => {
+              if (
+                edge.type === 'selfconnecting' &&
+                edge.source === id &&
+                edge.target === id
+              ) {
+                // Find the node to get its position
+                const node = nodes.find((n) => n.id === id);
+                const nodeX = node?.position?.x ?? 0;
+                const nodeY = node?.position?.y ?? 0;
+                return {
+                  ...edge,
+                  data: {
+                    ...edge.data,
+                    nodeWidth: width,
+                    nodeHeight: height,
+                    nodeX,
+                    nodeY,
+                  },
+                };
+              }
+              return edge;
+            }),
+          );
+        };
+
+        setNodes((nds: Node[]) => [
+          ...nds,
+          {
+            id: newId,
+            type,
+            position: center,
+            data: { ...data, onResize },
+            draggable: true,
+          },
+        ]);
+      },
+      [nodes, setNodes, setEdges],
+    );
+
+    // Update onAddEntity to use internal addNodeAtCenter
+    const handleAddEntity = useCallback(
+      (nodeTypeId = 'StateNode') => {
+        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 = 'state-machine-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) {
+            console.error('Error importing diagram:', 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 FullscreenElement).webkitRequestFullscreen) {
+          (diagramContainerRef.current as FullscreenElement).webkitRequestFullscreen?.();
+        } else if ((diagramContainerRef.current as FullscreenElement).msRequestFullscreen) {
+          (diagramContainerRef.current as FullscreenElement).msRequestFullscreen?.();
+        }
+      } else {
+        if (document.exitFullscreen) {
+          document.exitFullscreen();
+        } else if ((document as FullscreenDocument).webkitExitFullscreen) {
+          (document as FullscreenDocument).webkitExitFullscreen?.();
+        } else if ((document as FullscreenDocument).msExitFullscreen) {
+          (document as FullscreenDocument).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: 'customEdge',
+          sourceHandle: params.sourceHandle,
+          targetHandle: params.targetHandle,
+          data: {
+            label: `Transition ${edges.length + 1}`,
+          },
+          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) => {
+      if (node.type === 'bendpoint') {
+        setStateNodeSettings({ id: null });
+        setSelectedCompositeStateNodeId(null);
+        return;
+      }
+
+      if (node.type === 'CompositeStateNode') {
+        setSelectedCompositeStateNodeId(node.id);
+        setStateNodeSettings({ id: null });
+        setIsSettingsSidebarOpen(true);
+      } else if (
+        [
+          'StateNode',
+          'InitialStateNode',
+          'FinalStateNode',
+          'HistoryStateNode',
+          'JunctionPointNode',
+          'ChoicePointNode',
+          'EntryPointNode',
+          'ExitPointNode',
+          'ForkJoinNode',
+        ].includes(node.type || '')
+      ) {
+        setStateNodeSettings({ id: node.id });
+        setSelectedCompositeStateNodeId(null);
+        setIsSettingsSidebarOpen(true);
+      } else {
+        setStateNodeSettings({ id: null });
+        setSelectedCompositeStateNodeId(null);
+      }
+    }, []);
+
+    // 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 === 'customEdge') {
+          setEdgeSettings({ id: edge.id });
+          setIsSettingsSidebarOpen(true);
+        }
+      },
+      [],
+    );
+
+    // Function to update selected state node
+    const updateSelectedStateNode = useCallback(
+      (updates: Record<string, unknown>) => {
+        if (!stateNodeSettings.id) return;
+        setNodes((nds: Node[]) =>
+          nds.map((node: Node) => {
+            if (node.id === stateNodeSettings.id) {
+              return {
+                ...node,
+                data: {
+                  ...node.data,
+                  ...updates,
+                },
+              };
+            }
+            return node;
+          }),
+        );
+      },
+      [stateNodeSettings.id, setNodes],
+    );
+
+    // Close state node settings
+    const closeStateNodeSettings = useCallback(() => {
+      setStateNodeSettings({ id: null });
+      setIsSettingsSidebarOpen(false);
+    }, []);
+
+    // Update edge from EdgeSettings
+    const updateSelectedEdge = useCallback(
+      (updates: Record<string, unknown>) => {
+        if (!edgeSettings.id) return;
+        setEdges((eds: Edge[]) =>
+          eds.map((edge: Edge) => {
+            if (edge.id === edgeSettings.id) {
+              return {
+                ...edge,
+                source: typeof updates.source === 'string' ? updates.source : edge.source,
+                target: typeof updates.target === 'string' ? updates.target : edge.target,
+                data: {
+                  ...edge.data,
+                  ...updates,
+                },
+              };
+            }
+            return edge;
+          }) as Edge[],
+        );
+      },
+      [edgeSettings.id, setEdges],
+    );
+
+    // Close EdgeSettings
+    const closeEdgeSettings = useCallback(() => {
+      setEdgeSettings({ id: null });
+      setIsSettingsSidebarOpen(false);
+    }, []);
+
+    // Close CompositeStateNodeSettings
+    const closeCompositeStateNodeSettings = useCallback(() => {
+      setSelectedCompositeStateNodeId(null);
+      setIsSettingsSidebarOpen(false);
+    }, []);
+
+    // Function to update selected composite state node
+    const updateSelectedCompositeStateNode = useCallback(
+      (updates: Record<string, unknown>) => {
+        if (!selectedCompositeStateNodeId) return;
+        setNodes((nds: Node[]) =>
+          nds.map((node: Node) => {
+            if (node.id === selectedCompositeStateNodeId) {
+              return {
+                ...node,
+                data: {
+                  ...node.data,
+                  ...updates,
+                },
+              };
+            }
+            return node;
+          }),
+        );
+      },
+      [selectedCompositeStateNodeId, setNodes],
+    );
+
+    // Collapse all composite nodes on pane click
+    const handlePaneClick = useCallback(() => {
+      setNodes((nds: Node[]) =>
+        nds.map((node: Node) =>
+          node.type === 'CompositeStateNode'
+            ? { ...node, data: { ...node.data, expanded: false } }
+            : node,
+        ),
+      );
+      setSelectedCompositeStateNodeId(null);
+      setStateNodeSettings({ id: null });
+    }, [setNodes]);
+
+    // Memoize hidden node ids
+    const hiddenNodeIds = useMemo(() => getHiddenNodeIds(nodes), [nodes]);
+    // Memoize filtered nodes
+    const filteredNodes = useMemo(
+      () => nodes.filter((n) => !hiddenNodeIds.has(n.id)),
+      [nodes, hiddenNodeIds],
+    );
+    // Memoize collapsed composites
+    const collapsedComposites = useMemo(
+      () =>
+        nodes.filter((n) => n.type === 'CompositeStateNode' && !n.data.expanded),
+      [nodes],
+    );
+    // Memoize proxy edge logic (optimized for performance)
+    const { filteredEdges } = useMemo(() => {
+      // Precompute collapsed composite node ids for O(1) lookup
+      const collapsedCompositeIds = new Set(
+        nodes
+          .filter((n) => n.type === 'CompositeStateNode' && !n.data.expanded)
+          .map((n) => n.id),
+      );
+      // Precompute a map of nodeId -> node for O(1) lookup
+      const nodeIdMap = new Map(nodes.map((n) => [n.id, n]));
+      // Map childId to its collapsed composite parent
+      const childToCollapsedParent: Record<string, string> = {};
+      const collectChildren = (parentId: string) => {
+        nodes.forEach((child) => {
+          if (child.parentId === parentId) {
+            childToCollapsedParent[child.id] = parentId;
+            if (child.type === 'CompositeStateNode') {
+              collectChildren(child.id);
+            }
+          }
+        });
+      };
+      collapsedCompositeIds.forEach((compId) => collectChildren(compId));
+      // Generate proxy edges for edges crossing collapsed composite boundaries
+      const proxyEdges: Edge[] = [];
+      const realEdges: Edge[] = [];
+      edges.forEach((edge) => {
+        const sourceHidden = hiddenNodeIds.has(edge.source);
+        const targetHidden = hiddenNodeIds.has(edge.target);
+
+        // Self-loop on composite node: hide if collapsed
+        if (
+          edge.source === edge.target &&
+          collapsedCompositeIds.has(edge.source)
+        ) {
+          return; // skip this edge
+        }
+
+        // If both ends are hidden, skip (don't show)
+        if (sourceHidden && targetHidden) return;
+        // If neither end is hidden, show as normal
+        if (!sourceHidden && !targetHidden) {
+          realEdges.push(edge);
+          return;
+        }
+        // If one end is hidden and the other is visible, create a proxy edge
+        let proxySource = edge.source;
+        let proxyTarget = edge.target;
+        let proxyId = edge.id + '-proxy';
+        // If source is hidden, map to its collapsed composite parent
+        if (sourceHidden && childToCollapsedParent[edge.source]) {
+          proxySource = childToCollapsedParent[edge.source];
+        }
+        // If target is hidden, map to its collapsed composite parent
+        if (targetHidden && childToCollapsedParent[edge.target]) {
+          proxyTarget = childToCollapsedParent[edge.target];
+        }
+        // Avoid proxying to the same node (self-loop)
+        if (proxySource === proxyTarget) return;
+        proxyEdges.push({
+          ...edge,
+          id: proxyId,
+          source: proxySource,
+          target: proxyTarget,
+          style: { ...(edge.style || {}), strokeDasharray: '6 4', opacity: 0.7 },
+          data: { ...(edge.data || {}), isProxy: true },
+        });
+      });
+      return { filteredEdges: [...realEdges, ...proxyEdges] as Edge[] };
+    }, [edges, nodes, hiddenNodeIds]);
+
+    // Custom onNodesChange to snap PinPointNode to border of parent composite node
+    const customOnNodesChange = useCallback(
+      (changes: NodeChange[]) => {
+        setNodes((nds) => {
+          // Apply default node changes
+          let updated = applyNodeChanges(changes, nds);
+          // Snap PinPointNode nodes to border of parent composite node
+          updated = updated.map((node) => {
+            if (node.type === 'PinPointNode' && node.parentId) {
+              const parent = nds.find((n) => n.id === node.parentId);
+              if (!parent) return node;
+              const width = (parent.data as Record<string, unknown>).width as number || 400;
+              const height = (parent.data as Record<string, unknown>).height as number || 300;
+              const size = 22;
+              let { x, y } = node.position;
+              // Calculate center
+              const cx = x + size / 2;
+              const cy = y + size / 2;
+              // Find closest border (top, bottom, left, right)
+              const dTop = Math.abs(cy - 0);
+              const dBottom = Math.abs(cy - height);
+              const dLeft = Math.abs(cx - 0);
+              const dRight = Math.abs(cx - width);
+              const minDist = Math.min(dTop, dBottom, dLeft, dRight);
+              let newX = x,
+                newY = y;
+              if (minDist === dTop) {
+                newX =
+                  Math.max(-size / 2, Math.min(width - size / 2, cx)) - size / 2;
+                newY = -size / 2;
+              } else if (minDist === dBottom) {
+                newX =
+                  Math.max(-size / 2, Math.min(width - size / 2, cx)) - size / 2;
+                newY = height - size / 2;
+              } else if (minDist === dLeft) {
+                newX = -size / 2;
+                newY =
+                  Math.max(-size / 2, Math.min(height - size / 2, cy)) - size / 2;
+              } else if (minDist === dRight) {
+                newX = width - size / 2;
+                newY =
+                  Math.max(-size / 2, Math.min(height - size / 2, cy)) - size / 2;
+              }
+              return { ...node, position: { x: newX, y: newY } };
+            }
+            return node;
+          });
+          return updated;
+        });
+      },
+      [setNodes],
+    );
+
+    // Add layout direction change handler
+    const handleLayoutDirectionChange = useCallback(
+      (direction: 'DOWN' | 'RIGHT') => {
+        setLayoutDirection(direction);
+        getLayoutedElements();
+      },
+      [setLayoutDirection, getLayoutedElements],
+    );
+
+    // Add layout effect
+    useEffect(() => {
+      // Only layout if we have nodes/edges
+      if (nodes.length > 0 || edges.length > 0) {
+        getLayoutedElements();
+      }
+    }, [handleToggleCompositeExpand]);
+
+    // Add edge edit handlers
+    const handleEditEdge = useCallback((edge: Edge) => {
+      setEdgeSettings({ id: edge.id });
+    }, []);
+
+    const handleUpdateEdgeLabel = useCallback(
+      (edgeId: string, newLabel: string) => {
+        setEdges((eds: Edge[]) =>
+          eds.map((edge: Edge) =>
+            edge.id === edgeId
+              ? { ...edge, data: { ...edge.data, label: newLabel } }
+              : edge,
+          ),
+        );
+      },
+      [setEdges],
+    );
+
+    const handleToggleEdgeArrowHead = useCallback(
+      (edgeId: string, direction: 'left' | 'right') => {
+        setEdges((eds: Edge[]) =>
+          eds.map((edge: Edge) => {
+            if (edge.id === edgeId) {
+              const data = { ...edge.data };
+              if (direction === 'left') {
+                data.arrowHeadLeft = !data.arrowHeadLeft;
+              } else {
+                data.arrowHeadRight = !data.arrowHeadRight;
+              }
+              return { ...edge, data };
+            }
+            return edge;
+          }),
+        );
+      },
+      [setEdges],
+    );
+
+    // 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]);
+
+    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 StateMachineEdge[]}
+              nodes={nodes as StateMachineNode[]}
+              onEditEdge={handleEditEdge}
+              onUpdateEdgeLabel={handleUpdateEdgeLabel}
+              onToggleEdgeArrowHead={handleToggleEdgeArrowHead}
+            />
+          </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: 'white',
+            }}
+          >
+            <ReactFlow
+              nodes={filteredNodes}
+              edges={filteredEdges}
+              onNodesChange={customOnNodesChange}
+              onEdgesChange={onEdgesChange}
+              onConnect={onConnect}
+              nodeTypes={nodeTypes}
+              edgeTypes={stateMachineEdgeTypes}
+              fitView
+              fitViewOptions={{
+                padding: 0.2,
+                minZoom: 0.5,
+                maxZoom: 2,
+              }}
+              minZoom={0.1}
+              maxZoom={2}
+              defaultViewport={{ x: 0, y: 0, zoom: 1 }}
+              style={{
+                background: '#f8f9fa',
+              }}
+              proOptions={{ hideAttribution: true }}
+              selectNodesOnDrag={false}
+              connectOnClick={false}
+              onNodeClick={onNodeClick}
+              onNodeDoubleClick={onNodeDoubleClick}
+              onEdgeDoubleClick={onEdgeDoubleClick}
+              onPaneClick={handlePaneClick}
+            >
+              <Background
+                color="#f8f9fa"
+                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={() => handleLayoutDirectionChange('DOWN')}
+                onTreeLayout={() => handleLayoutDirectionChange('RIGHT')}
+              />
+
+              {/* AI Button positioned at bottom-right */}
+              <AiButton />
+
+              {/* Layout controls positioned at top-right */}
+
+              {/* Fullscreen button positioned at top-left */}
+            </ReactFlow>
+          </Box>
+        </Box>
+
+        {/* Settings panels */}
+        {stateNodeSettings.id && (
+          <StateNodeSettings
+            nodeId={stateNodeSettings.id}
+            data={nodeMap.get(stateNodeSettings.id)?.data || {}}
+            onUpdate={updateSelectedStateNode}
+            onClose={closeStateNodeSettings}
+            edges={edges as StateMachineEdge[]}
+            setEdges={(updater) => {
+              if (typeof updater === 'function') {
+                setEdges((prev) => updater(prev as StateMachineEdge[]) as Edge[]);
+              } else {
+                setEdges(updater as Edge[]);
+              }
+            }}
+            nodes={nodes as StateMachineNode[]}
+          />
+        )}
+
+        {edgeSettings.id && (
+          <EdgeSettings
+            data={{
+              ...(edgeMap.get(edgeSettings.id)?.data || {}),
+              source: edgeMap.get(edgeSettings.id)?.source,
+              target: edgeMap.get(edgeSettings.id)?.target,
+            }}
+            nodes={nodes as StateMachineNode[]}
+            onUpdate={updateSelectedEdge}
+            onClose={closeEdgeSettings}
+            onDelete={() => {
+              setEdges((eds: Edge[]) =>
+                eds.filter((edge: Edge) => edge.id !== edgeSettings.id),
+              );
+              closeEdgeSettings();
+            }}
+          />
+        )}
+
+        {selectedCompositeStateNodeId && (
+          <CompositeStateNodeSettings
+            nodeId={selectedCompositeStateNodeId}
+            data={nodeMap.get(selectedCompositeStateNodeId)?.data || {}}
+            nodes={nodes as StateMachineNode[]}
+            setNodes={(updater) => {
+              if (typeof updater === 'function') {
+                setNodes((prev) => updater(prev as StateMachineNode[]) as Node[]);
+              } else {
+                setNodes(updater as Node[]);
+              }
+            }}
+            onUpdate={updateSelectedCompositeStateNode}
+            onClose={closeCompositeStateNodeSettings}
+            edges={edges as StateMachineEdge[]}
+            setEdges={(updater) => {
+              if (typeof updater === 'function') {
+                setEdges((prev) => updater(prev as StateMachineEdge[]) as Edge[]);
+              } else {
+                setEdges(updater as Edge[]);
+              }
+            }}
+          />
+        )}
+      </Box>
+    );
+  };
+
+// Main wrapper component with ReactFlowProvider
+const StateMachineDiagramFlow: React.FC<StateMachineDiagramFlowProps> = (
+  props,
+) => {
+  return (
+    <ReactFlowProvider>
+      <StateMachineDiagramFlowContent {...props} />
+    </ReactFlowProvider>
+  );
+};
+
+export default StateMachineDiagramFlow;