pepperplus-ui/src/pages/VisProgPage/visualProgrammingUI/VisProgStores.tsx

// This program has been developed by students from the bachelor Computer Science at Utrecht
// University within the Software Project course.
// © Copyright Utrecht University (Department of Information and Computing Sciences)
import { create } from 'zustand';
import {
  applyNodeChanges,
  applyEdgeChanges,
  addEdge,
  reconnectEdge,
  type Node,
  type Edge,
  type XYPosition,
} from '@xyflow/react';
import '@xyflow/react/dist/style.css';
import {type ConnectionContext, validateConnectionWithRules} from "./HandleRuleLogic.ts";
import {editorWarningRegistry} from "./components/EditorWarnings.tsx";
import type { FlowState } from './VisProgTypes';
import {
  NodeDefaults,
  NodeConnections as NodeCs,
  NodeDisconnections as NodeDs,
  NodeDeletes
} from './NodeRegistry';
import { UndoRedo } from "./EditorUndoRedo.ts";


/**
 * A Function to create a new node with the correct default data and properties.
 *
 * @param id - The unique ID of the node.
 * @param type - The type of node to create (must exist in NodeDefaults).
 * @param position - The XY position of the node in the flow canvas.
 * @param data - The data object to initialize the node with.
 * @param deletable - Optional flag to indicate if the node can be deleted (can be deleted by default).
 * @returns A fully initialized Node object ready to be added to the flow.
 */
function createNode(id: string, type: string, position: XYPosition, data: Record<string, unknown>, deletable?: boolean) {
    const defaultData = NodeDefaults[type as keyof typeof NodeDefaults]
    return {
      id,
      type,
      position,
      deletable,
      data: {
        ...JSON.parse(JSON.stringify(defaultData)),
        ...data,
      },
    }
}

//* Initial nodes, created by using createNode. */
// Start and End don't need to apply the UUID, since they are technically never compiled into a program.
const startNode = createNode('start', 'start', {x: 110, y: 100}, {label: "Start"}, false)
const endNode = createNode('end', 'end', {x: 590, y: 100}, {label: "End"}, false)
const initialPhaseNode = createNode(crypto.randomUUID(), 'phase', {x:235, y:100}, {label: "Phase 1", children : [], isFirstPhase: false, nextPhaseId: null})

const initialNodes : Node[] = [startNode, endNode, initialPhaseNode];

// Initial edges, leave empty as setting initial edges...
// ...breaks logic that is dependent on connection events
const initialEdges: Edge[] = [];

/**
 * useFlowStore contains the implementation for all editor functionality
 * and stores the current state of the visual programming editor
 *
 *  * Provides:
 * - Node and edge state management
 * - Node creation, deletion, and updates
 * - Custom connection handling via NodeConnects
 * - Edge reconnection handling
 * - Undo Redo functionality through custom middleware
 */
const useFlowStore = create<FlowState>(UndoRedo((set, get) => ({
  nodes: initialNodes,
  edges: initialEdges,
  edgeReconnectSuccessful: true,
  scrollable: true,

  /**
   * handles changing the scrollable state of the editor,
   * this is used to control if scrolling is captured by the editor
   * or if it's available to other components within the reactFlowProvider
   * @param {boolean} val - the desired state
   */
  setScrollable: (val) => set({scrollable: val}),

   /**
   * Handles changes to nodes triggered by ReactFlow.
   */
  onNodesChange: (changes) => set({nodes: applyNodeChanges(changes, get().nodes)}),

  onNodesDelete: (deletedNodes) => {

    const allNodes = get().nodes;
    const deletedIds = new Set(deletedNodes.map(n => n.id));

    deletedNodes.forEach((node) => {
      get().unregisterNodeRules(node.id);
      get().unregisterWarningsForId(node.id);
    });
    const remainingNodes = allNodes.filter((node) => !deletedIds.has(node.id));
    
    // Validate only the survivors
    get().validateDuplicateNames(remainingNodes);
  },

  onEdgesDelete: (edges) => {
    // we make sure any affected nodes get updated to reflect removal of edges
    edges.forEach((edge) => {
        const nodes = get().nodes;

        const sourceNode = nodes.find((n) => n.id == edge.source);
        const targetNode = nodes.find((n) => n.id == edge.target);

        if (sourceNode) { NodeDs.Sources[sourceNode.type as keyof typeof NodeDs.Sources](sourceNode, edge.target); }
        if (targetNode) { NodeDs.Targets[targetNode.type as keyof typeof NodeDs.Targets](targetNode, edge.source); }
      });
  },
  /**
   * Handles changes to edges triggered by ReactFlow.
   */
  onEdgesChange: (changes) => {
    set({ edges: applyEdgeChanges(changes, get().edges) })
  },

  /**
   * Handles creating a new connection between nodes.
   * Updates edges and calls the node-specific connection functions.
   */
  onConnect: (connection) => {
    get().pushSnapshot();
    set({edges: addEdge(connection, get().edges)});

    // We make sure to perform any required data updates on the newly connected nodes
    const nodes = get().nodes;

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

    if (sourceNode) { NodeCs.Sources[sourceNode.type as keyof typeof NodeCs.Sources](sourceNode, connection.target); }
    if (targetNode) { NodeCs.Targets[targetNode.type as keyof typeof NodeCs.Targets](targetNode, connection.source); }
  },

  /**
   * Handles reconnecting an edge between nodes.
   */
  onReconnect: (oldEdge, newConnection) => {

    function createContext(
      source: {id: string, handleId: string},
      target: {id: string, handleId: string}
    ) : ConnectionContext {
      const edges = get().edges;
      const targetConnections = edges.filter(edge => edge.target === target.id && edge.targetHandle === target.handleId).length
      return {
        connectionCount: targetConnections,
        source: source,
        target: target
      }
    }

    // connection validation
    const context: ConnectionContext = oldEdge.source === newConnection.source
      ? createContext({id: newConnection.source, handleId: newConnection.sourceHandle!}, {id: newConnection.target, handleId: newConnection.targetHandle!})
      : createContext({id: newConnection.target, handleId: newConnection.targetHandle!}, {id: newConnection.source, handleId: newConnection.sourceHandle!});

    const result = validateConnectionWithRules(
      newConnection,
      context
    );

    if (!result.isSatisfied) {
      set({
        edges: get().edges.map(e =>
          e.id === oldEdge.id ? oldEdge : e
        ),
      });
      return;
    }

    // further reconnect logic
    set({ edgeReconnectSuccessful: true });
    set({ edges: reconnectEdge(oldEdge, newConnection, get().edges) });

    // We make sure to perform any required data updates on the newly reconnected nodes
    const nodes = get().nodes;

    const oldSourceNode = nodes.find((n) => n.id == oldEdge.source)!;
    const oldTargetNode = nodes.find((n) => n.id == oldEdge.target)!;
    const newSourceNode = nodes.find((n) => n.id == newConnection.source)!;
    const newTargetNode = nodes.find((n) => n.id == newConnection.target)!;

    if (oldSourceNode === newSourceNode && oldTargetNode === newTargetNode) return;

    NodeCs.Sources[newSourceNode.type as keyof typeof NodeCs.Sources](newSourceNode, newConnection.target);
    NodeCs.Targets[newTargetNode.type as keyof typeof NodeCs.Targets](newTargetNode, newConnection.source);

    NodeDs.Sources[oldSourceNode.type as keyof typeof NodeDs.Sources](oldSourceNode, oldEdge.target);
    NodeDs.Targets[oldTargetNode.type as keyof typeof NodeDs.Targets](oldTargetNode, oldEdge.source);
  },

  onReconnectStart: () => {
    get().pushSnapshot();
    set({ edgeReconnectSuccessful: false })
  },

  /**
   * handles potential dropping (deleting) of an edge
   * if it is not reconnected to a node after detaching it
   *
   * @param _evt - the event
   * @param edge - the described edge
   */
  onReconnectEnd: (_evt, edge) => {
    if (!get().edgeReconnectSuccessful) {
      // delete the edge from the flowState
      set({ edges: get().edges.filter((e) => e.id !== edge.id) });

      // update node data to reflect the dropped edge
      const nodes = get().nodes;

      const sourceNode = nodes.find((n) => n.id == edge.source)!;
      const targetNode = nodes.find((n) => n.id == edge.target)!;

      NodeDs.Sources[sourceNode.type as keyof typeof NodeDs.Sources](sourceNode, edge.target);
      NodeDs.Targets[targetNode.type as keyof typeof NodeDs.Targets](targetNode, edge.source);
    }
    set({ edgeReconnectSuccessful: true });
  },

  /**
   * Deletes a node by ID, respecting NodeDeletes rules.
   * Also removes all edges connected to that node.
   */
  deleteNode: (nodeId, deleteElements) => {
    get().pushSnapshot();

    // Let's find our node to check if they have a special deletion function
    const ourNode = get().nodes.find((n)=>n.id==nodeId);
    const ourFunction = Object.entries(NodeDeletes).find(([t])=>t==ourNode?.type)?.[1]



    // If there's no function, OR, our function tells us we can delete it, let's do so...
    if (ourFunction == undefined || ourFunction()) {
      if (deleteElements){
        deleteElements({
          nodes: get().nodes.filter((n) => n.id === nodeId),
          edges: get().edges.filter((e) => e.source !== nodeId && e.target === nodeId)}
        ).then(() => {
          get().unregisterNodeRules(nodeId);
          get().unregisterWarningsForId(nodeId);
          // Re-validate after deletion is finished
          get().validateDuplicateNames(get().nodes);
        });
      } else {
        const remainingNodes = get().nodes.filter((n) => n.id !== nodeId);
        get().validateDuplicateNames(remainingNodes); // Re-validate survivors
        set({
        nodes: remainingNodes,
        edges: get().edges.filter((e) => e.source !== nodeId && e.target !== nodeId),
        })
      }
    }
  },

  /**
   * Replaces the entire nodes array in the store.
   */
  setNodes: (nodes) => set({ nodes }),

  /**
   * Replaces the entire edges array in the store.
   */
  setEdges: (edges) => set({ edges }),

  /**
   * Updates the data of a node by merging new data with existing data.
   */
  updateNodeData: (nodeId, data) => {
    get().pushSnapshot();
    const updatedNodes = get().nodes.map((node) => {
      if (node.id === nodeId) {
        return { ...node, data: { ...node.data, ...data } };
      }
      return node;
    });

    get().validateDuplicateNames(updatedNodes); // Re-validate after update
    set({ nodes: updatedNodes });
  },

  //helper function to see if any of the nodes have duplicate names
  validateDuplicateNames: (nodes: Node[]) => {
    const nameMap = new Map<string, string[]>();

    // 1. Group IDs by their identifier (name, norm, or label)
    nodes.forEach((n) => {
      const name = (n.data.name || n.data.norm )?.toString().trim();
      if (name) {
        if (!nameMap.has(name)) nameMap.set(name, []);
        nameMap.get(name)!.push(n.id);
      }
    });

    // 2. Scan nodes and toggle the warning
    nodes.forEach((n) => {
      const name = (n.data.name || n.data.norm )?.toString().trim();
      const isDuplicate = name ? (nameMap.get(name)?.length || 0) > 1 : false;

      if (isDuplicate) {
        get().registerWarning({
          scope: { id: n.id },
          type: 'DUPLICATE_ELEMENT_NAME',
          severity: 'ERROR',
          description: `The name "${name}" is already used by another element.`
        });
      } else {
        // This clears the warning if the "twin" was deleted or renamed
        get().unregisterWarning(n.id, 'DUPLICATE_ELEMENT_NAME');
      }
    });
  },
  

  /**
   * Adds a new node to the flow store.
   */
  addNode: (node: Node) => {
    get().pushSnapshot();
    set({ nodes: [...get().nodes, node] });
  },

  // undo redo default values
  past: [],
  future: [],
  isBatchAction: false,

  // handleRuleRegistry definitions
  /**
   * stores registered rules for handle connection validation
   */
  ruleRegistry: new Map(),

  /**
   * gets the rules registered by that handle described by the given node and handle ids
   *
   * @param {string} targetNodeId
   * @param {string} targetHandleId
   * @returns {HandleRule[]}
   */
  getTargetRules: (targetNodeId, targetHandleId) => {
    const key = `${targetNodeId}:${targetHandleId}`;
    const rules = get().ruleRegistry.get(key);

    // helper function that handles a situation where no rules were registered
    const missingRulesResponse = () => {
      console.warn(
        `No rules were registered for the following handle "${key}"!
          returning and empty handleRule[] to avoid crashing`);
      return []
    }

    return rules
      ? rules
      : missingRulesResponse()
  },

  /**
   * registers a handle's connection rules
   *
   * @param {string} nodeId
   * @param {string} handleId
   * @param {HandleRule[]} rules
   */
  registerRules: (nodeId, handleId, rules) => {
      const registry = get().ruleRegistry;
      registry.set(`${nodeId}:${handleId}`, rules);
      set({ ruleRegistry: registry }) ;
  },

  /**
   * unregisters a handles connection rules
   *
   * @param {string} nodeId
   * @param {string} handleId
   */
  unregisterHandleRules: (nodeId, handleId) => {
    set( () => {
      const registry = get().ruleRegistry;
      registry.delete(`${nodeId}:${handleId}`);
      return { ruleRegistry: registry };
    })
  },

  /**
   * unregisters connection rules for all handles on the given node
   * used for cleaning up rules on node deletion
   *
   * @param {string} nodeId
   */
  unregisterNodeRules: (nodeId) => {
    set(() => {
      const registry = get().ruleRegistry;
      registry.forEach((_,key) => {
        if (key.startsWith(`${nodeId}:`)) registry.delete(key)
      })
      return { ruleRegistry: registry };
    })
  },
  ...editorWarningRegistry(get, set),
  }))
);



export default useFlowStore;