collect-system/apps/web/src/components/common/editor/graph/useGraphOperation.ts

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<string>();

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