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

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

    const edges = addEdge(connection, get().edges);
    const nodes = get().nodes;
    // connection has: { source, sourceHandle, target, targetHandle }
    // Let's find the source and target ID's.
    const sourceNode = nodes.find((n) => n.id == connection.source);
    const targetNode = nodes.find((n) => n.id == connection.target);

    // In case the nodes weren't found, return basic functionality.
    if ( sourceNode == undefined
      || targetNode == undefined
      || sourceNode.type == undefined
      || targetNode.type == undefined
    ){
      set({ nodes, edges });
      return;
    }

    // We should find out how their data changes by calling their respective functions.
    const sourceConnectFunction = NodeConnects[sourceNode.type as keyof typeof NodeConnects]
    const targetConnectFunction = NodeConnects[targetNode.type as keyof typeof NodeConnects]

    // We're going to have to update their data based on how they want to update it.
    sourceConnectFunction(sourceNode, targetNode, true)
    targetConnectFunction(targetNode, sourceNode, false)
    set({ nodes, edges });
  },

  /**
   * Handles reconnecting an edge between nodes.
   */
  onReconnect: (oldEdge, newConnection) => {
    get().edgeReconnectSuccessful = true;
    set({ edges: reconnectEdge(oldEdge, newConnection, get().edges) });
  },

  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 {{id: string}} edge - the described edge
   */
  onReconnectEnd: (_evt, edge) => {
    if (!get().edgeReconnectSuccessful) {
      set({ edges: get().edges.filter((e) => e.id !== edge.id) });
    }
    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;