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

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 { 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: (nodes) => nodes.forEach(node => get().unregisterNodeRules(node.id)),

  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) => {
    get().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) => {
    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()) {
      set({
      nodes: get().nodes.filter((n) => n.id !== nodeId),
      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();
    set({
      nodes: get().nodes.map((node) => {
        if (node.id === nodeId) {
          node = { ...node, data: { ...node.data, ...data }};
        }
        return node;
      }),
    });
  },

  /**
   * 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 };
    })
  }
  }))
);

export default useFlowStore;