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 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]
  const newData = {
      id: id,
      type: type,
      position: position,
      data: data,
      deletable: deletable,
  }
  return {...defaultData, ...newData}
}

//* Initial nodes, created by using createNode. */
const initialNodes : Node[] = [
    createNode('start', 'start', {x: 100, y: 100}, {label: "Start"}, false),
    createNode('end', 'end', {x: 500, y: 100}, {label: "End"}, false),
    createNode('phase-1', 'phase',  {x:200, y:100}, {label: "Phase 1", children : []}),
    createNode('norms-1', 'norm',  {x:-200, y:100}, {label: "Initial Norms", normList: ["Be a robot", "get good"]}),     
];

// * Initial edges * /
const initialEdges: Edge[] = [
  { id: 'start-phase-1', source: 'start', target: 'phase-1' },
  { id: 'phase-1-end', source: 'phase-1', target: 'end' },
];


/**
 * 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,

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

  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,
  }))
);

export default useFlowStore;