@flowuent-org/diagramming-core 1.3.10 → 1.4.1

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

@@ -1,1238 +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;
+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;