@flowuent-org/diagramming-core 1.3.10 → 1.4.1

package/packages/diagrams/src/lib/templates/collaborationDiagram/components/CollaborationDiagramFlow.tsx

@@ -1,731 +1,731 @@
-import React, {
-  useCallback,
-  useState,
-  useRef,
-  useEffect,
-  useMemo,
-} from 'react';
-import {
-  ReactFlow,
-  useNodesState,
-  useEdgesState,
-  Node,
-  ConnectionLineType,
-  Connection,
-  Edge,
-  NodeTypes,
-  EdgeTypes,
-} from '@xyflow/react';
-import '@xyflow/react/dist/style.css';
-import CollabObjectNode from '../nodes/CollabObjectNode';
-import CollabEdge from '../nodes/CollabEdge';
-import {
-  Box,
-  Paper,
-  Button,
-  ToggleButton,
-  ToggleButtonGroup,
-} from '@mui/material';
-import CollaborationControls from '../controls/CollaborationControls';
-import CollabObjectNodeSettings from './CollabObjectNodeSettings';
-import AddCollabEdgeDialog from './AddCollabEdgeDialog';
-import { CollabEdgeVariant } from '../../../types/collaboration-types';
-import UndoIcon from '@mui/icons-material/Undo';
-import RedoIcon from '@mui/icons-material/Redo';
-import DownloadIcon from '@mui/icons-material/Download';
-import UploadIcon from '@mui/icons-material/Upload';
-import ViewWeekIcon from '@mui/icons-material/ViewWeek';
-import ViewStreamIcon from '@mui/icons-material/ViewStream';
-import { useCollabLayout } from './CollabLayout';
-import ActorNode from '../../../atoms/ActorNode';
-import ActorNodeSettings from '../../useCaseDiagram/components/ActorNodeSettings';
-import {
-  CollabObjectNodeData,
-  CollabActorNodeData,
-  CollaborationNode,
-  CollaborationEdge,
-  CollaborationDiagramProps,
-  CollabEdgeData
-} from '../../../types/collaboration-types';
-import { getCollabObjectNodeSize } from './CollabLayout';
-import Sidebar from '../../../components/SideBar';
-import DiagramPanel from '../../../components/DiagramPanel';
-import AiButton from '../../../components/AiButton';
-import MiniMapCommon from '../../../components/MiniMapCommon';
-
-const nodeTypes: NodeTypes = {
-  collabObject: CollabObjectNode,
-  actor: ActorNode,
-};
-
-const edgeTypes: EdgeTypes = {
-  collabEdge: CollabEdge,
-};
-
-interface CollaborationDiagramFlowProps {
-  collabNodes: Node[];
-  collabEdges: Edge[];
-}
-
-const getRandomPosition = () => ({
-  x: 100 + Math.random() * 400,
-  y: 100 + Math.random() * 300,
-});
-
-// Utility: Deep compare node arrays by id, type, data, width, height
-function 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
-function 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;
-}
-
-// Utility: Debounce function with proper typing
-function debounce<T extends (...args: never[]) => void>(fn: T, delay: number) {
-  let timer: NodeJS.Timeout | null = null;
-  return (...args: Parameters<T>) => {
-    if (timer) clearTimeout(timer);
-    timer = setTimeout(() => fn(...args), delay);
-  };
-}
-
-// Utility to estimate width/height for actor node based on label
-function getActorNodeSize(label: string) {
-  const lines = label.split('\n');
-  const maxLineLength = Math.max(...lines.map((l) => l.length));
-  // Estimate width: 8px per char + padding, clamp to min/max (min 40, max 200)
-  const width = Math.max(40, Math.min(200, maxLineLength * 8 + 32));
-  // Height: 60px base + 18px per extra line
-  const height = 60 + (lines.length - 1) * 18;
-  return { width, height };
-}
-
-const CollaborationDiagramFlow: React.FC<CollaborationDiagramFlowProps> = ({
-  collabNodes,
-  collabEdges,
-}) => {
-  const [nodes, setNodes, onNodesChange] = useNodesState(collabNodes);
-  const [edges, setEdges, onEdgesChange] = useEdgesState(collabEdges);
-  const [settingsOpen, setSettingsOpen] = useState(false);
-  const [selectedNode, setSelectedNode] =
-    useState<Node<CollabObjectNodeData | CollabActorNodeData> | null>(null);
-  const [edgeDialogOpen, setEdgeDialogOpen] = useState(false);
-  const [pendingConnection, setPendingConnection] = useState<{
-    source: string;
-    target: string;
-  } | null>(null);
-  // Undo/Redo/Import/Export state
-  const [history, setHistory] = useState<{ nodes: Node[]; edges: Edge[] }[]>([]);
-  const [historyIndex, setHistoryIndex] = useState<number>(-1);
-  const isUndoRedoAction = useRef(false);
-  // Debounced history update
-  const debouncedUpdateHistory = useRef(
-    debounce(
-      (
-        nodes: Node[],
-        edges: Edge[],
-        history: { nodes: Node[]; edges: Edge[] }[],
-        historyIndex: number,
-        setHistory: React.Dispatch<React.SetStateAction<{ nodes: Node[]; edges: Edge[] }[]>>,
-        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;
-  // Initialize history with initial state
-  useEffect(() => {
-    if (history.length === 0) {
-      setHistory([{ nodes: [...nodes], edges: [...edges] }]);
-      setHistoryIndex(0);
-    }
-  }, []); // Only run on mount
-  // Debounced save to history on nodes/edges change
-  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]);
-  // Undo handler
-  const handleUndo = useCallback(() => {
-    if (historyIndex > 0) {
-      isUndoRedoAction.current = true;
-      const prev = history[historyIndex - 1];
-      setNodes(prev.nodes);
-      setEdges(prev.edges);
-      setHistoryIndex(historyIndex - 1);
-      setTimeout(() => {
-        isUndoRedoAction.current = false;
-      }, 0);
-    }
-  }, [history, historyIndex, setNodes, setEdges]);
-  // Redo handler
-  const handleRedo = useCallback(() => {
-    if (historyIndex < history.length - 1) {
-      isUndoRedoAction.current = true;
-      const next = history[historyIndex + 1];
-      setNodes(next.nodes);
-      setEdges(next.edges);
-      setHistoryIndex(historyIndex + 1);
-      setTimeout(() => {
-        isUndoRedoAction.current = false;
-      }, 0);
-    }
-  }, [history, historyIndex, setNodes, setEdges]);
-  // Keyboard shortcuts for undo/redo
-  useEffect(() => {
-    const handleKeyDown = (e: KeyboardEvent) => {
-      // Ctrl+Z or Cmd+Z (undo)
-      if ((e.ctrlKey || e.metaKey) && e.key === 'z' && !e.shiftKey) {
-        e.preventDefault();
-        handleUndo();
-      }
-      // 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]);
-  // 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]);
-  // Export/import diagram
-  const exportDiagram = useCallback(() => {
-    const diagramData = { nodes, edges };
-    const dataStr = JSON.stringify(diagramData, null, 2);
-    const dataUri =
-      'data:application/json;charset=utf-8,' + encodeURIComponent(dataStr);
-    const exportFileDefaultName = 'collaboration-diagram.json';
-    const linkElement = document.createElement('a');
-    linkElement.setAttribute('href', dataUri);
-    linkElement.setAttribute('download', exportFileDefaultName);
-    linkElement.click();
-  }, [nodes, edges]);
-  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],
-  );
-
-  // Handlers for adding nodes
-  const handleAddObject = useCallback(() => {
-    const label = `Object ${nodes.filter((n) => n.type === 'collabObject' && n.data.variant !== 'multiple').length + 1}`;
-    const { width, height } = getCollabObjectNodeSize(label);
-    const newNode: Node = {
-      id: `collabObject-${Date.now()}-${Math.floor(Math.random() * 10000)}`,
-      type: 'collabObject',
-      position: getRandomPosition(),
-      data: {
-        label,
-        variant: 'normal',
-        width,
-        height,
-        infiniteResize: true,
-      },
-      width,
-      height,
-      style: {
-        width,
-        height,
-      },
-    };
-    setNodes((nds) => [...nds, newNode]);
-  }, [nodes, setNodes]);
-
-  const handleAddMultipleObjects = useCallback(() => {
-    const label = `Objects ${nodes.filter((n) => n.type === 'collabObject' && n.data.variant === 'multiple').length + 1}`;
-    const { width, height } = getCollabObjectNodeSize(label);
-    const newNode: Node = {
-      id: `collabObject-${Date.now()}-${Math.floor(Math.random() * 10000)}`,
-      type: 'collabObject',
-      position: getRandomPosition(),
-      data: {
-        label,
-        variant: 'multiple',
-        width,
-        height,
-        infiniteResize: true,
-      },
-      width,
-      height,
-      style: {
-        width,
-        height,
-      },
-    };
-    setNodes((nds) => [...nds, newNode]);
-  }, [nodes, setNodes]);
-
-  const handleAddActor = useCallback(() => {
-    const label = `Actor ${nodes.filter((n) => n.type === 'actor').length + 1}`;
-    const { width, height } = getActorNodeSize(label);
-    const newNode: Node = {
-      id: `actor-${Date.now()}-${Math.floor(Math.random() * 10000)}`,
-      type: 'actor',
-      position: getRandomPosition(),
-      data: {
-        label,
-        width,
-        height,
-        infiniteResize: true,
-      },
-      width,
-      height,
-      style: {
-        width,
-        height,
-      },
-    };
-    setNodes((nds) => [...nds, newNode]);
-  }, [nodes, setNodes]);
-
-  // Helper to calculate parallel edges info
-  const getParallelEdgeInfo = (edgesArr: Edge[], edge: Edge) => {
-    // Find all edges between the same two nodes (in either direction)
-    const parallelEdges = edgesArr.filter(
-      (e: Edge) =>
-        (e.source === edge.source && e.target === edge.target) ||
-        (e.source === edge.target && e.target === edge.source),
-    );
-    // Sort for consistent ordering
-    parallelEdges.sort((a: Edge, b: Edge) => a.id.localeCompare(b.id));
-    const parallelCount = parallelEdges.length;
-    const parallelIndex = parallelEdges.findIndex((e: Edge) => e.id === edge.id);
-    return { parallelIndex, parallelCount };
-  };
-
-  // Handlers for adding edges
-  const handleAddEdge = useCallback(
-    (
-      sourceId: string,
-      targetId: string,
-      variant: CollabEdgeVariant,
-      label?: string,
-      arrowHeadSource?: 'none' | 'v',
-      arrowHeadTarget?: 'none' | 'v',
-    ) => {
-      setEdges((eds) => [
-        ...eds,
-        {
-          id: `collabEdge-${variant}-${eds.length + 1}-${Date.now()}`,
-          source: sourceId,
-          target: targetId,
-          type: 'collabEdge',
-          data: { variant, label, arrowHeadSource, arrowHeadTarget },
-        },
-      ]);
-    },
-    [setEdges],
-  );
-
-  // Modified onConnect handler to open dialog
-  const onConnect = useCallback((params: Connection) => {
-    if (params.source && params.target) {
-      setPendingConnection({ source: params.source, target: params.target });
-      setEdgeDialogOpen(true);
-    }
-  }, []);
-
-  // Handler to open dialog from controls
-  const handleOpenAddEdgeDialog = useCallback(() => {
-    // Default to first two nodes if available
-    setEditingEdge(null);
-    if (nodes.length >= 2) {
-      setPendingConnection({ source: nodes[0].id, target: nodes[1].id });
-    } else if (nodes.length === 1) {
-      setPendingConnection({ source: nodes[0].id, target: nodes[0].id });
-    } else {
-      setPendingConnection(null);
-    }
-    setEdgeDialogOpen(true);
-  }, [nodes]);
-
-  // Handler for editing an existing message (edge)
-  const handleEditMessage = useCallback((edge: Edge) => {
-    setEditingEdge(edge as Edge<CollabEdgeData>);
-    setEdgeDialogOpen(true);
-  }, []);
-
-  // Edge click handler for ReactFlow
-  const handleEdgeClick = useCallback(
-    (event: React.MouseEvent, edge: Edge) => {
-      event.stopPropagation(); // Prevents deselecting
-      handleEditMessage(edge);
-    },
-    [handleEditMessage],
-  );
-
-  // Double-click handler to open settings for collabObject nodes
-  const handleNodeDoubleClick = useCallback(
-    (event: React.MouseEvent, node: Node) => {
-      if (node.type === 'collabObject' || node.type === 'actor') {
-        setSelectedNode(node as Node<CollabObjectNodeData | CollabActorNodeData>);
-        setSettingsOpen(true);
-      }
-    },
-    [],
-  );
-
-  // Update node label
-  const handleUpdateNode = useCallback(
-    (updates: Partial<CollabObjectNodeData | CollabActorNodeData>) => {
-      if (!selectedNode) return;
-      setNodes((nds) =>
-        nds.map((n) => {
-          if (n.id === selectedNode.id) {
-            // If label is being updated, recalculate dimensions
-            const newWidth = updates.label
-              ? n.type === 'collabObject'
-                ? getCollabObjectNodeSize(updates.label).width
-                : getActorNodeSize(updates.label).width
-              : n.width || 0;
-            const newHeight = updates.label
-              ? n.type === 'collabObject'
-                ? getCollabObjectNodeSize(updates.label).height
-                : getActorNodeSize(updates.label).height
-              : n.height || 0;
-
-            const updatedNode: Node = {
-              ...n,
-              data: {
-                ...n.data,
-                ...updates,
-                width: newWidth,
-                height: newHeight,
-              },
-              width: newWidth,
-              height: newHeight,
-              style: {
-                ...n.style,
-                width: newWidth,
-                height: newHeight,
-              },
-            };
-            return updatedNode;
-          }
-          return n;
-        }),
-      );
-    },
-    [selectedNode, setNodes],
-  );
-
-  // Close settings panel
-  const handleCloseSettings = useCallback(() => {
-    setSettingsOpen(false);
-    setSelectedNode(null);
-  }, []);
-
-  // Add this handler for self-message toggle
-  const handleSelfMessageChange = useCallback(
-    (checked: boolean) => {
-      if (!selectedNode) return;
-      if (checked) {
-        // Add self-message edge if not already present
-        setEdges((eds) => {
-          const exists = eds.some(
-            (e) =>
-              e.source === selectedNode.id &&
-              e.target === selectedNode.id &&
-              e.data?.variant === 'self-message',
-          );
-          if (exists) return eds;
-          return [
-            ...eds,
-            {
-              id: `collabEdge-self-message-${eds.length + 1}-${Date.now()}`,
-              source: selectedNode.id,
-              target: selectedNode.id,
-              type: 'collabEdge',
-              data: { variant: 'self-message' },
-            },
-          ];
-        });
-      } else {
-        // Remove self-message edge for this node
-        setEdges((eds) =>
-          eds.filter(
-            (e) =>
-              !(
-                e.source === selectedNode.id &&
-                e.target === selectedNode.id &&
-                e.data?.variant === 'self-message'
-              ),
-          ),
-        );
-      }
-    },
-    [selectedNode, setEdges],
-  );
-
-  const [editingEdge, setEditingEdge] = useState<Edge<CollabEdgeData> | null>(null);
-
-  // --- Layout integration ---
-  const { getLayoutedElements, layoutDirection, setLayoutDirection } =
-    useCollabLayout(nodes as CollaborationNode[], edges as CollaborationEdge[], setNodes);
-  // Manual layout trigger
-  const handleManualLayout = useCallback(() => {
-    getLayoutedElements();
-  }, [getLayoutedElements]);
-  // Layout direction change
-  const handleLayoutDirectionChange = (direction: 'DOWN' | 'RIGHT') => {
-    setLayoutDirection(direction);
-    getLayoutedElements();
-  };
-
-  // Memoize mapped edges for ReactFlow to avoid unnecessary re-renders
-  const mappedEdges = React.useMemo(() => {
-    return edges.map((e: Edge) => {
-      const { parallelIndex, parallelCount } = getParallelEdgeInfo(edges, e);
-      return {
-        ...e,
-        data: { ...e.data, parallelIndex, parallelCount, type: 'collabEdge' },
-      };
-    });
-  }, [edges]);
-
-  return (
-    <Box sx={{ display: 'flex', width: '100%', height: '100vh' }}>
-      <Sidebar>
-        <CollaborationControls
-          onAddObject={handleAddObject}
-          onAddMultipleObjects={handleAddMultipleObjects}
-          onAddActor={handleAddActor}
-          onAddMessage={handleOpenAddEdgeDialog}
-          nodes={nodes as CollaborationNode[]}
-          edges={edges as CollaborationEdge[]}
-          onEditMessage={handleEditMessage}
-        />
-      </Sidebar>
-      <Box sx={{ flex: 1, height: '100%', position: 'relative' }}>
-        <ReactFlow
-          nodes={nodes}
-          edges={mappedEdges}
-          onNodesChange={onNodesChange}
-          onEdgesChange={onEdgesChange}
-          onConnect={onConnect}
-          nodeTypes={nodeTypes}
-          edgeTypes={edgeTypes}
-          fitView
-          connectionLineType={ConnectionLineType.Straight}
-          onNodeDoubleClick={handleNodeDoubleClick}
-          onEdgeDoubleClick={handleEdgeClick}
-          style={{
-            background: '#1a1a1a',
-          }}
-        />
-        <MiniMapCommon />
-        {selectedNode?.type === 'actor' && (
-          <ActorNodeSettings
-            open={settingsOpen}
-            node={selectedNode as Node<CollabActorNodeData>}
-            onUpdate={handleUpdateNode}
-            onClose={handleCloseSettings}
-            onSelfMessageChange={handleSelfMessageChange}
-            hasSelfMessage={
-              !!(
-                selectedNode &&
-                edges.some(
-                  (e) =>
-                    e.source === selectedNode.id &&
-                    e.target === selectedNode.id &&
-                    e.data?.variant === 'self-message',
-                )
-              )
-            }
-          />
-        )}
-        {selectedNode?.type === 'collabObject' && (
-          <CollabObjectNodeSettings
-            open={settingsOpen}
-            node={selectedNode}
-            onUpdate={handleUpdateNode}
-            onClose={handleCloseSettings}
-            onSelfMessageChange={handleSelfMessageChange}
-            hasSelfMessage={
-              !!(
-                selectedNode &&
-                edges.some(
-                  (e) =>
-                    e.source === selectedNode.id &&
-                    e.target === selectedNode.id &&
-                    e.data?.variant === 'self-message',
-                )
-              )
-            }
-          />
-        )}
-        <AddCollabEdgeDialog
-          open={edgeDialogOpen}
-          onClose={() => {
-            setEdgeDialogOpen(false);
-            setEditingEdge(null);
-            setPendingConnection(null);
-          }}
-          onAddEdge={(
-            sourceId,
-            targetId,
-            variant,
-            label,
-            arrowHeadSource,
-            arrowHeadTarget,
-          ) => {
-            if (editingEdge) {
-              // Edit existing edge
-              setEdges((eds) =>
-                eds.map((e) =>
-                  e.id === editingEdge.id
-                    ? {
-                      ...e,
-                      source: sourceId,
-                      target: targetId,
-                      data: {
-                        ...e.data,
-                        variant,
-                        label,
-                        arrowHeadSource,
-                        arrowHeadTarget,
-                      },
-                    }
-                    : e,
-                ),
-              );
-            } else {
-              // Add new edge
-              handleAddEdge(
-                sourceId,
-                targetId,
-                variant,
-                label,
-                arrowHeadSource,
-                arrowHeadTarget,
-              );
-            }
-            setEdgeDialogOpen(false);
-            setEditingEdge(null);
-            setPendingConnection(null);
-          }}
-          nodes={nodes as CollaborationNode[]}
-          initialSourceId={pendingConnection?.source}
-          initialTargetId={pendingConnection?.target}
-          editingEdge={editingEdge}
-          onDeleteEdge={
-            editingEdge
-              ? () => {
-                setEdges((eds) => eds.filter((e) => e.id !== editingEdge.id));
-                setEdgeDialogOpen(false);
-                setEditingEdge(null);
-                setPendingConnection(null);
-              }
-              : undefined
-          }
-        />
-        {/* Bottom-right panel for controls */}
-        <DiagramPanel
-          historyIndex={historyIndex}
-          historyLength={history.length}
-          onUndo={handleUndo}
-          onRedo={handleRedo}
-          onClear={handleClear}
-          onImport={() => { }}
-          isSettingsOpen={false}
-          onToggleSettings={() => { }}
-          showCustomControls={true}
-          onGridLayout={() => handleLayoutDirectionChange('DOWN')}
-          onTreeLayout={() => handleLayoutDirectionChange('RIGHT')}
-        />
-        <AiButton />
-
-      </Box>
-    </Box>
-  );
-};
-
-export default CollaborationDiagramFlow;
+import React, {
+  useCallback,
+  useState,
+  useRef,
+  useEffect,
+  useMemo,
+} from 'react';
+import {
+  ReactFlow,
+  useNodesState,
+  useEdgesState,
+  Node,
+  ConnectionLineType,
+  Connection,
+  Edge,
+  NodeTypes,
+  EdgeTypes,
+} from '@xyflow/react';
+import '@xyflow/react/dist/style.css';
+import CollabObjectNode from '../nodes/CollabObjectNode';
+import CollabEdge from '../nodes/CollabEdge';
+import {
+  Box,
+  Paper,
+  Button,
+  ToggleButton,
+  ToggleButtonGroup,
+} from '@mui/material';
+import CollaborationControls from '../controls/CollaborationControls';
+import CollabObjectNodeSettings from './CollabObjectNodeSettings';
+import AddCollabEdgeDialog from './AddCollabEdgeDialog';
+import { CollabEdgeVariant } from '../../../types/collaboration-types';
+import UndoIcon from '@mui/icons-material/Undo';
+import RedoIcon from '@mui/icons-material/Redo';
+import DownloadIcon from '@mui/icons-material/Download';
+import UploadIcon from '@mui/icons-material/Upload';
+import ViewWeekIcon from '@mui/icons-material/ViewWeek';
+import ViewStreamIcon from '@mui/icons-material/ViewStream';
+import { useCollabLayout } from './CollabLayout';
+import ActorNode from '../../../atoms/ActorNode';
+import ActorNodeSettings from '../../useCaseDiagram/components/ActorNodeSettings';
+import {
+  CollabObjectNodeData,
+  CollabActorNodeData,
+  CollaborationNode,
+  CollaborationEdge,
+  CollaborationDiagramProps,
+  CollabEdgeData
+} from '../../../types/collaboration-types';
+import { getCollabObjectNodeSize } from './CollabLayout';
+import Sidebar from '../../../components/SideBar';
+import DiagramPanel from '../../../components/DiagramPanel';
+import AiButton from '../../../components/AiButton';
+import MiniMapCommon from '../../../components/MiniMapCommon';
+
+const nodeTypes: NodeTypes = {
+  collabObject: CollabObjectNode,
+  actor: ActorNode,
+};
+
+const edgeTypes: EdgeTypes = {
+  collabEdge: CollabEdge,
+};
+
+interface CollaborationDiagramFlowProps {
+  collabNodes: Node[];
+  collabEdges: Edge[];
+}
+
+const getRandomPosition = () => ({
+  x: 100 + Math.random() * 400,
+  y: 100 + Math.random() * 300,
+});
+
+// Utility: Deep compare node arrays by id, type, data, width, height
+function 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
+function 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;
+}
+
+// Utility: Debounce function with proper typing
+function debounce<T extends (...args: never[]) => void>(fn: T, delay: number) {
+  let timer: NodeJS.Timeout | null = null;
+  return (...args: Parameters<T>) => {
+    if (timer) clearTimeout(timer);
+    timer = setTimeout(() => fn(...args), delay);
+  };
+}
+
+// Utility to estimate width/height for actor node based on label
+function getActorNodeSize(label: string) {
+  const lines = label.split('\n');
+  const maxLineLength = Math.max(...lines.map((l) => l.length));
+  // Estimate width: 8px per char + padding, clamp to min/max (min 40, max 200)
+  const width = Math.max(40, Math.min(200, maxLineLength * 8 + 32));
+  // Height: 60px base + 18px per extra line
+  const height = 60 + (lines.length - 1) * 18;
+  return { width, height };
+}
+
+const CollaborationDiagramFlow: React.FC<CollaborationDiagramFlowProps> = ({
+  collabNodes,
+  collabEdges,
+}) => {
+  const [nodes, setNodes, onNodesChange] = useNodesState(collabNodes);
+  const [edges, setEdges, onEdgesChange] = useEdgesState(collabEdges);
+  const [settingsOpen, setSettingsOpen] = useState(false);
+  const [selectedNode, setSelectedNode] =
+    useState<Node<CollabObjectNodeData | CollabActorNodeData> | null>(null);
+  const [edgeDialogOpen, setEdgeDialogOpen] = useState(false);
+  const [pendingConnection, setPendingConnection] = useState<{
+    source: string;
+    target: string;
+  } | null>(null);
+  // Undo/Redo/Import/Export state
+  const [history, setHistory] = useState<{ nodes: Node[]; edges: Edge[] }[]>([]);
+  const [historyIndex, setHistoryIndex] = useState<number>(-1);
+  const isUndoRedoAction = useRef(false);
+  // Debounced history update
+  const debouncedUpdateHistory = useRef(
+    debounce(
+      (
+        nodes: Node[],
+        edges: Edge[],
+        history: { nodes: Node[]; edges: Edge[] }[],
+        historyIndex: number,
+        setHistory: React.Dispatch<React.SetStateAction<{ nodes: Node[]; edges: Edge[] }[]>>,
+        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;
+  // Initialize history with initial state
+  useEffect(() => {
+    if (history.length === 0) {
+      setHistory([{ nodes: [...nodes], edges: [...edges] }]);
+      setHistoryIndex(0);
+    }
+  }, []); // Only run on mount
+  // Debounced save to history on nodes/edges change
+  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]);
+  // Undo handler
+  const handleUndo = useCallback(() => {
+    if (historyIndex > 0) {
+      isUndoRedoAction.current = true;
+      const prev = history[historyIndex - 1];
+      setNodes(prev.nodes);
+      setEdges(prev.edges);
+      setHistoryIndex(historyIndex - 1);
+      setTimeout(() => {
+        isUndoRedoAction.current = false;
+      }, 0);
+    }
+  }, [history, historyIndex, setNodes, setEdges]);
+  // Redo handler
+  const handleRedo = useCallback(() => {
+    if (historyIndex < history.length - 1) {
+      isUndoRedoAction.current = true;
+      const next = history[historyIndex + 1];
+      setNodes(next.nodes);
+      setEdges(next.edges);
+      setHistoryIndex(historyIndex + 1);
+      setTimeout(() => {
+        isUndoRedoAction.current = false;
+      }, 0);
+    }
+  }, [history, historyIndex, setNodes, setEdges]);
+  // Keyboard shortcuts for undo/redo
+  useEffect(() => {
+    const handleKeyDown = (e: KeyboardEvent) => {
+      // Ctrl+Z or Cmd+Z (undo)
+      if ((e.ctrlKey || e.metaKey) && e.key === 'z' && !e.shiftKey) {
+        e.preventDefault();
+        handleUndo();
+      }
+      // 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]);
+  // 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]);
+  // Export/import diagram
+  const exportDiagram = useCallback(() => {
+    const diagramData = { nodes, edges };
+    const dataStr = JSON.stringify(diagramData, null, 2);
+    const dataUri =
+      'data:application/json;charset=utf-8,' + encodeURIComponent(dataStr);
+    const exportFileDefaultName = 'collaboration-diagram.json';
+    const linkElement = document.createElement('a');
+    linkElement.setAttribute('href', dataUri);
+    linkElement.setAttribute('download', exportFileDefaultName);
+    linkElement.click();
+  }, [nodes, edges]);
+  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],
+  );
+
+  // Handlers for adding nodes
+  const handleAddObject = useCallback(() => {
+    const label = `Object ${nodes.filter((n) => n.type === 'collabObject' && n.data.variant !== 'multiple').length + 1}`;
+    const { width, height } = getCollabObjectNodeSize(label);
+    const newNode: Node = {
+      id: `collabObject-${Date.now()}-${Math.floor(Math.random() * 10000)}`,
+      type: 'collabObject',
+      position: getRandomPosition(),
+      data: {
+        label,
+        variant: 'normal',
+        width,
+        height,
+        infiniteResize: true,
+      },
+      width,
+      height,
+      style: {
+        width,
+        height,
+      },
+    };
+    setNodes((nds) => [...nds, newNode]);
+  }, [nodes, setNodes]);
+
+  const handleAddMultipleObjects = useCallback(() => {
+    const label = `Objects ${nodes.filter((n) => n.type === 'collabObject' && n.data.variant === 'multiple').length + 1}`;
+    const { width, height } = getCollabObjectNodeSize(label);
+    const newNode: Node = {
+      id: `collabObject-${Date.now()}-${Math.floor(Math.random() * 10000)}`,
+      type: 'collabObject',
+      position: getRandomPosition(),
+      data: {
+        label,
+        variant: 'multiple',
+        width,
+        height,
+        infiniteResize: true,
+      },
+      width,
+      height,
+      style: {
+        width,
+        height,
+      },
+    };
+    setNodes((nds) => [...nds, newNode]);
+  }, [nodes, setNodes]);
+
+  const handleAddActor = useCallback(() => {
+    const label = `Actor ${nodes.filter((n) => n.type === 'actor').length + 1}`;
+    const { width, height } = getActorNodeSize(label);
+    const newNode: Node = {
+      id: `actor-${Date.now()}-${Math.floor(Math.random() * 10000)}`,
+      type: 'actor',
+      position: getRandomPosition(),
+      data: {
+        label,
+        width,
+        height,
+        infiniteResize: true,
+      },
+      width,
+      height,
+      style: {
+        width,
+        height,
+      },
+    };
+    setNodes((nds) => [...nds, newNode]);
+  }, [nodes, setNodes]);
+
+  // Helper to calculate parallel edges info
+  const getParallelEdgeInfo = (edgesArr: Edge[], edge: Edge) => {
+    // Find all edges between the same two nodes (in either direction)
+    const parallelEdges = edgesArr.filter(
+      (e: Edge) =>
+        (e.source === edge.source && e.target === edge.target) ||
+        (e.source === edge.target && e.target === edge.source),
+    );
+    // Sort for consistent ordering
+    parallelEdges.sort((a: Edge, b: Edge) => a.id.localeCompare(b.id));
+    const parallelCount = parallelEdges.length;
+    const parallelIndex = parallelEdges.findIndex((e: Edge) => e.id === edge.id);
+    return { parallelIndex, parallelCount };
+  };
+
+  // Handlers for adding edges
+  const handleAddEdge = useCallback(
+    (
+      sourceId: string,
+      targetId: string,
+      variant: CollabEdgeVariant,
+      label?: string,
+      arrowHeadSource?: 'none' | 'v',
+      arrowHeadTarget?: 'none' | 'v',
+    ) => {
+      setEdges((eds) => [
+        ...eds,
+        {
+          id: `collabEdge-${variant}-${eds.length + 1}-${Date.now()}`,
+          source: sourceId,
+          target: targetId,
+          type: 'collabEdge',
+          data: { variant, label, arrowHeadSource, arrowHeadTarget },
+        },
+      ]);
+    },
+    [setEdges],
+  );
+
+  // Modified onConnect handler to open dialog
+  const onConnect = useCallback((params: Connection) => {
+    if (params.source && params.target) {
+      setPendingConnection({ source: params.source, target: params.target });
+      setEdgeDialogOpen(true);
+    }
+  }, []);
+
+  // Handler to open dialog from controls
+  const handleOpenAddEdgeDialog = useCallback(() => {
+    // Default to first two nodes if available
+    setEditingEdge(null);
+    if (nodes.length >= 2) {
+      setPendingConnection({ source: nodes[0].id, target: nodes[1].id });
+    } else if (nodes.length === 1) {
+      setPendingConnection({ source: nodes[0].id, target: nodes[0].id });
+    } else {
+      setPendingConnection(null);
+    }
+    setEdgeDialogOpen(true);
+  }, [nodes]);
+
+  // Handler for editing an existing message (edge)
+  const handleEditMessage = useCallback((edge: Edge) => {
+    setEditingEdge(edge as Edge<CollabEdgeData>);
+    setEdgeDialogOpen(true);
+  }, []);
+
+  // Edge click handler for ReactFlow
+  const handleEdgeClick = useCallback(
+    (event: React.MouseEvent, edge: Edge) => {
+      event.stopPropagation(); // Prevents deselecting
+      handleEditMessage(edge);
+    },
+    [handleEditMessage],
+  );
+
+  // Double-click handler to open settings for collabObject nodes
+  const handleNodeDoubleClick = useCallback(
+    (event: React.MouseEvent, node: Node) => {
+      if (node.type === 'collabObject' || node.type === 'actor') {
+        setSelectedNode(node as Node<CollabObjectNodeData | CollabActorNodeData>);
+        setSettingsOpen(true);
+      }
+    },
+    [],
+  );
+
+  // Update node label
+  const handleUpdateNode = useCallback(
+    (updates: Partial<CollabObjectNodeData | CollabActorNodeData>) => {
+      if (!selectedNode) return;
+      setNodes((nds) =>
+        nds.map((n) => {
+          if (n.id === selectedNode.id) {
+            // If label is being updated, recalculate dimensions
+            const newWidth = updates.label
+              ? n.type === 'collabObject'
+                ? getCollabObjectNodeSize(updates.label).width
+                : getActorNodeSize(updates.label).width
+              : n.width || 0;
+            const newHeight = updates.label
+              ? n.type === 'collabObject'
+                ? getCollabObjectNodeSize(updates.label).height
+                : getActorNodeSize(updates.label).height
+              : n.height || 0;
+
+            const updatedNode: Node = {
+              ...n,
+              data: {
+                ...n.data,
+                ...updates,
+                width: newWidth,
+                height: newHeight,
+              },
+              width: newWidth,
+              height: newHeight,
+              style: {
+                ...n.style,
+                width: newWidth,
+                height: newHeight,
+              },
+            };
+            return updatedNode;
+          }
+          return n;
+        }),
+      );
+    },
+    [selectedNode, setNodes],
+  );
+
+  // Close settings panel
+  const handleCloseSettings = useCallback(() => {
+    setSettingsOpen(false);
+    setSelectedNode(null);
+  }, []);
+
+  // Add this handler for self-message toggle
+  const handleSelfMessageChange = useCallback(
+    (checked: boolean) => {
+      if (!selectedNode) return;
+      if (checked) {
+        // Add self-message edge if not already present
+        setEdges((eds) => {
+          const exists = eds.some(
+            (e) =>
+              e.source === selectedNode.id &&
+              e.target === selectedNode.id &&
+              e.data?.variant === 'self-message',
+          );
+          if (exists) return eds;
+          return [
+            ...eds,
+            {
+              id: `collabEdge-self-message-${eds.length + 1}-${Date.now()}`,
+              source: selectedNode.id,
+              target: selectedNode.id,
+              type: 'collabEdge',
+              data: { variant: 'self-message' },
+            },
+          ];
+        });
+      } else {
+        // Remove self-message edge for this node
+        setEdges((eds) =>
+          eds.filter(
+            (e) =>
+              !(
+                e.source === selectedNode.id &&
+                e.target === selectedNode.id &&
+                e.data?.variant === 'self-message'
+              ),
+          ),
+        );
+      }
+    },
+    [selectedNode, setEdges],
+  );
+
+  const [editingEdge, setEditingEdge] = useState<Edge<CollabEdgeData> | null>(null);
+
+  // --- Layout integration ---
+  const { getLayoutedElements, layoutDirection, setLayoutDirection } =
+    useCollabLayout(nodes as CollaborationNode[], edges as CollaborationEdge[], setNodes);
+  // Manual layout trigger
+  const handleManualLayout = useCallback(() => {
+    getLayoutedElements();
+  }, [getLayoutedElements]);
+  // Layout direction change
+  const handleLayoutDirectionChange = (direction: 'DOWN' | 'RIGHT') => {
+    setLayoutDirection(direction);
+    getLayoutedElements();
+  };
+
+  // Memoize mapped edges for ReactFlow to avoid unnecessary re-renders
+  const mappedEdges = React.useMemo(() => {
+    return edges.map((e: Edge) => {
+      const { parallelIndex, parallelCount } = getParallelEdgeInfo(edges, e);
+      return {
+        ...e,
+        data: { ...e.data, parallelIndex, parallelCount, type: 'collabEdge' },
+      };
+    });
+  }, [edges]);
+
+  return (
+    <Box sx={{ display: 'flex', width: '100%', height: '100vh' }}>
+      <Sidebar>
+        <CollaborationControls
+          onAddObject={handleAddObject}
+          onAddMultipleObjects={handleAddMultipleObjects}
+          onAddActor={handleAddActor}
+          onAddMessage={handleOpenAddEdgeDialog}
+          nodes={nodes as CollaborationNode[]}
+          edges={edges as CollaborationEdge[]}
+          onEditMessage={handleEditMessage}
+        />
+      </Sidebar>
+      <Box sx={{ flex: 1, height: '100%', position: 'relative' }}>
+        <ReactFlow
+          nodes={nodes}
+          edges={mappedEdges}
+          onNodesChange={onNodesChange}
+          onEdgesChange={onEdgesChange}
+          onConnect={onConnect}
+          nodeTypes={nodeTypes}
+          edgeTypes={edgeTypes}
+          fitView
+          connectionLineType={ConnectionLineType.Straight}
+          onNodeDoubleClick={handleNodeDoubleClick}
+          onEdgeDoubleClick={handleEdgeClick}
+          style={{
+            background: '#1a1a1a',
+          }}
+        />
+        <MiniMapCommon />
+        {selectedNode?.type === 'actor' && (
+          <ActorNodeSettings
+            open={settingsOpen}
+            node={selectedNode as Node<CollabActorNodeData>}
+            onUpdate={handleUpdateNode}
+            onClose={handleCloseSettings}
+            onSelfMessageChange={handleSelfMessageChange}
+            hasSelfMessage={
+              !!(
+                selectedNode &&
+                edges.some(
+                  (e) =>
+                    e.source === selectedNode.id &&
+                    e.target === selectedNode.id &&
+                    e.data?.variant === 'self-message',
+                )
+              )
+            }
+          />
+        )}
+        {selectedNode?.type === 'collabObject' && (
+          <CollabObjectNodeSettings
+            open={settingsOpen}
+            node={selectedNode}
+            onUpdate={handleUpdateNode}
+            onClose={handleCloseSettings}
+            onSelfMessageChange={handleSelfMessageChange}
+            hasSelfMessage={
+              !!(
+                selectedNode &&
+                edges.some(
+                  (e) =>
+                    e.source === selectedNode.id &&
+                    e.target === selectedNode.id &&
+                    e.data?.variant === 'self-message',
+                )
+              )
+            }
+          />
+        )}
+        <AddCollabEdgeDialog
+          open={edgeDialogOpen}
+          onClose={() => {
+            setEdgeDialogOpen(false);
+            setEditingEdge(null);
+            setPendingConnection(null);
+          }}
+          onAddEdge={(
+            sourceId,
+            targetId,
+            variant,
+            label,
+            arrowHeadSource,
+            arrowHeadTarget,
+          ) => {
+            if (editingEdge) {
+              // Edit existing edge
+              setEdges((eds) =>
+                eds.map((e) =>
+                  e.id === editingEdge.id
+                    ? {
+                      ...e,
+                      source: sourceId,
+                      target: targetId,
+                      data: {
+                        ...e.data,
+                        variant,
+                        label,
+                        arrowHeadSource,
+                        arrowHeadTarget,
+                      },
+                    }
+                    : e,
+                ),
+              );
+            } else {
+              // Add new edge
+              handleAddEdge(
+                sourceId,
+                targetId,
+                variant,
+                label,
+                arrowHeadSource,
+                arrowHeadTarget,
+              );
+            }
+            setEdgeDialogOpen(false);
+            setEditingEdge(null);
+            setPendingConnection(null);
+          }}
+          nodes={nodes as CollaborationNode[]}
+          initialSourceId={pendingConnection?.source}
+          initialTargetId={pendingConnection?.target}
+          editingEdge={editingEdge}
+          onDeleteEdge={
+            editingEdge
+              ? () => {
+                setEdges((eds) => eds.filter((e) => e.id !== editingEdge.id));
+                setEdgeDialogOpen(false);
+                setEditingEdge(null);
+                setPendingConnection(null);
+              }
+              : undefined
+          }
+        />
+        {/* Bottom-right panel for controls */}
+        <DiagramPanel
+          historyIndex={historyIndex}
+          historyLength={history.length}
+          onUndo={handleUndo}
+          onRedo={handleRedo}
+          onClear={handleClear}
+          onImport={() => { }}
+          isSettingsOpen={false}
+          onToggleSettings={() => { }}
+          showCustomControls={true}
+          onGridLayout={() => handleLayoutDirectionChange('DOWN')}
+          onTreeLayout={() => handleLayoutDirectionChange('RIGHT')}
+        />
+        <AiButton />
+
+      </Box>
+    </Box>
+  );
+};
+
+export default CollaborationDiagramFlow;