import { useCallback } from 'react' import produce from 'immer' import type { EdgeMouseHandler, OnEdgesChange, } from 'reactflow' import { useStoreApi, } from 'reactflow' import type { Node, } from '../types' import { getNodesConnectedSourceOrTargetHandleIdsMap } from '../utils' import { useNodesSyncDraft } from './use-nodes-sync-draft' import { useNodesReadOnly } from './use-workflow' import { WorkflowHistoryEvent, useWorkflowHistory } from './use-workflow-history' export const useEdgesInteractions = () => { const store = useStoreApi() const { handleSyncWorkflowDraft } = useNodesSyncDraft() const { getNodesReadOnly } = useNodesReadOnly() const { saveStateToHistory } = useWorkflowHistory() const handleEdgeEnter = useCallback((_, edge) => { if (getNodesReadOnly()) return const { edges, setEdges, } = store.getState() const newEdges = produce(edges, (draft) => { const currentEdge = draft.find(e => e.id === edge.id)! currentEdge.data._hovering = true }) setEdges(newEdges) }, [store, getNodesReadOnly]) const handleEdgeLeave = useCallback((_, edge) => { if (getNodesReadOnly()) return const { edges, setEdges, } = store.getState() const newEdges = produce(edges, (draft) => { const currentEdge = draft.find(e => e.id === edge.id)! currentEdge.data._hovering = false }) setEdges(newEdges) }, [store, getNodesReadOnly]) const handleEdgeDeleteByDeleteBranch = useCallback((nodeId: string, branchId: string) => { if (getNodesReadOnly()) return const { getNodes, setNodes, edges, setEdges, } = store.getState() const currentEdgeIndex = edges.findIndex(edge => edge.source === nodeId && edge.sourceHandle === branchId) if (currentEdgeIndex < 0) return const currentEdge = edges[currentEdgeIndex] const newNodes = produce(getNodes(), (draft: Node[]) => { const sourceNode = draft.find(node => node.id === currentEdge.source) const targetNode = draft.find(node => node.id === currentEdge.target) if (sourceNode) sourceNode.data._connectedSourceHandleIds = sourceNode.data._connectedSourceHandleIds?.filter(handleId => handleId !== currentEdge.sourceHandle) if (targetNode) targetNode.data._connectedTargetHandleIds = targetNode.data._connectedTargetHandleIds?.filter(handleId => handleId !== currentEdge.targetHandle) }) setNodes(newNodes) const newEdges = produce(edges, (draft) => { draft.splice(currentEdgeIndex, 1) }) setEdges(newEdges) handleSyncWorkflowDraft() saveStateToHistory(WorkflowHistoryEvent.EdgeDeleteByDeleteBranch) }, [getNodesReadOnly, store, handleSyncWorkflowDraft, saveStateToHistory]) const handleEdgeDelete = useCallback(() => { if (getNodesReadOnly()) return const { getNodes, setNodes, edges, setEdges, } = store.getState() const currentEdgeIndex = edges.findIndex(edge => edge.selected) if (currentEdgeIndex < 0) return const currentEdge = edges[currentEdgeIndex] const nodes = getNodes() const nodesConnectedSourceOrTargetHandleIdsMap = getNodesConnectedSourceOrTargetHandleIdsMap( [ { type: 'remove', edge: currentEdge }, ], nodes, ) const newNodes = produce(nodes, (draft: Node[]) => { draft.forEach((node) => { if (nodesConnectedSourceOrTargetHandleIdsMap[node.id]) { node.data = { ...node.data, ...nodesConnectedSourceOrTargetHandleIdsMap[node.id], } } }) }) setNodes(newNodes) const newEdges = produce(edges, (draft) => { draft.splice(currentEdgeIndex, 1) }) setEdges(newEdges) handleSyncWorkflowDraft() saveStateToHistory(WorkflowHistoryEvent.EdgeDelete) }, [getNodesReadOnly, store, handleSyncWorkflowDraft, saveStateToHistory]) const handleEdgesChange = useCallback((changes) => { if (getNodesReadOnly()) return const { edges, setEdges, } = store.getState() const newEdges = produce(edges, (draft) => { changes.forEach((change) => { if (change.type === 'select') draft.find(edge => edge.id === change.id)!.selected = change.selected }) }) setEdges(newEdges) }, [store, getNodesReadOnly]) const handleEdgeCancelRunningStatus = useCallback(() => { const { edges, setEdges, } = store.getState() const newEdges = produce(edges, (draft) => { draft.forEach((edge) => { edge.data._run = false }) }) setEdges(newEdges) }, [store]) return { handleEdgeEnter, handleEdgeLeave, handleEdgeDeleteByDeleteBranch, handleEdgeDelete, handleEdgesChange, handleEdgeCancelRunningStatus, } }