@eventcatalogtest/studio/stores/flow-store.ts

A drag and drop UI for distributed systems that keeps your diagrams where they belong – in your repo

import { Edge, ReactFlowInstance, Node, applyNodeChanges, NodeChange, addEdge, Connection, MarkerType, ReactFlowJsonObject } from '@xyflow/react';
import { create } from 'zustand';
import { persist, createJSONStorage } from 'zustand/middleware';
import ELK from 'elkjs/lib/elk.bundled.js';

import { getEdgeOptionsForConnection, isValidConnection } from '@/utils/node-rule-engine';
import { useNotificationsStore } from './notifications-store';
import { getNodeId } from '@/utils/react-flow';

import slugify from 'slugify';
import { nanoid } from 'nanoid';

interface HistoryState {
  nodes: Node[];
  edges: Edge[];
}

interface Store {
  reactFlowInstance: ReactFlowInstance | null;
  nodes: Node[];
  edges: Edge[];
  history: HistoryState[];
  historyIndex: number;
  actions: {
    addNode: (node: Node) => void;
    addEdge: (edge: Edge) => void;
    deleteNode: (id: string) => void;
    deleteEdge: (id: string) => void;
    setReactFlowInstance: (instance: ReactFlowInstance) => void;
    onNodesChange: (changes: NodeChange<any>[]) => void;
    setEdges: (change: any) => void;
    onConnect: (connection: Connection) => void;
    duplicateNode: (id: string) => void;
    updateNode: (node: Node) => void;
    updateEdge: (edge: Edge) => void;
    reset: () => void;
    exportFlow: () => any;
    importFlow: (data: ReactFlowJsonObject) => void;
    saveCurrentFlowToDesign: () => void;
    autoLayout: () => void;
    undo: () => void;
    redo: () => void;
    canUndo: () => boolean;
    canRedo: () => boolean;
    saveToHistory: () => void;
    checkAndAssignParentIds: (nodes: Node[]) => Node[];
    sortNodesByParentChild: (nodes: Node[]) => Node[];
  };
}

const initialNodes: Node[] = [];
// const initialEdges = [{ id: 'e1-2', source: '1', target: '2' }];
const initialEdges: Edge[] = [];

// TODO: Get the initial edges and nodes from localstorage?


const useFlowStore = create<Store>()(
  persist(
    (set, get) => ({
      reactFlowInstance: null,
      nodes: initialNodes,
      edges: initialEdges,
      history: [{ nodes: initialNodes, edges: initialEdges }],
      historyIndex: 0,

      actions: {
        addNode: (node: Node) => {
          get().actions.saveToHistory();
          const newNodes = [...get().nodes, node];
          const sortedNodes = get().actions.sortNodesByParentChild(newNodes);
          set({ nodes: sortedNodes });
        },
        addEdge: (edge: Edge) => {
          get().actions.saveToHistory();
          set((state) => ({ edges: [...state.edges, edge] }));
        },
        setEdges: (change: any) => {
          let newChange = typeof change === "function" ? change(get().edges) : change;
          set({
            edges: newChange,
          });
          // get().updateCurrentFlow({ edges: newChange });
        },
        duplicateNode: (id: string) => {
          const node = get().nodes.find((node) => node.id === id);
          const nodes = get().nodes;
          // Set all previous nodes to selected false
          nodes.forEach((node) => {
            node.selected = false;
          });

          // Get the window mouse position
          if (node) {
            get().actions.addNode({
              ...node,
              id: getNodeId(node.type ?? ''),
              position: {
                x: node.position.x + 100,
                y: node.position.y + 150,
              },
              selected: false,
            });
          }
        },
        deleteNode: (id: string) => {
          get().actions.saveToHistory();
          set((state) => ({ nodes: state.nodes.filter((node) => node.id !== id) }));
        },
        deleteEdge: (id: string) => {
          get().actions.saveToHistory();
          set((state) => ({ edges: state.edges.filter((edge) => edge.id !== id) }));
        },
        setReactFlowInstance: (instance: ReactFlowInstance) => set({ reactFlowInstance: instance }),
        onNodesChange: (changes: NodeChange<any>[]) => {
          // Check if there are position changes that ended (drag end)
          const hasPositionEnd = changes.some(change =>
            change.type === 'position' && 'dragging' in change && !change.dragging
          );

          // Check for dimensions changes (resize end)
          const hasDimensionsEnd = changes.some(change =>
            change.type === 'dimensions' && 'resizing' in change && !change.resizing
          );

          if (hasPositionEnd || hasDimensionsEnd) {
            get().actions.saveToHistory();
          }

          // Apply the changes first
          const updatedNodes = applyNodeChanges(changes, get().nodes);

          // Check for nodes that need parentId assignment after drag end
          if (hasPositionEnd) {
            const finalNodes = get().actions.checkAndAssignParentIds(updatedNodes);
            // Sort nodes so parent nodes come before their children
            const sortedNodes = get().actions.sortNodesByParentChild(finalNodes);
            set({ nodes: sortedNodes });
          } else {
            set({ nodes: updatedNodes });
          }
        },
        reset: () => {
          set({
            nodes: [],
            edges: [],
          });
        },
        onConnect: (connection: Connection) => {
          let newEdges: Edge[] = [];

          const sourceNode = get().nodes.find((node) => node.id === connection.source);
          const targetNode = get().nodes.find((node) => node.id === connection.target);

          const validConnection = isValidConnection(sourceNode?.type ?? '', targetNode?.type ?? '', connection);

          if (!validConnection) {
            useNotificationsStore.getState().addNotification({
              message: `Invalid connection between ${sourceNode?.type} and ${targetNode?.type}`,
              type: 'error',
              duration: 5000,
            });
            return;
          }

          const edgeOptions = getEdgeOptionsForConnection(sourceNode?.type ?? '', targetNode?.type ?? '', connection);

          get().actions.setEdges((oldEdges: Edge[]) => {
            newEdges = addEdge(
              {
                ...connection,
                ...edgeOptions,
                animated: true,
                type: 'animatedMessage',
                data: {
                  source: sourceNode?.type,
                  target: targetNode?.type,
                  message: {
                    collection: 'events',
                    opacity: 1,
                  },
                },
              },
              oldEdges,
            );

            return newEdges;
          });
        },
        updateNode: (node: Node) => {
          const nodes = get().nodes;
          set((state) => ({ nodes: state.nodes.map((n) => n.id === node.id ? node : n) }))
        },
        updateEdge: (edge: Edge) => {
          get().actions.saveToHistory();
          set((state) => ({ edges: state.edges.map((e) => e.id === edge.id ? edge : e) }));
        },
        exportFlow: () => {
          const flowData = get().reactFlowInstance?.toObject();
          if (!flowData) return null;

          // Get current design name from design store
          let designName = 'Untitled Design';
          let designId = '';
          try {
            // Dynamically import to avoid circular dependency
            const { useDesignStore } = require('./design-store');
            const currentDesign = useDesignStore.getState().currentDesign;
            if (currentDesign?.name) {
              designName = currentDesign.name;
            }
            if (currentDesign?.id) {
              designId = currentDesign.id;
            }
          } catch (error) {
            // Fallback if design store is not available
            console.warn('Could not access design store for name');
          }

          // Get current document
          let document = null;
          try {
            const { useDocumentStore } = require('./document-store');
            document = useDocumentStore.getState().getCurrentDocument();
          } catch (error) {
            console.warn('Could not access document store');
          }

          // Add metadata fields to the root of the object
          const dataWithMetadata = {
            ...flowData,
            creationDate: new Date().toISOString(),
            name: designName,
            version: '1.0',
            source: 'https://app.eventcatalog.dev',
            document,
            appState: {},
            id: `${slugify(designName, { lower: true })}`
          };

          return dataWithMetadata;
        },
        importFlow: (data: ReactFlowJsonObject & { creationDate?: string; name?: string; version?: string; source?: string; appState?: any; document?: any }) => {
          get().reactFlowInstance?.setViewport(data.viewport);
          // set the nodes and edges
          set({
            nodes: data.nodes,
            edges: data.edges,
          });

          // Handle document import - load document into current document store
          try {
            const { useDocumentStore } = require('./document-store');
            if (data.document) {
              // Use setTimeout to ensure document import happens after the flow import is complete
              setTimeout(() => {
                // Check if it's the new object format
                if (data.document.id && data.document.content) {
                  useDocumentStore.getState().setCurrentDocument(data.document);
                }
                // Legacy support: if it's still an array format, convert it
                else if (Array.isArray(data.document) && data.document.length > 0) {
                  const convertedDoc = {
                    id: `doc-${Date.now()}`,
                    name: 'Imported Document',
                    content: data.document,
                    createdAt: new Date().toISOString(),
                    updatedAt: new Date().toISOString(),
                  };
                  useDocumentStore.getState().setCurrentDocument(convertedDoc);
                }
              }, 100);
            }
          } catch (error) {
            console.warn('Could not import document:', error);
          }

          // Handle metadata if present
          if (data.creationDate || data.name || data.version) {
            console.log('Loaded file metadata:', {
              creationDate: data.creationDate,
              name: data.name,
              version: data.version,
              source: data.source,
              appState: data.appState,
              hasDocument: !!(data.document?.length)
            });
          }
        },
        saveCurrentFlowToDesign: () => {
          const flowData = get().actions.exportFlow();
          // @ts-ignore
          if (flowData) {
            // Dynamically import to avoid circular dependency
            import('./design-store').then(({ useDesignStore }) => {
              useDesignStore.getState().updateCurrentDesignData(flowData);
            });
          }
        },
        autoLayout: async () => {
          const { nodes, edges } = get();

          if (nodes.length === 0) {
            useNotificationsStore.getState().addNotification({
              message: 'No nodes to layout',
              type: 'warning',
              duration: 3000,
            });
            return;
          }

          // Save current state before auto-layout
          get().actions.saveToHistory();

          const elk = new ELK();

          // Convert ReactFlow nodes and edges to ELK format
          const elkNodes = nodes.map(node => ({
            id: node.id,
            width: node.measured?.width || 200,
            height: node.measured?.height || 100,
          }));

          const elkEdges = edges.map(edge => ({
            id: edge.id,
            sources: [edge.source],
            targets: [edge.target],
          }));

          const elkGraph = {
            id: 'root',
            layoutOptions: {
              'elk.algorithm': 'layered',
              'elk.direction': 'RIGHT',
              'elk.spacing.nodeNode': '80',
              'elk.layered.spacing.nodeNodeBetweenLayers': '120',
              'elk.layered.spacing.edgeNodeBetweenLayers': '30',
            },
            children: elkNodes,
            edges: elkEdges,
          };

          try {
            const layoutedGraph = await elk.layout(elkGraph);

            // Update node positions based on ELK layout
            const layoutedNodes = nodes.map(node => {
              const elkNode = layoutedGraph.children?.find(n => n.id === node.id);
              if (elkNode && elkNode.x !== undefined && elkNode.y !== undefined) {
                return {
                  ...node,
                  position: { x: elkNode.x, y: elkNode.y },
                };
              }
              return node;
            });

            set({ nodes: layoutedNodes });

            useNotificationsStore.getState().addNotification({
              message: 'Nodes arranged automatically',
              type: 'info',
              duration: 3000,
            });

          } catch (error) {
            useNotificationsStore.getState().addNotification({
              message: 'Failed to auto-layout nodes',
              type: 'error',
              duration: 5000,
            });
          }
        },
        saveToHistory: () => {
          const { nodes, edges, history, historyIndex } = get();
          // Remove any future history if we're not at the end
          const newHistory = history.slice(0, historyIndex + 1);
          // Add current state to history
          newHistory.push({ nodes: [...nodes], edges: [...edges] });
          // Limit history to 50 entries
          const limitedHistory = newHistory.slice(-50);
          set({
            history: limitedHistory,
            historyIndex: limitedHistory.length - 1,
          });
        },
        undo: () => {
          const { history, historyIndex } = get();
          if (historyIndex > 0) {
            const previousState = history[historyIndex - 1];
            if (previousState) {
              set({
                nodes: [...previousState.nodes],
                edges: [...previousState.edges],
                historyIndex: historyIndex - 1,
              });
            }
          }
        },
        redo: () => {
          const { history, historyIndex } = get();
          if (historyIndex < history.length - 1) {
            const nextState = history[historyIndex + 1];
            if (nextState) {
              set({
                nodes: [...nextState.nodes],
                edges: [...nextState.edges],
                historyIndex: historyIndex + 1,
              });
            }
          }
        },
        canUndo: () => {
          const { historyIndex } = get();
          return historyIndex > 0;
        },
        canRedo: () => {
          const { history, historyIndex } = get();
          return historyIndex < history.length - 1;
        },
        checkAndAssignParentIds: (nodes: Node[]) => {
          // Helper function to check if a point is inside a rectangle
          const isPointInside = (point: { x: number; y: number }, rect: { x: number; y: number; width: number; height: number }) => {
            return point.x >= rect.x &&
              point.x <= rect.x + rect.width &&
              point.y >= rect.y &&
              point.y <= rect.y + rect.height;
          };

          // Find all domain nodes (potential parents)
          const domainNodes = nodes.filter(node => node.type === 'domain');

          // Process all non-domain nodes to check if they should be assigned to a domain
          return nodes.map(node => {
            // Skip domain nodes themselves
            if (node.type === 'domain') {
              return node;
            }

            // Calculate the node's absolute position (accounting for parentId)
            let nodeAbsolutePosition = { ...node.position };
            if (node.parentId) {
              const currentParent = nodes.find(n => n.id === node.parentId);
              if (currentParent) {
                nodeAbsolutePosition = {
                  x: node.position.x + currentParent.position.x,
                  y: node.position.y + currentParent.position.y
                };
              }
            }

            // Find all domains that contain this node
            const containingDomains = domainNodes.filter(domain => {
              // Calculate domain bounds (accounting for node dimensions)
              const widthValue = domain.measured?.width || domain.style?.width || 400;
              const heightValue = domain.measured?.height || domain.style?.height || 300;

              const domainBounds = {
                x: domain.position.x,
                y: domain.position.y,
                width: typeof widthValue === 'number' ? widthValue : parseInt(String(widthValue)) || 400,
                height: typeof heightValue === 'number' ? heightValue : parseInt(String(heightValue)) || 300
              };

              // Check if the node's center point is inside the domain
              const nodeCenter = {
                x: nodeAbsolutePosition.x + ((node.measured?.width || 200) / 2),
                y: nodeAbsolutePosition.y + ((node.measured?.height || 100) / 2)
              };

              console.log(`Checking node ${node.id} (center: ${nodeCenter.x}, ${nodeCenter.y}) against domain ${domain.id} (bounds: ${domainBounds.x}, ${domainBounds.y}, ${domainBounds.width}, ${domainBounds.height})`);

              return isPointInside(nodeCenter, domainBounds);
            });

            // If multiple domains contain the node, pick the smallest one (most specific)
            const parentDomain = containingDomains.length > 0
              ? containingDomains.reduce((smallest, current) => {
                const smallestWidthValue = smallest.measured?.width || smallest.style?.width || 400;
                const smallestHeightValue = smallest.measured?.height || smallest.style?.height || 300;
                const currentWidthValue = current.measured?.width || current.style?.width || 400;
                const currentHeightValue = current.measured?.height || current.style?.height || 300;

                const smallestWidth = typeof smallestWidthValue === 'number' ? smallestWidthValue : parseInt(String(smallestWidthValue)) || 400;
                const smallestHeight = typeof smallestHeightValue === 'number' ? smallestHeightValue : parseInt(String(smallestHeightValue)) || 300;
                const currentWidth = typeof currentWidthValue === 'number' ? currentWidthValue : parseInt(String(currentWidthValue)) || 400;
                const currentHeight = typeof currentHeightValue === 'number' ? currentHeightValue : parseInt(String(currentHeightValue)) || 300;

                const smallestArea = smallestWidth * smallestHeight;
                const currentArea = currentWidth * currentHeight;
                return currentArea < smallestArea ? current : smallest;
              })
              : null;

            // Update parentId based on whether node is inside a domain
            if (parentDomain && node.parentId !== parentDomain.id) {
              // Convert absolute position to relative position within the parent
              const relativePosition = {
                x: node.position.x - parentDomain.position.x,
                y: node.position.y - parentDomain.position.y
              };

              return {
                ...node,
                parentId: parentDomain.id,
                position: relativePosition,
                extent: 'parent' as const
              };
            } else if (!parentDomain && node.parentId) {
              // Node was moved out of its parent domain
              const parentNode = nodes.find(n => n.id === node.parentId);
              if (parentNode) {
                // Convert relative position back to absolute position
                const absolutePosition = {
                  x: node.position.x + parentNode.position.x,
                  y: node.position.y + parentNode.position.y
                };

                return {
                  ...node,
                  parentId: undefined,
                  position: absolutePosition,
                  extent: undefined
                };
              }
            }

            return node;
          });
        },
        sortNodesByParentChild: (nodes: Node[]) => {
          // Separate parent nodes (domains) and child nodes
          const parentNodes = nodes.filter(node => node.type === 'domain');
          const childNodes = nodes.filter(node => node.type !== 'domain');

          // Return parent nodes first, then child nodes
          return [...parentNodes, ...childNodes];
        },
      },
    }),
    {
      name: 'flow-storage', // unique name for localStorage key
      storage: createJSONStorage(() => localStorage),
      partialize: (state) => ({ nodes: state.nodes, edges: state.edges }), // only persist nodes and edges
    }
  )
);

export default useFlowStore;
export const useFlowStoreActions = () => useFlowStore((state) => state.actions);