@flowuent-org/diagramming-core 1.3.10 → 1.4.1

package/packages/diagrams/src/lib/templates/sequence/components/SequenceDiagramFlow.tsx

@@ -1,1361 +1,1361 @@
-import React, { useCallback, useState, useRef, useEffect } from 'react';
-import {
-  ReactFlow,
-  Background,
-  useNodesState,
-  useEdgesState,
-  Panel,
-  BackgroundVariant,
-  Connection,
-  NodeChange,
-  Node,
-  Edge,
-  MiniMap,
-  NodeMouseHandler,
-  EdgeMouseHandler,
-} from '@xyflow/react';
-import '@xyflow/react/dist/style.css';
-
-import {
-  Box,
-  Button,
-  TextField,
-  Typography,
-  IconButton,
-  Paper,
-} 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 { useTranslation } from 'react-i18next';
-import {
-  SequenceDiagramProps,
-  ThemeData,
-  SequenceNode,
-  SequenceEdge,
-  SequenceHistoryEntry,
-  SequenceEntityShape,
-  CombinedFragmentType,
-  SequenceMessageType,
-  SequenceEntityNode,
-  SequenceEntityData,
-  FragmentSettingsUpdateData,
-  EntitySettingsUpdateData,
-  ActivationBarNode,
-  CombinedFragmentNode,
-  SequenceMessageData,
-  ActivationBarData,
-  CombinedFragmentData
-} from '../../../types/sequence-types';
-import { getTemplateData } from '../data/template-data';
-import { sequenceEdgeTypes, sequenceNodeTypes } from './SequenceComponents';
-import DiagramControls from '../controls/DiagramControls';
-import FragmentSettings from './FragmentSettings';
-import EntitySettings from './EntitySettings';
-import AddMessageDialog from './AddMessageDialog';
-import Sidebar from '../../../components/SideBar';
-import AiButton from '../../../components/AiButton';
-import MiniMapCommon from '../../../components/MiniMapCommon';
-import DiagramPanel from '../../../components/DiagramPanel';
-
-// Types for browser fullscreen API extensions
-interface HTMLElementWithFullscreen extends HTMLElement {
-  webkitRequestFullscreen?: () => Promise<void>;
-  msRequestFullscreen?: () => Promise<void>;
-}
-
-interface DocumentWithFullscreen extends Document {
-  webkitExitFullscreen?: () => Promise<void>;
-  msExitFullscreen?: () => Promise<void>;
-}
-
-// Main component that manages the diagram state
-
-const SequenceDiagramFlow: React.FC<SequenceDiagramProps> = ({
-  sequenceNodes,
-  sequenceEdges,
-}) => {
-  const { t } = useTranslation();
-
-  // Initialize nodes and edges with our data - use xy/flow base types with proper generics
-  const [nodes, setNodes, onNodesChange] = useNodesState(sequenceNodes);
-
-  const [edges, setEdges, onEdgesChange] = useEdgesState(sequenceEdges);
-
-  // History state for undo/redo functionality
-  const [history, setHistory] = useState<SequenceHistoryEntry[]>([]);
-  const [historyIndex, setHistoryIndex] = useState<number>(-1);
-  const isUndoRedoAction = useRef(false);
-
-  // Initialize history with initial state
-  useEffect(() => {
-    // Only initialize if history is empty
-    if (history.length === 0) {
-      const initialState: SequenceHistoryEntry = {
-        nodes: [...nodes] as SequenceNode[],
-        edges: [...edges] as SequenceEdge[],
-      };
-      setHistory([initialState]);
-      setHistoryIndex(0);
-    }
-  }, []);
-
-  // Save current state to history when nodes or edges change
-  useEffect(() => {
-    // Don't record history when performing undo/redo actions or when we're still initializing
-    if (isUndoRedoAction.current || history.length === 0) {
-      return;
-    }
-
-    // Debounce: Only update history if nodes/edges actually changed (deep compare)
-    const newHistoryEntry: SequenceHistoryEntry = {
-      nodes: [...nodes] as SequenceNode[],
-      edges: [...edges] as SequenceEdge[],
-    };
-    const currentState = history[historyIndex];
-    if (
-      JSON.stringify(currentState.nodes) ===
-      JSON.stringify(newHistoryEntry.nodes) &&
-      JSON.stringify(currentState.edges) ===
-      JSON.stringify(newHistoryEntry.edges)
-    ) {
-      return; // No change, don't add to history
-    }
-
-    // Trim future history if we're in the middle of the history stack
-    const newHistory = history.slice(0, historyIndex + 1);
-    newHistory.push(newHistoryEntry);
-    setHistory(newHistory);
-    setHistoryIndex(newHistory.length - 1);
-  }, [nodes, edges, history, historyIndex]);
-
-  // 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]);
-
-  // 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]);
-
-  const [theme, setTheme] = useState<ThemeData>({
-    primary: '#61dafb',
-    secondary: '#2a2a2a',
-  });
-  const [isFullscreen, setIsFullscreen] = useState<boolean>(false);
-  const [selectedFragmentId, setSelectedFragmentId] = useState<string | null>(
-    null,
-  );
-  const [fragmentSettings, setFragmentSettings] = useState({
-    width: 400,
-    height: 150,
-  });
-  const [showFragmentSettings, setShowFragmentSettings] =
-    useState<boolean>(false);
-  const diagramContainerRef = useRef<HTMLDivElement>(null);
-  const [editingNode, setEditingNode] = useState<string | null>(null);
-  const [editingEdge, setEditingEdge] = useState<string | null>(null);
-  const [editingLabel, setEditingLabel] = useState<string>('');
-  const [editingConditions, setEditingConditions] = useState<string[]>([]);
-  const [editingType, setEditingType] = useState<'node' | 'edge' | 'fragment'>(
-    'node',
-  );
-  const [selectedEntityId, setSelectedEntityId] = useState<string | null>(null);
-  const [showEntitySettings, setShowEntitySettings] = useState<boolean>(false);
-  const [showMessageDialog, setShowMessageDialog] = useState<boolean>(false);
-  const ENTITY_FIXED_Y = 50; // Fixed y position for all entities
-
-  // Type-safe node access helpers
-  const getSequenceEntityNodes = useCallback(() => {
-    return nodes.filter((node): node is Node<SequenceEntityData> =>
-      node.type === 'sequenceEntity'
-    );
-  }, [nodes]);
-
-  const getCombinedFragmentNodes = useCallback(() => {
-    return nodes.filter((node): node is Node<CombinedFragmentData> =>
-      node.type === 'combinedFragment'
-    );
-  }, [nodes]);
-
-  // Intercept node position changes to lock y for sequenceEntity nodes
-  const onNodesChangeLockedY = useCallback(
-    (changes: NodeChange[]) => {
-      let hasPositionChange = false;
-      const updatedNodes = nodes.map((node) => {
-        // Find if there's a position change for this node
-        const change = changes.find(
-          (c: NodeChange) => c.type === 'position' && c.id === node.id,
-        );
-        if (
-          change &&
-          change.type === 'position' &&
-          node.type === 'sequenceEntity' &&
-          change.position
-        ) {
-          hasPositionChange = true;
-          // Lock y position
-          return {
-            ...node,
-            position: {
-              x: change.position.x,
-              y: ENTITY_FIXED_Y, // always fixed
-            },
-          };
-        }
-        return node;
-      });
-      if (hasPositionChange) {
-        setNodes(updatedNodes);
-      } else {
-        onNodesChange(changes);
-      }
-    },
-    [nodes, setNodes, onNodesChange],
-  );
-
-  // Handler to add a new entity
-  const handleAddEntity = useCallback(
-    (shape: SequenceEntityShape = 'rectangle') => {
-      console.log('Adding new entity with shape:', shape);
-      console.log('Current nodes:', nodes);
-
-      // Find all entity nodes
-      const entityNodes = getSequenceEntityNodes();
-      const entityCount = entityNodes.length;
-      const newEntityId = `entity-${entityCount + 1}`;
-
-      // Find the rightmost x position among entities
-      const rightmostX =
-        entityNodes.length > 0
-          ? Math.max(...entityNodes.map((node) => node.position.x))
-          : 100; // Default starting x
-
-      const xOffset = 150; // Space between entities
-      const xPosition =
-        entityNodes.length > 0 ? rightmostX + xOffset : rightmostX;
-
-      // Create a new entity node with extended properties
-      const newEntity: Node<SequenceEntityData> = {
-        id: newEntityId,
-        type: 'sequenceEntity',
-        position: { x: xPosition, y: 50 },
-        data: {
-          label: `Entity ${entityCount + 1}`,
-          color: theme.primary,
-          backgroundColor: theme.secondary,
-          lifelineLength: 500, // Use same height as others
-          shape: shape, // Use the provided shape
-          fontSize: 12,
-          borderWidth: 2,
-          description: '',
-          collapsed: false,
-        },
-        draggable: true,
-        zIndex: 5, // Higher z-index to appear above other elements
-      };
-
-      // Adjust position based on entity shape
-      let offsetX = 55;
-      switch (shape) {
-        case 'circle':
-          offsetX = 40;
-          break;
-        case 'database':
-          offsetX = 45;
-          break;
-        case 'actor':
-          offsetX = 60;
-          break;
-        default:
-          offsetX = 55;
-          break;
-      }
-
-      const activationBar: Node<ActivationBarData> = {
-        id: `${newEntityId}-activation`,
-        type: 'activationBar',
-        position: { x: xPosition + offsetX, y: 180 }, // Position below the entity
-        data: {
-          entityId: newEntityId,
-          color: theme.secondary,
-          borderColor: theme.primary,
-          length: 150,
-        },
-        draggable: true,
-        zIndex: 1, // Lower z-index to appear behind entity boxes but above combined fragments
-      };
-
-      setNodes((nds) => [...nds, newEntity, activationBar]);
-    },
-    [nodes, setNodes, theme, getSequenceEntityNodes],
-  );
-
-  // Filter entities for Add Message Dialog
-  const sequenceEntities = useCallback((): SequenceEntityNode[] => {
-    return getSequenceEntityNodes().map((node) => ({
-      id: node.id,
-      type: 'sequenceEntity',
-      data: node.data,
-      position: node.position,
-      draggable: node.draggable,
-      zIndex: node.zIndex,
-    }));
-  }, [getSequenceEntityNodes]);
-
-  // Handler to open add message dialog
-  const handleOpenAddMessageDialog = useCallback(() => {
-    // Filter to get only entity nodes
-    const entityNodes = getSequenceEntityNodes();
-
-    if (entityNodes.length < 1) {
-      alert(t('sequenceDiagram.errors.needEntitiesForMessage'));
-      return;
-    }
-
-    setShowMessageDialog(true);
-  }, [getSequenceEntityNodes]);
-
-  // Handler to create a new message with specified source, target, and type
-  const handleAddCustomMessage = useCallback(
-    (
-      sourceId: string,
-      targetId: string,
-      messageType: SequenceMessageType,
-      label: string,
-    ) => {
-      // Find the source and target nodes
-      const sourceNode = nodes.find((n) => n.id === sourceId);
-      const targetNode = nodes.find((n) => n.id === targetId);
-
-      if (!sourceNode || !targetNode) {
-        console.error(t('sequenceDiagram.errors.sourceTargetNotFound'));
-        return;
-      }
-
-      // Calculate a reasonable Y position for the new message
-      const messageY =
-        Math.max(
-          ...edges.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
-          180, // Default minimum Y position
-        ) + 60; // Add some spacing
-
-      // Calculate handle positions based on nodes
-      const sourceX = sourceNode.position?.x || 0;
-      const targetX = targetNode.position?.x || 0;
-
-      // Calculate source and target handles
-      let sourceHandle = 'right';
-      let targetHandle = 'left';
-
-      if (messageType === 'self') {
-        // For self messages, always use the right side
-        sourceHandle = 'right';
-        targetHandle = 'right';
-      } else if (sourceX > targetX) {
-        // If source is to the right of target (server-to-client)
-        sourceHandle = 'left';
-        targetHandle = 'right';
-      } else {
-        // If source is to the left of target (client-to-server)
-        sourceHandle = 'right';
-        targetHandle = 'left';
-      }
-
-      // Create a message from the source to the target entity
-      const newEdge: Edge<SequenceMessageData> = {
-        id: `seq-edge-${edges.length + 1}`,
-        source: sourceId,
-        target: targetId,
-        type: 'sequenceMessage',
-        sourceHandle: sourceHandle,
-        targetHandle: targetHandle,
-        data: {
-          label: label,
-          type: messageType,
-          y: messageY,
-        },
-        animated: messageType === 'async',
-      };
-
-      setEdges((eds) => [...eds, newEdge]);
-
-      setNodes((nds) => {
-        return nds.map((node) => {
-          if (
-            node.type === 'sequenceEntity' &&
-            (node.id === sourceId || node.id === targetId)
-          ) {
-            const nodeData = node.data as SequenceEntityData;
-            // Find all messages (edges) connected to this node, including the new one
-            const connectedMessages = [...edges, newEdge].filter(
-              (edge) => edge.source === node.id || edge.target === node.id,
-            );
-            // Find the max Y position among these messages
-            const maxY = Math.max(
-              ...connectedMessages.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
-              nodeData.lifelineLength || 0, // at least current length
-            );
-            const newLength = Math.max(maxY + 40, nodeData.lifelineLength || 0);
-            // Only update if needed
-            if (newLength > (nodeData.lifelineLength || 0)) {
-              return {
-                ...node,
-                data: {
-                  ...nodeData,
-                  lifelineLength: newLength,
-                },
-              };
-            }
-          }
-          return node;
-        });
-      });
-    },
-    [nodes, edges, setEdges, setNodes],
-  );
-
-  // Handler for the original add message button (now opens the dialog)
-  const handleAddMessage = useCallback(() => {
-    handleOpenAddMessageDialog();
-  }, [handleOpenAddMessageDialog]);
-
-  // 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 = t('sequenceDiagram.export.filename');
-
-    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(t('sequenceDiagram.errors.importFailed'));
-        }
-      };
-    },
-    [setNodes, setEdges],
-  );
-
-  // Function to handle theme changes
-  const handleThemeChange = useCallback(
-    (newTheme: ThemeData) => {
-      setTheme(newTheme);
-
-      // Update nodes with new theme colors
-      setNodes((nds) =>
-        nds.map((node) => {
-          if (node.type === 'sequenceEntity') {
-            const nodeData = node.data;
-            return {
-              ...node,
-              data: {
-                ...nodeData,
-                color: newTheme.primary,
-                backgroundColor: newTheme.secondary,
-              },
-            };
-          } else if (node.type === 'activationBar') {
-            const nodeData = node.data;
-            return {
-              ...node,
-              data: {
-                ...nodeData,
-                color: newTheme.secondary,
-                borderColor: newTheme.primary,
-              },
-            };
-          }
-          return node;
-        }),
-      );
-    },
-    [setNodes],
-  );
-
-  // Toggle fullscreen mode
-  const toggleFullscreen = useCallback(() => {
-    if (!diagramContainerRef.current) return;
-
-    const element = diagramContainerRef.current as HTMLElementWithFullscreen;
-    const doc = document as DocumentWithFullscreen;
-
-    if (!isFullscreen) {
-      if (element.requestFullscreen) {
-        element.requestFullscreen();
-      } else if (element.webkitRequestFullscreen) {
-        element.webkitRequestFullscreen();
-      } else if (element.msRequestFullscreen) {
-        element.msRequestFullscreen();
-      }
-    } else {
-      if (doc.exitFullscreen) {
-        doc.exitFullscreen();
-      } else if (doc.webkitExitFullscreen) {
-        doc.webkitExitFullscreen();
-      } else if (doc.msExitFullscreen) {
-        doc.msExitFullscreen();
-      }
-    }
-
-    setIsFullscreen(!isFullscreen);
-  }, [isFullscreen]);
-
-  // Listen for fullscreen change events from browser
-  React.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,
-      );
-    };
-  }, []);
-
-  // Properly handle connections between nodes
-  const onConnect = useCallback(
-    (params: Connection) => {
-      console.log('Connection params:', params);
-      if (!params.source || !params.target) return;
-
-      // Find the source and target nodes
-      const sourceNode = nodes.find((n) => n.id === params.source);
-      const targetNode = nodes.find((n) => n.id === params.target);
-
-      if (!sourceNode || !targetNode) return;
-
-      // Only allow connections between sequence entities
-      if (
-        sourceNode.type !== 'sequenceEntity' ||
-        targetNode.type !== 'sequenceEntity'
-      ) {
-        alert(t('sequenceDiagram.errors.onlyConnectEntities'));
-        return;
-      }
-
-      // Calculate a Y position for the new message
-      const messageY =
-        Math.max(
-          ...edges.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
-          180, // Default minimum Y position
-        ) + 60; // Add some spacing
-
-      // Calculate handle positions based on nodes
-      const sourceX = sourceNode.position?.x || 0;
-      const targetX = targetNode.position?.x || 0;
-
-      // Determine correct handles based on relative positions
-      let sourceHandle = params.sourceHandle;
-      let targetHandle = params.targetHandle;
-
-      // If handles aren't explicitly provided by the connection event,
-      // calculate them based on node positions
-      if (!sourceHandle || !targetHandle) {
-        if (sourceX > targetX) {
-          // Source is to the right of target (server-to-client)
-          sourceHandle = 'left';
-          targetHandle = 'right';
-        } else {
-          // Source is to the left of target (client-to-server)
-          sourceHandle = 'right';
-          targetHandle = 'left';
-        }
-      }
-
-      // Create a new edge for the message
-      const newEdge: Edge<SequenceMessageData> = {
-        id: `seq-edge-${edges.length + 1}`,
-        source: params.source,
-        target: params.target,
-        type: 'sequenceMessage',
-        sourceHandle: sourceHandle,
-        targetHandle: targetHandle,
-        data: {
-          label: `Message ${edges.length + 1}`,
-          type: 'sync', // Default to sync message
-          y: messageY,
-        },
-        animated: false,
-      };
-
-      setEdges((eds) => [...eds, newEdge]);
-
-      setNodes((nds) => {
-        return nds.map((node) => {
-          if (
-            node.type === 'sequenceEntity' &&
-            (node.id === params.source || node.id === params.target)
-          ) {
-            const nodeData = node.data as SequenceEntityData;
-            // Find all messages (edges) connected to this node, including the new one
-            const connectedMessages = [...edges, newEdge].filter(
-              (edge) => edge.source === node.id || edge.target === node.id,
-            );
-            // Find the max Y position among these messages
-            const maxY = Math.max(
-              ...connectedMessages.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
-              nodeData.lifelineLength || 0, // at least current length
-            );
-            const newLength = Math.max(maxY + 40, nodeData.lifelineLength || 0);
-            // Only update if needed
-            if (newLength > (nodeData.lifelineLength || 0)) {
-              return {
-                ...node,
-                data: {
-                  ...nodeData,
-                  lifelineLength: newLength,
-                },
-              };
-            }
-          }
-          return node;
-        });
-      });
-    },
-    [nodes, edges, setEdges, setNodes],
-  );
-
-  // Handle node click events - use ReactFlow's NodeMouseHandler type
-  const onNodeClick: NodeMouseHandler = useCallback(
-    (event, node) => {
-      // If the node is a combined fragment, select it for editing
-      if (node.type === 'combinedFragment') {
-        setSelectedFragmentId(node.id);
-
-        // Set current fragment dimensions
-        const width = node.style?.width || 400;
-        const height = node.style?.height || 150;
-        setFragmentSettings({ width: Number(width), height: Number(height) });
-
-        // Show the settings panel
-        setShowFragmentSettings(true);
-        setShowEntitySettings(false);
-      }
-      // If the node is an entity, select it for settings
-      else if (node.type === 'sequenceEntity') {
-        setSelectedEntityId(node.id);
-        setShowEntitySettings(true);
-        setShowFragmentSettings(false);
-      } else {
-        // If clicking on a non-fragment node while settings are open, close them
-        if (showFragmentSettings) {
-          setShowFragmentSettings(false);
-          setSelectedFragmentId(null);
-        }
-        if (showEntitySettings) {
-          setShowEntitySettings(false);
-          setSelectedEntityId(null);
-        }
-      }
-    },
-    [showFragmentSettings, showEntitySettings],
-  );
-
-  // Handle node double-click for label editing - use ReactFlow's NodeMouseHandler type
-  const onNodeDoubleClick: NodeMouseHandler = useCallback(
-    (event, node) => {
-      event.preventDefault();
-      event.stopPropagation();
-
-      if (node.type === 'sequenceEntity') {
-        const nodeData = node.data as SequenceEntityData;
-        setEditingNode(node.id);
-        setEditingEdge(null);
-        setEditingLabel(nodeData.label || '');
-        setEditingType('node');
-      } else if (node.type === 'combinedFragment') {
-        const nodeData = node.data as CombinedFragmentData;
-        setEditingNode(node.id);
-        setEditingEdge(null);
-        setEditingLabel(nodeData.label || '');
-        setEditingType('fragment');
-        setEditingConditions(nodeData.conditions || []);
-      }
-    },
-    [],
-  );
-
-  // Handle edge double-click for label editing - use ReactFlow's EdgeMouseHandler type
-  const onEdgeDoubleClick: EdgeMouseHandler = useCallback(
-    (event, edge) => {
-      event.preventDefault();
-      event.stopPropagation();
-
-      if (edge.type === 'sequenceMessage') {
-        const edgeData = edge.data as SequenceMessageData;
-        setEditingEdge(edge.id);
-        setEditingNode(null);
-        setEditingLabel(edgeData?.label || '');
-        setEditingType('edge');
-      }
-    },
-    [],
-  );
-
-  // Handle label input change
-  const handleLabelChange = useCallback(
-    (e: React.ChangeEvent<HTMLInputElement>) => {
-      setEditingLabel(e.target.value);
-    },
-    [],
-  );
-
-  // Handle condition change
-  const handleConditionChange = (idx: number, value: string) => {
-    const updated = [...editingConditions];
-    updated[idx] = value;
-    setEditingConditions(updated);
-  };
-
-  // Handle add condition
-  const handleAddCondition = () => {
-    setEditingConditions([...editingConditions, '']);
-  };
-
-  // Handle remove condition
-  const handleRemoveCondition = (idx: number) => {
-    setEditingConditions(editingConditions.filter((_, i) => i !== idx));
-  };
-
-  // Save the edited label
-  const saveLabel = useCallback(() => {
-    if (editingNode && editingType === 'node') {
-      setNodes((nds) =>
-        nds.map((node) => {
-          if (node.id === editingNode) {
-            return {
-              ...node,
-              data: {
-                ...node.data,
-                label: editingLabel,
-              },
-            };
-          }
-          return node;
-        }),
-      );
-      setEditingNode(null);
-    } else if (editingNode && editingType === 'fragment') {
-      setNodes((nds) =>
-        nds.map((node) => {
-          if (node.id === editingNode) {
-            return {
-              ...node,
-              data: {
-                ...node.data,
-                label: editingLabel,
-                conditions: editingConditions,
-              },
-            };
-          }
-          return node;
-        }),
-      );
-      setEditingNode(null);
-      setEditingType('node');
-    } else if (editingEdge && editingType === 'edge') {
-      setEdges((eds) =>
-        eds.map((edge) => {
-          if (edge.id === editingEdge) {
-            return {
-              ...edge,
-              data: {
-                ...edge.data,
-                label: editingLabel,
-              },
-            };
-          }
-          return edge;
-        }),
-      );
-      setEditingEdge(null);
-    }
-
-    // Reset editing state
-    setEditingType('node');
-  }, [
-    editingNode,
-    editingEdge,
-    editingType,
-    editingLabel,
-    editingConditions,
-    setNodes,
-    setEdges,
-  ]);
-
-  // Handle keyboard events for the input field
-  const handleKeyDown = useCallback(
-    (e: React.KeyboardEvent) => {
-      if (e.key === 'Enter') {
-        saveLabel();
-      } else if (e.key === 'Escape') {
-        setEditingNode(null);
-        setEditingEdge(null);
-      }
-    },
-    [saveLabel],
-  );
-
-  // Cancel editing
-  const cancelEditing = useCallback(() => {
-    setEditingNode(null);
-    setEditingEdge(null);
-    setEditingType('node');
-  }, []);
-
-  // Handler to add a new alt fragment
-  const handleAddAltFragment = useCallback(
-    (fragmentType: CombinedFragmentType = 'alt') => {
-      console.log(`Adding new ${fragmentType} fragment...`);
-
-      // Find all fragment nodes
-      const fragmentNodes = getCombinedFragmentNodes();
-      const fragmentCount = fragmentNodes.length;
-      const newFragmentId = `${fragmentType}-fragment-${fragmentCount + 1}`;
-
-      // Find the lowest y position among fragments
-      const defaultX = 250; // Fixed x position for all fragments
-      const defaultY = 300; // Default starting y for fragments
-
-      const lowestY =
-        fragmentNodes.length > 0
-          ? Math.max(...fragmentNodes.map((node) => node.position.y))
-          : defaultY;
-
-      const yOffset = 200; // Space between fragments vertically
-      const centerX = defaultX;
-      const centerY = fragmentNodes.length > 0 ? lowestY + yOffset : lowestY;
-
-      // Set up conditions based on fragment type
-      let conditions: string[] = [];
-      switch (fragmentType) {
-        case 'alt':
-          conditions = ['[condition]', '[else]'];
-          break;
-        case 'opt':
-          conditions = ['[condition]'];
-          break;
-        case 'loop':
-          conditions = ['[loop condition]'];
-          break;
-        case 'par':
-          conditions = ['[parallel 1]', '[parallel 2]'];
-          break;
-        default:
-          conditions = ['[condition]'];
-      }
-
-      // Create a new fragment node
-      const newFragment: Node<CombinedFragmentData> = {
-        id: newFragmentId,
-        type: 'combinedFragment',
-        position: { x: centerX, y: centerY },
-        style: {
-          width: fragmentSettings.width,
-          height: fragmentSettings.height,
-        },
-        data: {
-          label: fragmentType,
-          conditions: conditions,
-          style: {
-            width: fragmentSettings.width,
-            height: fragmentSettings.height,
-          },
-          color: '#000000',
-          backgroundColor: '#ffffff',
-          lifelineLength: 0,
-        },
-        draggable: true,
-        zIndex: 0, // Lower z-index to appear behind other elements
-      };
-
-      setNodes((nds) => [...nds, newFragment]);
-      setSelectedFragmentId(newFragmentId);
-      setShowFragmentSettings(true);
-    },
-    [fragmentSettings, getCombinedFragmentNodes, setNodes],
-  );
-
-  // Function to update selected fragment
-  const updateSelectedFragment = useCallback(
-    (fragmentId: string, updates: FragmentSettingsUpdateData) => {
-      setNodes((nds) =>
-        nds.map((node) => {
-          if (node.id === fragmentId) {
-            // Safely handle the style spread by ensuring the fragment data has a style property
-            const fragmentData = node.data;
-            const currentStyle = fragmentData.style || {};
-            // Apply updates to both style and data.style to ensure consistency
-            const updatedNode = {
-              ...node,
-              style: { ...node.style, ...updates },
-              data: {
-                ...fragmentData,
-                style: { ...currentStyle, ...updates },
-              },
-            };
-            return updatedNode;
-          }
-          return node;
-        }),
-      );
-    },
-    [setNodes],
-  );
-
-  const handleFragmentWidthChange = useCallback(
-    (width: number) => {
-      setFragmentSettings((prev) => ({ ...prev, width }));
-
-      if (selectedFragmentId) {
-        updateSelectedFragment(selectedFragmentId, { width });
-      }
-    },
-    [selectedFragmentId, updateSelectedFragment],
-  );
-
-  const handleFragmentHeightChange = useCallback(
-    (height: number) => {
-      setFragmentSettings((prev) => ({ ...prev, height }));
-
-      if (selectedFragmentId) {
-        updateSelectedFragment(selectedFragmentId, { height });
-      }
-    },
-    [selectedFragmentId, updateSelectedFragment],
-  );
-
-  const closeFragmentSettings = useCallback(() => {
-    setShowFragmentSettings(false);
-    setSelectedFragmentId(null);
-  }, []);
-
-  // Function to update selected entity
-  const updateSelectedEntity = useCallback(
-    (updates: EntitySettingsUpdateData) => {
-      if (!selectedEntityId) return;
-
-      setNodes((nds) =>
-        nds.map((node) => {
-          if (node.id === selectedEntityId) {
-            return {
-              ...node,
-              data: {
-                ...node.data,
-                ...updates,
-              },
-            };
-          }
-          return node;
-        }),
-      );
-    },
-    [selectedEntityId, setNodes],
-  );
-
-  // Close entity settings
-  const closeEntitySettings = useCallback(() => {
-    setShowEntitySettings(false);
-    setSelectedEntityId(null);
-  }, []);
-
-  // Handler to load a template diagram
-  const loadTemplateData = useCallback(
-    (templateType: string) => {
-      console.log(`Loading template: ${templateType}`);
-
-      const template = getTemplateData(templateType);
-
-      if (template.nodes.length > 0) {
-        setNodes(template.nodes);
-        setEdges(template.edges);
-      }
-    },
-    [setNodes, setEdges],
-  );
-
-  return (
-    <>
-      <Box sx={{ display: 'flex', width: '100%', height: '100vh' }}>
-        <Sidebar>
-          <DiagramControls
-            onAddEntity={handleAddEntity}
-            onAddMessage={handleAddMessage}
-            onAddAltFragment={handleAddAltFragment}
-            onUndo={handleUndo}
-            onRedo={handleRedo}
-            canUndo={historyIndex > 0}
-            canRedo={historyIndex < history.length - 1}
-            onLoadTemplate={loadTemplateData}
-          />
-        </Sidebar>
-        <Box sx={{ flex: 1, height: '100%', position: 'relative' }}>
-          <div
-            ref={diagramContainerRef}
-            className={`diagram-container ${isFullscreen ? 'fullscreen' : ''}`}
-            style={{
-              width: '100%',
-              height: '100%',
-              border: '1px solid #333',
-              position: 'relative',
-              backgroundColor: '#1a1a1a',
-              color: '#ffffff',
-            }}
-          >
-            <ReactFlow<Node, Edge>
-              nodes={nodes}
-              edges={edges}
-              onNodesChange={onNodesChangeLockedY}
-              onEdgesChange={onEdgesChange}
-              onConnect={onConnect}
-              nodeTypes={sequenceNodeTypes}
-              edgeTypes={sequenceEdgeTypes}
-              fitView
-              attributionPosition="bottom-left"
-              connectOnClick={false}
-              onNodeClick={onNodeClick}
-              onNodeDoubleClick={onNodeDoubleClick}
-              onEdgeDoubleClick={onEdgeDoubleClick}
-              style={{
-                backgroundColor: '#1a1a1a',
-              }}
-            >
-              <Background
-                color="#333"
-                variant={BackgroundVariant.Dots}
-                gap={16}
-              />
-            </ReactFlow>
-
-            {/* Components outside ReactFlow */}
-            <DiagramPanel
-              historyIndex={historyIndex}
-              historyLength={history.length}
-              onUndo={handleUndo}
-              onRedo={handleRedo}
-              onClear={() => {
-                const clearedState: SequenceHistoryEntry = { nodes: [], edges: [] };
-                const newHistory = history.slice(0, historyIndex + 1);
-                newHistory.push(clearedState);
-                setNodes([]);
-                setEdges([]);
-                setHistory(newHistory);
-                setHistoryIndex(newHistory.length - 1);
-              }}
-              onImport={() => { }}
-              isSettingsOpen={false}
-              onToggleSettings={() => { }}
-            />
-
-            <AiButton />
-
-            <MiniMapCommon />
-
-            {/* Fragment settings panel */}
-            {showFragmentSettings && (
-              <FragmentSettings
-                fragmentId={selectedFragmentId}
-                width={fragmentSettings.width}
-                height={fragmentSettings.height}
-                onWidthChange={handleFragmentWidthChange}
-                onHeightChange={handleFragmentHeightChange}
-                onClose={closeFragmentSettings}
-              />
-            )}
-
-            {/* Entity settings panel */}
-            {showEntitySettings && selectedEntityId && (
-              <EntitySettings
-                entityId={selectedEntityId}
-                data={
-                  (nodes.find((node) =>
-                    node.id === selectedEntityId && node.type === 'sequenceEntity'
-                  )?.data as SequenceEntityData) || {
-                    label: '',
-                    color: '#61dafb',
-                    backgroundColor: '#2a2a2a',
-                    lifelineLength: 500
-                  }
-                }
-                onUpdate={updateSelectedEntity}
-                onClose={closeEntitySettings}
-              />
-            )}
-
-            {/* Label editing overlay */}
-            {(editingNode || editingEdge) && (
-              <Box
-                sx={{
-                  position: 'absolute',
-                  zIndex: 10,
-                  top: 10,
-                  left: '50%',
-                  transform: 'translateX(-50%)',
-                  bgcolor: '#242424',
-                  border: '1px solid #333',
-                  color: '#fff',
-                  p: 2,
-                  borderRadius: 1,
-                  boxShadow: '0 4px 12px rgba(0, 0, 0, 0.3)',
-                  minWidth: 320,
-                  '& .MuiTextField-root': {
-                    '& .MuiInputLabel-root': {
-                      color: '#999',
-                    },
-                    '& .MuiOutlinedInput-root': {
-                      color: '#fff',
-                      '& fieldset': {
-                        borderColor: '#444',
-                      },
-                      '&:hover fieldset': {
-                        borderColor: '#61dafb',
-                      },
-                      '&.Mui-focused fieldset': {
-                        borderColor: '#61dafb',
-                      },
-                    },
-                  },
-                  '& .MuiButton-contained': {
-                    bgcolor: '#61dafb',
-                    color: '#000',
-                    '&:hover': {
-                      bgcolor: '#4db6d6',
-                    },
-                  },
-                  '& .MuiButton-outlined': {
-                    borderColor: '#333',
-                    color: '#fff',
-                    '&:hover': {
-                      borderColor: '#61dafb',
-                      bgcolor: 'rgba(97, 218, 251, 0.1)',
-                    },
-                  },
-                  '& .MuiIconButton-root': {
-                    color: '#999',
-                    '&:hover': {
-                      color: '#fff',
-                      bgcolor: 'rgba(255, 255, 255, 0.1)',
-                    },
-                  },
-                }}
-              >
-                <Typography sx={{ mb: 1, fontWeight: 'bold', color: '#fff' }}>
-                  {editingType === 'node'
-                    ? t('sequenceFlow.editLabels.entityLabel')
-                    : editingType === 'edge'
-                      ? t('sequenceFlow.editLabels.messageLabel')
-                      : t('sequenceFlow.editLabels.editFragment')}
-                </Typography>
-                {editingType === 'fragment' ? (
-                  <>
-                    <Box display="flex" alignItems="center" gap={2} mb={2}>
-                      <TextField
-                        type="text"
-                        value={editingLabel}
-                        onChange={handleLabelChange}
-                        onKeyDown={handleKeyDown}
-                        autoFocus
-                        size="small"
-                        sx={{ width: 200 }}
-                        label={t('fragmentSettings.fragmentName')}
-                      />
-                    </Box>
-                    {/* Conditions editing */}
-                    <Box mb={2}>
-                      <Typography
-                        variant="body2"
-                        fontWeight="bold"
-                        mb={0.5}
-                        sx={{ color: '#61dafb' }}
-                      >
-                        {t('fragmentSettings.conditions')}
-                      </Typography>
-                      {editingConditions.map((cond, idx) => (
-                        <Box
-                          key={idx}
-                          display="flex"
-                          alignItems="center"
-                          mb={1}
-                        >
-                          <TextField
-                            value={cond}
-                            onChange={(e) =>
-                              handleConditionChange(idx, e.target.value)
-                            }
-                            size="small"
-                            fullWidth
-                            sx={{ mr: 1 }}
-                          />
-                          <IconButton
-                            onClick={() => handleRemoveCondition(idx)}
-                            size="small"
-                          >
-                            <CloseIcon fontSize="small" />
-                          </IconButton>
-                        </Box>
-                      ))}
-                      <Button
-                        onClick={handleAddCondition}
-                        size="small"
-                        variant="outlined"
-                      >
-                        {t('fragmentSettings.addCondition')}
-                      </Button>
-                    </Box>
-                    <Box display="flex" gap={2}>
-                      <Button
-                        variant="contained"
-                        size="small"
-                        onClick={saveLabel}
-                        sx={{ textTransform: 'none' }}
-                      >
-                        {t('common.save')}
-                      </Button>
-                      <Button
-                        variant="outlined"
-                        size="small"
-                        onClick={cancelEditing}
-                        sx={{ textTransform: 'none' }}
-                      >
-                        {t('common.cancel')}
-                      </Button>
-                    </Box>
-                  </>
-                ) : (
-                  <Box display="flex" alignItems="center" gap={2}>
-                    <TextField
-                      type="text"
-                      value={editingLabel}
-                      onChange={handleLabelChange}
-                      onKeyDown={handleKeyDown}
-                      autoFocus
-                      size="small"
-                      sx={{ width: 200 }}
-                    />
-                    <Button
-                      variant="contained"
-                      size="small"
-                      onClick={saveLabel}
-                      sx={{ textTransform: 'none' }}
-                    >
-                      {t('common.save')}
-                    </Button>
-                    <Button
-                      variant="outlined"
-                      size="small"
-                      onClick={cancelEditing}
-                      sx={{ textTransform: 'none' }}
-                    >
-                      {t('common.cancel')}
-                    </Button>
-                  </Box>
-                )}
-              </Box>
-            )}
-
-            {/* Add Message Dialog */}
-            <AddMessageDialog
-              open={showMessageDialog}
-              onClose={() => setShowMessageDialog(false)}
-              onAddMessage={handleAddCustomMessage}
-              entities={sequenceEntities()}
-            />
-          </div>
-        </Box>
-      </Box>
-    </>
-  );
-};
-
-export default SequenceDiagramFlow;
+import React, { useCallback, useState, useRef, useEffect } from 'react';
+import {
+  ReactFlow,
+  Background,
+  useNodesState,
+  useEdgesState,
+  Panel,
+  BackgroundVariant,
+  Connection,
+  NodeChange,
+  Node,
+  Edge,
+  MiniMap,
+  NodeMouseHandler,
+  EdgeMouseHandler,
+} from '@xyflow/react';
+import '@xyflow/react/dist/style.css';
+
+import {
+  Box,
+  Button,
+  TextField,
+  Typography,
+  IconButton,
+  Paper,
+} 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 { useTranslation } from 'react-i18next';
+import {
+  SequenceDiagramProps,
+  ThemeData,
+  SequenceNode,
+  SequenceEdge,
+  SequenceHistoryEntry,
+  SequenceEntityShape,
+  CombinedFragmentType,
+  SequenceMessageType,
+  SequenceEntityNode,
+  SequenceEntityData,
+  FragmentSettingsUpdateData,
+  EntitySettingsUpdateData,
+  ActivationBarNode,
+  CombinedFragmentNode,
+  SequenceMessageData,
+  ActivationBarData,
+  CombinedFragmentData
+} from '../../../types/sequence-types';
+import { getTemplateData } from '../data/template-data';
+import { sequenceEdgeTypes, sequenceNodeTypes } from './SequenceComponents';
+import DiagramControls from '../controls/DiagramControls';
+import FragmentSettings from './FragmentSettings';
+import EntitySettings from './EntitySettings';
+import AddMessageDialog from './AddMessageDialog';
+import Sidebar from '../../../components/SideBar';
+import AiButton from '../../../components/AiButton';
+import MiniMapCommon from '../../../components/MiniMapCommon';
+import DiagramPanel from '../../../components/DiagramPanel';
+
+// Types for browser fullscreen API extensions
+interface HTMLElementWithFullscreen extends HTMLElement {
+  webkitRequestFullscreen?: () => Promise<void>;
+  msRequestFullscreen?: () => Promise<void>;
+}
+
+interface DocumentWithFullscreen extends Document {
+  webkitExitFullscreen?: () => Promise<void>;
+  msExitFullscreen?: () => Promise<void>;
+}
+
+// Main component that manages the diagram state
+
+const SequenceDiagramFlow: React.FC<SequenceDiagramProps> = ({
+  sequenceNodes,
+  sequenceEdges,
+}) => {
+  const { t } = useTranslation();
+
+  // Initialize nodes and edges with our data - use xy/flow base types with proper generics
+  const [nodes, setNodes, onNodesChange] = useNodesState(sequenceNodes);
+
+  const [edges, setEdges, onEdgesChange] = useEdgesState(sequenceEdges);
+
+  // History state for undo/redo functionality
+  const [history, setHistory] = useState<SequenceHistoryEntry[]>([]);
+  const [historyIndex, setHistoryIndex] = useState<number>(-1);
+  const isUndoRedoAction = useRef(false);
+
+  // Initialize history with initial state
+  useEffect(() => {
+    // Only initialize if history is empty
+    if (history.length === 0) {
+      const initialState: SequenceHistoryEntry = {
+        nodes: [...nodes] as SequenceNode[],
+        edges: [...edges] as SequenceEdge[],
+      };
+      setHistory([initialState]);
+      setHistoryIndex(0);
+    }
+  }, []);
+
+  // Save current state to history when nodes or edges change
+  useEffect(() => {
+    // Don't record history when performing undo/redo actions or when we're still initializing
+    if (isUndoRedoAction.current || history.length === 0) {
+      return;
+    }
+
+    // Debounce: Only update history if nodes/edges actually changed (deep compare)
+    const newHistoryEntry: SequenceHistoryEntry = {
+      nodes: [...nodes] as SequenceNode[],
+      edges: [...edges] as SequenceEdge[],
+    };
+    const currentState = history[historyIndex];
+    if (
+      JSON.stringify(currentState.nodes) ===
+      JSON.stringify(newHistoryEntry.nodes) &&
+      JSON.stringify(currentState.edges) ===
+      JSON.stringify(newHistoryEntry.edges)
+    ) {
+      return; // No change, don't add to history
+    }
+
+    // Trim future history if we're in the middle of the history stack
+    const newHistory = history.slice(0, historyIndex + 1);
+    newHistory.push(newHistoryEntry);
+    setHistory(newHistory);
+    setHistoryIndex(newHistory.length - 1);
+  }, [nodes, edges, history, historyIndex]);
+
+  // 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]);
+
+  // 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]);
+
+  const [theme, setTheme] = useState<ThemeData>({
+    primary: '#61dafb',
+    secondary: '#2a2a2a',
+  });
+  const [isFullscreen, setIsFullscreen] = useState<boolean>(false);
+  const [selectedFragmentId, setSelectedFragmentId] = useState<string | null>(
+    null,
+  );
+  const [fragmentSettings, setFragmentSettings] = useState({
+    width: 400,
+    height: 150,
+  });
+  const [showFragmentSettings, setShowFragmentSettings] =
+    useState<boolean>(false);
+  const diagramContainerRef = useRef<HTMLDivElement>(null);
+  const [editingNode, setEditingNode] = useState<string | null>(null);
+  const [editingEdge, setEditingEdge] = useState<string | null>(null);
+  const [editingLabel, setEditingLabel] = useState<string>('');
+  const [editingConditions, setEditingConditions] = useState<string[]>([]);
+  const [editingType, setEditingType] = useState<'node' | 'edge' | 'fragment'>(
+    'node',
+  );
+  const [selectedEntityId, setSelectedEntityId] = useState<string | null>(null);
+  const [showEntitySettings, setShowEntitySettings] = useState<boolean>(false);
+  const [showMessageDialog, setShowMessageDialog] = useState<boolean>(false);
+  const ENTITY_FIXED_Y = 50; // Fixed y position for all entities
+
+  // Type-safe node access helpers
+  const getSequenceEntityNodes = useCallback(() => {
+    return nodes.filter((node): node is Node<SequenceEntityData> =>
+      node.type === 'sequenceEntity'
+    );
+  }, [nodes]);
+
+  const getCombinedFragmentNodes = useCallback(() => {
+    return nodes.filter((node): node is Node<CombinedFragmentData> =>
+      node.type === 'combinedFragment'
+    );
+  }, [nodes]);
+
+  // Intercept node position changes to lock y for sequenceEntity nodes
+  const onNodesChangeLockedY = useCallback(
+    (changes: NodeChange[]) => {
+      let hasPositionChange = false;
+      const updatedNodes = nodes.map((node) => {
+        // Find if there's a position change for this node
+        const change = changes.find(
+          (c: NodeChange) => c.type === 'position' && c.id === node.id,
+        );
+        if (
+          change &&
+          change.type === 'position' &&
+          node.type === 'sequenceEntity' &&
+          change.position
+        ) {
+          hasPositionChange = true;
+          // Lock y position
+          return {
+            ...node,
+            position: {
+              x: change.position.x,
+              y: ENTITY_FIXED_Y, // always fixed
+            },
+          };
+        }
+        return node;
+      });
+      if (hasPositionChange) {
+        setNodes(updatedNodes);
+      } else {
+        onNodesChange(changes);
+      }
+    },
+    [nodes, setNodes, onNodesChange],
+  );
+
+  // Handler to add a new entity
+  const handleAddEntity = useCallback(
+    (shape: SequenceEntityShape = 'rectangle') => {
+      console.log('Adding new entity with shape:', shape);
+      console.log('Current nodes:', nodes);
+
+      // Find all entity nodes
+      const entityNodes = getSequenceEntityNodes();
+      const entityCount = entityNodes.length;
+      const newEntityId = `entity-${entityCount + 1}`;
+
+      // Find the rightmost x position among entities
+      const rightmostX =
+        entityNodes.length > 0
+          ? Math.max(...entityNodes.map((node) => node.position.x))
+          : 100; // Default starting x
+
+      const xOffset = 150; // Space between entities
+      const xPosition =
+        entityNodes.length > 0 ? rightmostX + xOffset : rightmostX;
+
+      // Create a new entity node with extended properties
+      const newEntity: Node<SequenceEntityData> = {
+        id: newEntityId,
+        type: 'sequenceEntity',
+        position: { x: xPosition, y: 50 },
+        data: {
+          label: `Entity ${entityCount + 1}`,
+          color: theme.primary,
+          backgroundColor: theme.secondary,
+          lifelineLength: 500, // Use same height as others
+          shape: shape, // Use the provided shape
+          fontSize: 12,
+          borderWidth: 2,
+          description: '',
+          collapsed: false,
+        },
+        draggable: true,
+        zIndex: 5, // Higher z-index to appear above other elements
+      };
+
+      // Adjust position based on entity shape
+      let offsetX = 55;
+      switch (shape) {
+        case 'circle':
+          offsetX = 40;
+          break;
+        case 'database':
+          offsetX = 45;
+          break;
+        case 'actor':
+          offsetX = 60;
+          break;
+        default:
+          offsetX = 55;
+          break;
+      }
+
+      const activationBar: Node<ActivationBarData> = {
+        id: `${newEntityId}-activation`,
+        type: 'activationBar',
+        position: { x: xPosition + offsetX, y: 180 }, // Position below the entity
+        data: {
+          entityId: newEntityId,
+          color: theme.secondary,
+          borderColor: theme.primary,
+          length: 150,
+        },
+        draggable: true,
+        zIndex: 1, // Lower z-index to appear behind entity boxes but above combined fragments
+      };
+
+      setNodes((nds) => [...nds, newEntity, activationBar]);
+    },
+    [nodes, setNodes, theme, getSequenceEntityNodes],
+  );
+
+  // Filter entities for Add Message Dialog
+  const sequenceEntities = useCallback((): SequenceEntityNode[] => {
+    return getSequenceEntityNodes().map((node) => ({
+      id: node.id,
+      type: 'sequenceEntity',
+      data: node.data,
+      position: node.position,
+      draggable: node.draggable,
+      zIndex: node.zIndex,
+    }));
+  }, [getSequenceEntityNodes]);
+
+  // Handler to open add message dialog
+  const handleOpenAddMessageDialog = useCallback(() => {
+    // Filter to get only entity nodes
+    const entityNodes = getSequenceEntityNodes();
+
+    if (entityNodes.length < 1) {
+      alert(t('sequenceDiagram.errors.needEntitiesForMessage'));
+      return;
+    }
+
+    setShowMessageDialog(true);
+  }, [getSequenceEntityNodes]);
+
+  // Handler to create a new message with specified source, target, and type
+  const handleAddCustomMessage = useCallback(
+    (
+      sourceId: string,
+      targetId: string,
+      messageType: SequenceMessageType,
+      label: string,
+    ) => {
+      // Find the source and target nodes
+      const sourceNode = nodes.find((n) => n.id === sourceId);
+      const targetNode = nodes.find((n) => n.id === targetId);
+
+      if (!sourceNode || !targetNode) {
+        console.error(t('sequenceDiagram.errors.sourceTargetNotFound'));
+        return;
+      }
+
+      // Calculate a reasonable Y position for the new message
+      const messageY =
+        Math.max(
+          ...edges.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
+          180, // Default minimum Y position
+        ) + 60; // Add some spacing
+
+      // Calculate handle positions based on nodes
+      const sourceX = sourceNode.position?.x || 0;
+      const targetX = targetNode.position?.x || 0;
+
+      // Calculate source and target handles
+      let sourceHandle = 'right';
+      let targetHandle = 'left';
+
+      if (messageType === 'self') {
+        // For self messages, always use the right side
+        sourceHandle = 'right';
+        targetHandle = 'right';
+      } else if (sourceX > targetX) {
+        // If source is to the right of target (server-to-client)
+        sourceHandle = 'left';
+        targetHandle = 'right';
+      } else {
+        // If source is to the left of target (client-to-server)
+        sourceHandle = 'right';
+        targetHandle = 'left';
+      }
+
+      // Create a message from the source to the target entity
+      const newEdge: Edge<SequenceMessageData> = {
+        id: `seq-edge-${edges.length + 1}`,
+        source: sourceId,
+        target: targetId,
+        type: 'sequenceMessage',
+        sourceHandle: sourceHandle,
+        targetHandle: targetHandle,
+        data: {
+          label: label,
+          type: messageType,
+          y: messageY,
+        },
+        animated: messageType === 'async',
+      };
+
+      setEdges((eds) => [...eds, newEdge]);
+
+      setNodes((nds) => {
+        return nds.map((node) => {
+          if (
+            node.type === 'sequenceEntity' &&
+            (node.id === sourceId || node.id === targetId)
+          ) {
+            const nodeData = node.data as SequenceEntityData;
+            // Find all messages (edges) connected to this node, including the new one
+            const connectedMessages = [...edges, newEdge].filter(
+              (edge) => edge.source === node.id || edge.target === node.id,
+            );
+            // Find the max Y position among these messages
+            const maxY = Math.max(
+              ...connectedMessages.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
+              nodeData.lifelineLength || 0, // at least current length
+            );
+            const newLength = Math.max(maxY + 40, nodeData.lifelineLength || 0);
+            // Only update if needed
+            if (newLength > (nodeData.lifelineLength || 0)) {
+              return {
+                ...node,
+                data: {
+                  ...nodeData,
+                  lifelineLength: newLength,
+                },
+              };
+            }
+          }
+          return node;
+        });
+      });
+    },
+    [nodes, edges, setEdges, setNodes],
+  );
+
+  // Handler for the original add message button (now opens the dialog)
+  const handleAddMessage = useCallback(() => {
+    handleOpenAddMessageDialog();
+  }, [handleOpenAddMessageDialog]);
+
+  // 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 = t('sequenceDiagram.export.filename');
+
+    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(t('sequenceDiagram.errors.importFailed'));
+        }
+      };
+    },
+    [setNodes, setEdges],
+  );
+
+  // Function to handle theme changes
+  const handleThemeChange = useCallback(
+    (newTheme: ThemeData) => {
+      setTheme(newTheme);
+
+      // Update nodes with new theme colors
+      setNodes((nds) =>
+        nds.map((node) => {
+          if (node.type === 'sequenceEntity') {
+            const nodeData = node.data;
+            return {
+              ...node,
+              data: {
+                ...nodeData,
+                color: newTheme.primary,
+                backgroundColor: newTheme.secondary,
+              },
+            };
+          } else if (node.type === 'activationBar') {
+            const nodeData = node.data;
+            return {
+              ...node,
+              data: {
+                ...nodeData,
+                color: newTheme.secondary,
+                borderColor: newTheme.primary,
+              },
+            };
+          }
+          return node;
+        }),
+      );
+    },
+    [setNodes],
+  );
+
+  // Toggle fullscreen mode
+  const toggleFullscreen = useCallback(() => {
+    if (!diagramContainerRef.current) return;
+
+    const element = diagramContainerRef.current as HTMLElementWithFullscreen;
+    const doc = document as DocumentWithFullscreen;
+
+    if (!isFullscreen) {
+      if (element.requestFullscreen) {
+        element.requestFullscreen();
+      } else if (element.webkitRequestFullscreen) {
+        element.webkitRequestFullscreen();
+      } else if (element.msRequestFullscreen) {
+        element.msRequestFullscreen();
+      }
+    } else {
+      if (doc.exitFullscreen) {
+        doc.exitFullscreen();
+      } else if (doc.webkitExitFullscreen) {
+        doc.webkitExitFullscreen();
+      } else if (doc.msExitFullscreen) {
+        doc.msExitFullscreen();
+      }
+    }
+
+    setIsFullscreen(!isFullscreen);
+  }, [isFullscreen]);
+
+  // Listen for fullscreen change events from browser
+  React.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,
+      );
+    };
+  }, []);
+
+  // Properly handle connections between nodes
+  const onConnect = useCallback(
+    (params: Connection) => {
+      console.log('Connection params:', params);
+      if (!params.source || !params.target) return;
+
+      // Find the source and target nodes
+      const sourceNode = nodes.find((n) => n.id === params.source);
+      const targetNode = nodes.find((n) => n.id === params.target);
+
+      if (!sourceNode || !targetNode) return;
+
+      // Only allow connections between sequence entities
+      if (
+        sourceNode.type !== 'sequenceEntity' ||
+        targetNode.type !== 'sequenceEntity'
+      ) {
+        alert(t('sequenceDiagram.errors.onlyConnectEntities'));
+        return;
+      }
+
+      // Calculate a Y position for the new message
+      const messageY =
+        Math.max(
+          ...edges.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
+          180, // Default minimum Y position
+        ) + 60; // Add some spacing
+
+      // Calculate handle positions based on nodes
+      const sourceX = sourceNode.position?.x || 0;
+      const targetX = targetNode.position?.x || 0;
+
+      // Determine correct handles based on relative positions
+      let sourceHandle = params.sourceHandle;
+      let targetHandle = params.targetHandle;
+
+      // If handles aren't explicitly provided by the connection event,
+      // calculate them based on node positions
+      if (!sourceHandle || !targetHandle) {
+        if (sourceX > targetX) {
+          // Source is to the right of target (server-to-client)
+          sourceHandle = 'left';
+          targetHandle = 'right';
+        } else {
+          // Source is to the left of target (client-to-server)
+          sourceHandle = 'right';
+          targetHandle = 'left';
+        }
+      }
+
+      // Create a new edge for the message
+      const newEdge: Edge<SequenceMessageData> = {
+        id: `seq-edge-${edges.length + 1}`,
+        source: params.source,
+        target: params.target,
+        type: 'sequenceMessage',
+        sourceHandle: sourceHandle,
+        targetHandle: targetHandle,
+        data: {
+          label: `Message ${edges.length + 1}`,
+          type: 'sync', // Default to sync message
+          y: messageY,
+        },
+        animated: false,
+      };
+
+      setEdges((eds) => [...eds, newEdge]);
+
+      setNodes((nds) => {
+        return nds.map((node) => {
+          if (
+            node.type === 'sequenceEntity' &&
+            (node.id === params.source || node.id === params.target)
+          ) {
+            const nodeData = node.data as SequenceEntityData;
+            // Find all messages (edges) connected to this node, including the new one
+            const connectedMessages = [...edges, newEdge].filter(
+              (edge) => edge.source === node.id || edge.target === node.id,
+            );
+            // Find the max Y position among these messages
+            const maxY = Math.max(
+              ...connectedMessages.map((edge) => (edge.data as SequenceMessageData)?.y || 0),
+              nodeData.lifelineLength || 0, // at least current length
+            );
+            const newLength = Math.max(maxY + 40, nodeData.lifelineLength || 0);
+            // Only update if needed
+            if (newLength > (nodeData.lifelineLength || 0)) {
+              return {
+                ...node,
+                data: {
+                  ...nodeData,
+                  lifelineLength: newLength,
+                },
+              };
+            }
+          }
+          return node;
+        });
+      });
+    },
+    [nodes, edges, setEdges, setNodes],
+  );
+
+  // Handle node click events - use ReactFlow's NodeMouseHandler type
+  const onNodeClick: NodeMouseHandler = useCallback(
+    (event, node) => {
+      // If the node is a combined fragment, select it for editing
+      if (node.type === 'combinedFragment') {
+        setSelectedFragmentId(node.id);
+
+        // Set current fragment dimensions
+        const width = node.style?.width || 400;
+        const height = node.style?.height || 150;
+        setFragmentSettings({ width: Number(width), height: Number(height) });
+
+        // Show the settings panel
+        setShowFragmentSettings(true);
+        setShowEntitySettings(false);
+      }
+      // If the node is an entity, select it for settings
+      else if (node.type === 'sequenceEntity') {
+        setSelectedEntityId(node.id);
+        setShowEntitySettings(true);
+        setShowFragmentSettings(false);
+      } else {
+        // If clicking on a non-fragment node while settings are open, close them
+        if (showFragmentSettings) {
+          setShowFragmentSettings(false);
+          setSelectedFragmentId(null);
+        }
+        if (showEntitySettings) {
+          setShowEntitySettings(false);
+          setSelectedEntityId(null);
+        }
+      }
+    },
+    [showFragmentSettings, showEntitySettings],
+  );
+
+  // Handle node double-click for label editing - use ReactFlow's NodeMouseHandler type
+  const onNodeDoubleClick: NodeMouseHandler = useCallback(
+    (event, node) => {
+      event.preventDefault();
+      event.stopPropagation();
+
+      if (node.type === 'sequenceEntity') {
+        const nodeData = node.data as SequenceEntityData;
+        setEditingNode(node.id);
+        setEditingEdge(null);
+        setEditingLabel(nodeData.label || '');
+        setEditingType('node');
+      } else if (node.type === 'combinedFragment') {
+        const nodeData = node.data as CombinedFragmentData;
+        setEditingNode(node.id);
+        setEditingEdge(null);
+        setEditingLabel(nodeData.label || '');
+        setEditingType('fragment');
+        setEditingConditions(nodeData.conditions || []);
+      }
+    },
+    [],
+  );
+
+  // Handle edge double-click for label editing - use ReactFlow's EdgeMouseHandler type
+  const onEdgeDoubleClick: EdgeMouseHandler = useCallback(
+    (event, edge) => {
+      event.preventDefault();
+      event.stopPropagation();
+
+      if (edge.type === 'sequenceMessage') {
+        const edgeData = edge.data as SequenceMessageData;
+        setEditingEdge(edge.id);
+        setEditingNode(null);
+        setEditingLabel(edgeData?.label || '');
+        setEditingType('edge');
+      }
+    },
+    [],
+  );
+
+  // Handle label input change
+  const handleLabelChange = useCallback(
+    (e: React.ChangeEvent<HTMLInputElement>) => {
+      setEditingLabel(e.target.value);
+    },
+    [],
+  );
+
+  // Handle condition change
+  const handleConditionChange = (idx: number, value: string) => {
+    const updated = [...editingConditions];
+    updated[idx] = value;
+    setEditingConditions(updated);
+  };
+
+  // Handle add condition
+  const handleAddCondition = () => {
+    setEditingConditions([...editingConditions, '']);
+  };
+
+  // Handle remove condition
+  const handleRemoveCondition = (idx: number) => {
+    setEditingConditions(editingConditions.filter((_, i) => i !== idx));
+  };
+
+  // Save the edited label
+  const saveLabel = useCallback(() => {
+    if (editingNode && editingType === 'node') {
+      setNodes((nds) =>
+        nds.map((node) => {
+          if (node.id === editingNode) {
+            return {
+              ...node,
+              data: {
+                ...node.data,
+                label: editingLabel,
+              },
+            };
+          }
+          return node;
+        }),
+      );
+      setEditingNode(null);
+    } else if (editingNode && editingType === 'fragment') {
+      setNodes((nds) =>
+        nds.map((node) => {
+          if (node.id === editingNode) {
+            return {
+              ...node,
+              data: {
+                ...node.data,
+                label: editingLabel,
+                conditions: editingConditions,
+              },
+            };
+          }
+          return node;
+        }),
+      );
+      setEditingNode(null);
+      setEditingType('node');
+    } else if (editingEdge && editingType === 'edge') {
+      setEdges((eds) =>
+        eds.map((edge) => {
+          if (edge.id === editingEdge) {
+            return {
+              ...edge,
+              data: {
+                ...edge.data,
+                label: editingLabel,
+              },
+            };
+          }
+          return edge;
+        }),
+      );
+      setEditingEdge(null);
+    }
+
+    // Reset editing state
+    setEditingType('node');
+  }, [
+    editingNode,
+    editingEdge,
+    editingType,
+    editingLabel,
+    editingConditions,
+    setNodes,
+    setEdges,
+  ]);
+
+  // Handle keyboard events for the input field
+  const handleKeyDown = useCallback(
+    (e: React.KeyboardEvent) => {
+      if (e.key === 'Enter') {
+        saveLabel();
+      } else if (e.key === 'Escape') {
+        setEditingNode(null);
+        setEditingEdge(null);
+      }
+    },
+    [saveLabel],
+  );
+
+  // Cancel editing
+  const cancelEditing = useCallback(() => {
+    setEditingNode(null);
+    setEditingEdge(null);
+    setEditingType('node');
+  }, []);
+
+  // Handler to add a new alt fragment
+  const handleAddAltFragment = useCallback(
+    (fragmentType: CombinedFragmentType = 'alt') => {
+      console.log(`Adding new ${fragmentType} fragment...`);
+
+      // Find all fragment nodes
+      const fragmentNodes = getCombinedFragmentNodes();
+      const fragmentCount = fragmentNodes.length;
+      const newFragmentId = `${fragmentType}-fragment-${fragmentCount + 1}`;
+
+      // Find the lowest y position among fragments
+      const defaultX = 250; // Fixed x position for all fragments
+      const defaultY = 300; // Default starting y for fragments
+
+      const lowestY =
+        fragmentNodes.length > 0
+          ? Math.max(...fragmentNodes.map((node) => node.position.y))
+          : defaultY;
+
+      const yOffset = 200; // Space between fragments vertically
+      const centerX = defaultX;
+      const centerY = fragmentNodes.length > 0 ? lowestY + yOffset : lowestY;
+
+      // Set up conditions based on fragment type
+      let conditions: string[] = [];
+      switch (fragmentType) {
+        case 'alt':
+          conditions = ['[condition]', '[else]'];
+          break;
+        case 'opt':
+          conditions = ['[condition]'];
+          break;
+        case 'loop':
+          conditions = ['[loop condition]'];
+          break;
+        case 'par':
+          conditions = ['[parallel 1]', '[parallel 2]'];
+          break;
+        default:
+          conditions = ['[condition]'];
+      }
+
+      // Create a new fragment node
+      const newFragment: Node<CombinedFragmentData> = {
+        id: newFragmentId,
+        type: 'combinedFragment',
+        position: { x: centerX, y: centerY },
+        style: {
+          width: fragmentSettings.width,
+          height: fragmentSettings.height,
+        },
+        data: {
+          label: fragmentType,
+          conditions: conditions,
+          style: {
+            width: fragmentSettings.width,
+            height: fragmentSettings.height,
+          },
+          color: '#000000',
+          backgroundColor: '#ffffff',
+          lifelineLength: 0,
+        },
+        draggable: true,
+        zIndex: 0, // Lower z-index to appear behind other elements
+      };
+
+      setNodes((nds) => [...nds, newFragment]);
+      setSelectedFragmentId(newFragmentId);
+      setShowFragmentSettings(true);
+    },
+    [fragmentSettings, getCombinedFragmentNodes, setNodes],
+  );
+
+  // Function to update selected fragment
+  const updateSelectedFragment = useCallback(
+    (fragmentId: string, updates: FragmentSettingsUpdateData) => {
+      setNodes((nds) =>
+        nds.map((node) => {
+          if (node.id === fragmentId) {
+            // Safely handle the style spread by ensuring the fragment data has a style property
+            const fragmentData = node.data;
+            const currentStyle = fragmentData.style || {};
+            // Apply updates to both style and data.style to ensure consistency
+            const updatedNode = {
+              ...node,
+              style: { ...node.style, ...updates },
+              data: {
+                ...fragmentData,
+                style: { ...currentStyle, ...updates },
+              },
+            };
+            return updatedNode;
+          }
+          return node;
+        }),
+      );
+    },
+    [setNodes],
+  );
+
+  const handleFragmentWidthChange = useCallback(
+    (width: number) => {
+      setFragmentSettings((prev) => ({ ...prev, width }));
+
+      if (selectedFragmentId) {
+        updateSelectedFragment(selectedFragmentId, { width });
+      }
+    },
+    [selectedFragmentId, updateSelectedFragment],
+  );
+
+  const handleFragmentHeightChange = useCallback(
+    (height: number) => {
+      setFragmentSettings((prev) => ({ ...prev, height }));
+
+      if (selectedFragmentId) {
+        updateSelectedFragment(selectedFragmentId, { height });
+      }
+    },
+    [selectedFragmentId, updateSelectedFragment],
+  );
+
+  const closeFragmentSettings = useCallback(() => {
+    setShowFragmentSettings(false);
+    setSelectedFragmentId(null);
+  }, []);
+
+  // Function to update selected entity
+  const updateSelectedEntity = useCallback(
+    (updates: EntitySettingsUpdateData) => {
+      if (!selectedEntityId) return;
+
+      setNodes((nds) =>
+        nds.map((node) => {
+          if (node.id === selectedEntityId) {
+            return {
+              ...node,
+              data: {
+                ...node.data,
+                ...updates,
+              },
+            };
+          }
+          return node;
+        }),
+      );
+    },
+    [selectedEntityId, setNodes],
+  );
+
+  // Close entity settings
+  const closeEntitySettings = useCallback(() => {
+    setShowEntitySettings(false);
+    setSelectedEntityId(null);
+  }, []);
+
+  // Handler to load a template diagram
+  const loadTemplateData = useCallback(
+    (templateType: string) => {
+      console.log(`Loading template: ${templateType}`);
+
+      const template = getTemplateData(templateType);
+
+      if (template.nodes.length > 0) {
+        setNodes(template.nodes);
+        setEdges(template.edges);
+      }
+    },
+    [setNodes, setEdges],
+  );
+
+  return (
+    <>
+      <Box sx={{ display: 'flex', width: '100%', height: '100vh' }}>
+        <Sidebar>
+          <DiagramControls
+            onAddEntity={handleAddEntity}
+            onAddMessage={handleAddMessage}
+            onAddAltFragment={handleAddAltFragment}
+            onUndo={handleUndo}
+            onRedo={handleRedo}
+            canUndo={historyIndex > 0}
+            canRedo={historyIndex < history.length - 1}
+            onLoadTemplate={loadTemplateData}
+          />
+        </Sidebar>
+        <Box sx={{ flex: 1, height: '100%', position: 'relative' }}>
+          <div
+            ref={diagramContainerRef}
+            className={`diagram-container ${isFullscreen ? 'fullscreen' : ''}`}
+            style={{
+              width: '100%',
+              height: '100%',
+              border: '1px solid #333',
+              position: 'relative',
+              backgroundColor: '#1a1a1a',
+              color: '#ffffff',
+            }}
+          >
+            <ReactFlow<Node, Edge>
+              nodes={nodes}
+              edges={edges}
+              onNodesChange={onNodesChangeLockedY}
+              onEdgesChange={onEdgesChange}
+              onConnect={onConnect}
+              nodeTypes={sequenceNodeTypes}
+              edgeTypes={sequenceEdgeTypes}
+              fitView
+              attributionPosition="bottom-left"
+              connectOnClick={false}
+              onNodeClick={onNodeClick}
+              onNodeDoubleClick={onNodeDoubleClick}
+              onEdgeDoubleClick={onEdgeDoubleClick}
+              style={{
+                backgroundColor: '#1a1a1a',
+              }}
+            >
+              <Background
+                color="#333"
+                variant={BackgroundVariant.Dots}
+                gap={16}
+              />
+            </ReactFlow>
+
+            {/* Components outside ReactFlow */}
+            <DiagramPanel
+              historyIndex={historyIndex}
+              historyLength={history.length}
+              onUndo={handleUndo}
+              onRedo={handleRedo}
+              onClear={() => {
+                const clearedState: SequenceHistoryEntry = { nodes: [], edges: [] };
+                const newHistory = history.slice(0, historyIndex + 1);
+                newHistory.push(clearedState);
+                setNodes([]);
+                setEdges([]);
+                setHistory(newHistory);
+                setHistoryIndex(newHistory.length - 1);
+              }}
+              onImport={() => { }}
+              isSettingsOpen={false}
+              onToggleSettings={() => { }}
+            />
+
+            <AiButton />
+
+            <MiniMapCommon />
+
+            {/* Fragment settings panel */}
+            {showFragmentSettings && (
+              <FragmentSettings
+                fragmentId={selectedFragmentId}
+                width={fragmentSettings.width}
+                height={fragmentSettings.height}
+                onWidthChange={handleFragmentWidthChange}
+                onHeightChange={handleFragmentHeightChange}
+                onClose={closeFragmentSettings}
+              />
+            )}
+
+            {/* Entity settings panel */}
+            {showEntitySettings && selectedEntityId && (
+              <EntitySettings
+                entityId={selectedEntityId}
+                data={
+                  (nodes.find((node) =>
+                    node.id === selectedEntityId && node.type === 'sequenceEntity'
+                  )?.data as SequenceEntityData) || {
+                    label: '',
+                    color: '#61dafb',
+                    backgroundColor: '#2a2a2a',
+                    lifelineLength: 500
+                  }
+                }
+                onUpdate={updateSelectedEntity}
+                onClose={closeEntitySettings}
+              />
+            )}
+
+            {/* Label editing overlay */}
+            {(editingNode || editingEdge) && (
+              <Box
+                sx={{
+                  position: 'absolute',
+                  zIndex: 10,
+                  top: 10,
+                  left: '50%',
+                  transform: 'translateX(-50%)',
+                  bgcolor: '#242424',
+                  border: '1px solid #333',
+                  color: '#fff',
+                  p: 2,
+                  borderRadius: 1,
+                  boxShadow: '0 4px 12px rgba(0, 0, 0, 0.3)',
+                  minWidth: 320,
+                  '& .MuiTextField-root': {
+                    '& .MuiInputLabel-root': {
+                      color: '#999',
+                    },
+                    '& .MuiOutlinedInput-root': {
+                      color: '#fff',
+                      '& fieldset': {
+                        borderColor: '#444',
+                      },
+                      '&:hover fieldset': {
+                        borderColor: '#61dafb',
+                      },
+                      '&.Mui-focused fieldset': {
+                        borderColor: '#61dafb',
+                      },
+                    },
+                  },
+                  '& .MuiButton-contained': {
+                    bgcolor: '#61dafb',
+                    color: '#000',
+                    '&:hover': {
+                      bgcolor: '#4db6d6',
+                    },
+                  },
+                  '& .MuiButton-outlined': {
+                    borderColor: '#333',
+                    color: '#fff',
+                    '&:hover': {
+                      borderColor: '#61dafb',
+                      bgcolor: 'rgba(97, 218, 251, 0.1)',
+                    },
+                  },
+                  '& .MuiIconButton-root': {
+                    color: '#999',
+                    '&:hover': {
+                      color: '#fff',
+                      bgcolor: 'rgba(255, 255, 255, 0.1)',
+                    },
+                  },
+                }}
+              >
+                <Typography sx={{ mb: 1, fontWeight: 'bold', color: '#fff' }}>
+                  {editingType === 'node'
+                    ? t('sequenceFlow.editLabels.entityLabel')
+                    : editingType === 'edge'
+                      ? t('sequenceFlow.editLabels.messageLabel')
+                      : t('sequenceFlow.editLabels.editFragment')}
+                </Typography>
+                {editingType === 'fragment' ? (
+                  <>
+                    <Box display="flex" alignItems="center" gap={2} mb={2}>
+                      <TextField
+                        type="text"
+                        value={editingLabel}
+                        onChange={handleLabelChange}
+                        onKeyDown={handleKeyDown}
+                        autoFocus
+                        size="small"
+                        sx={{ width: 200 }}
+                        label={t('fragmentSettings.fragmentName')}
+                      />
+                    </Box>
+                    {/* Conditions editing */}
+                    <Box mb={2}>
+                      <Typography
+                        variant="body2"
+                        fontWeight="bold"
+                        mb={0.5}
+                        sx={{ color: '#61dafb' }}
+                      >
+                        {t('fragmentSettings.conditions')}
+                      </Typography>
+                      {editingConditions.map((cond, idx) => (
+                        <Box
+                          key={idx}
+                          display="flex"
+                          alignItems="center"
+                          mb={1}
+                        >
+                          <TextField
+                            value={cond}
+                            onChange={(e) =>
+                              handleConditionChange(idx, e.target.value)
+                            }
+                            size="small"
+                            fullWidth
+                            sx={{ mr: 1 }}
+                          />
+                          <IconButton
+                            onClick={() => handleRemoveCondition(idx)}
+                            size="small"
+                          >
+                            <CloseIcon fontSize="small" />
+                          </IconButton>
+                        </Box>
+                      ))}
+                      <Button
+                        onClick={handleAddCondition}
+                        size="small"
+                        variant="outlined"
+                      >
+                        {t('fragmentSettings.addCondition')}
+                      </Button>
+                    </Box>
+                    <Box display="flex" gap={2}>
+                      <Button
+                        variant="contained"
+                        size="small"
+                        onClick={saveLabel}
+                        sx={{ textTransform: 'none' }}
+                      >
+                        {t('common.save')}
+                      </Button>
+                      <Button
+                        variant="outlined"
+                        size="small"
+                        onClick={cancelEditing}
+                        sx={{ textTransform: 'none' }}
+                      >
+                        {t('common.cancel')}
+                      </Button>
+                    </Box>
+                  </>
+                ) : (
+                  <Box display="flex" alignItems="center" gap={2}>
+                    <TextField
+                      type="text"
+                      value={editingLabel}
+                      onChange={handleLabelChange}
+                      onKeyDown={handleKeyDown}
+                      autoFocus
+                      size="small"
+                      sx={{ width: 200 }}
+                    />
+                    <Button
+                      variant="contained"
+                      size="small"
+                      onClick={saveLabel}
+                      sx={{ textTransform: 'none' }}
+                    >
+                      {t('common.save')}
+                    </Button>
+                    <Button
+                      variant="outlined"
+                      size="small"
+                      onClick={cancelEditing}
+                      sx={{ textTransform: 'none' }}
+                    >
+                      {t('common.cancel')}
+                    </Button>
+                  </Box>
+                )}
+              </Box>
+            )}
+
+            {/* Add Message Dialog */}
+            <AddMessageDialog
+              open={showMessageDialog}
+              onClose={() => setShowMessageDialog(false)}
+              onAddMessage={handleAddCustomMessage}
+              entities={sequenceEntities()}
+            />
+          </div>
+        </Box>
+      </Box>
+    </>
+  );
+};
+
+export default SequenceDiagramFlow;