import { useCallback, useMemo } from "react"; import { nanoid } from 'nanoid'; import { shallow } from 'zustand/shallow'; import { throttle } from 'lodash'; import { Edge, Node, useReactFlow } from "@xyflow/react"; import { GraphState } from "./types"; import useGraphStore from "./store"; // Store selector const selector = (store: GraphState) => ({ nodes: store.present.nodes, edges: store.present.edges, setNodes: store.setNodes, setEdges: store.setEdges, record: store.record }); // Helper functions const createNode = (label: string): Node => ({ id: nanoid(6), type: 'graph-node', data: { label }, position: { x: 0, y: 0 }, }); const createEdge = (source: string, target: string): Edge => ({ id: nanoid(6), source, target, type: 'graph-edge', }); export function useGraphOperation() { const store = useGraphStore(selector, shallow); const { addEdges, addNodes } = useReactFlow(); const selectedNodes = useMemo(() => store.nodes.filter(node => node.selected), [store.nodes] ); // Find parent node ID for a given node const findParentId = useCallback((nodeId: string) => { const parentEdge = store.edges.find(edge => edge.target === nodeId); return parentEdge?.source; }, [store.edges]); // Update node selection const updateNodeSelection = useCallback((nodeIds: string[]) => { return store.nodes.map(node => ({ ...node, selected: nodeIds.includes(node.id) })); }, [store.nodes]); // Create new node and connect it const createConnectedNode = useCallback((parentId: string, deselectOthers = true) => { const newNode = createNode(`新节点${store.nodes.length}`); const newEdge = createEdge(parentId, newNode.id); store.record(() => { addNodes({ ...newNode, selected: true }); addEdges(newEdge); if (deselectOthers) { store.setNodes(updateNodeSelection([newNode.id])); } }); }, [store, addNodes, addEdges, updateNodeSelection]); // Handle node creation operations const handleCreateChildNodes = useCallback(() => { if (selectedNodes.length === 0) return; throttle(() => { selectedNodes.forEach(node => { if (node.id) createConnectedNode(node.id); }); }, 300)(); }, [selectedNodes, createConnectedNode]); const handleCreateSiblingNodes = useCallback(() => { if (selectedNodes.length === 0) return; throttle(() => { selectedNodes.forEach(node => { const parentId = findParentId(node.id) || node.id; createConnectedNode(parentId); }); }, 300)(); }, [selectedNodes, findParentId, createConnectedNode]); const handleDeleteNodes = useCallback(() => { if (selectedNodes.length === 0) return; const nodesToDelete = new Set(); // Collect all nodes to delete including children const collectNodesToDelete = (nodeId: string) => { nodesToDelete.add(nodeId); store.edges .filter(edge => edge.source === nodeId) .forEach(edge => collectNodesToDelete(edge.target)); }; selectedNodes.forEach(node => collectNodesToDelete(node.id)); store.record(() => { // Filter out deleted nodes and their edges const remainingNodes = store.nodes.filter(node => !nodesToDelete.has(node.id)); const remainingEdges = store.edges.filter(edge => !nodesToDelete.has(edge.source) && !nodesToDelete.has(edge.target) ); // Select next node (sibling or parent of first deleted node) const firstDeletedNode = selectedNodes[0]; const parentId = findParentId(firstDeletedNode.id); let nextSelectedId: string | undefined; if (parentId) { const siblingEdge = store.edges.find(edge => edge.source === parentId && !nodesToDelete.has(edge.target) && edge.target !== firstDeletedNode.id ); nextSelectedId = siblingEdge?.target || parentId; } // Update nodes with new selection and set the remaining nodes const updatedNodes = remainingNodes.map(node => ({ ...node, selected: node.id === nextSelectedId })); store.setNodes(updatedNodes); store.setEdges(remainingEdges); }); }, [selectedNodes, store, findParentId]); return { handleCreateChildNodes, handleCreateSiblingNodes, handleDeleteNodes }; }