import type { DataSourceNodeType } from '@/app/components/workflow/nodes/data-source/types'
import type { Node, ValueSelector } from '@/app/components/workflow/types'
import { uniqBy } from 'es-toolkit/compat'
import { useCallback } from 'react'
import { getOutgoers, useStoreApi } from 'reactflow'
import {
  findUsedVarNodes,
  updateNodeVars,
} from '@/app/components/workflow/nodes/_base/components/variable/utils'
import { BlockEnum } from '@/app/components/workflow/types'

export const usePipeline = () => {
  const store = useStoreApi()

  const getAllDatasourceNodes = useCallback(() => {
    const { getNodes } = store.getState()
    const nodes = getNodes() as Node<DataSourceNodeType>[]
    const datasourceNodes = nodes.filter((node) => node.data.type === BlockEnum.DataSource)

    return datasourceNodes
  }, [store])

  const getAllNodesInSameBranch = useCallback(
    (nodeId: string) => {
      const { getNodes, edges } = store.getState()
      const nodes = getNodes()
      const list: Node[] = []

      const traverse = (root: Node, callback: (node: Node) => void) => {
        if (root) {
          const outgoers = getOutgoers(root, nodes, edges)

          if (outgoers.length) {
            outgoers.forEach((node) => {
              callback(node)
              traverse(node, callback)
            })
          }
        }
      }

      if (nodeId === 'shared') {
        const allDatasourceNodes = getAllDatasourceNodes()

        if (allDatasourceNodes.length === 0) return []

        list.push(...allDatasourceNodes)

        allDatasourceNodes.forEach((node) => {
          traverse(node, (childNode) => {
            list.push(childNode)
          })
        })
      } else {
        const currentNode = nodes.find((node) => node.id === nodeId)!

        if (!currentNode) return []

        list.push(currentNode)

        traverse(currentNode, (node) => {
          list.push(node)
        })
      }

      return uniqBy(list, 'id')
    },
    [getAllDatasourceNodes, store],
  )

  const isVarUsedInNodes = useCallback(
    (varSelector: ValueSelector) => {
      const nodeId = varSelector[1] // Assuming the first element is always 'VARIABLE_PREFIX'(rag)
      const afterNodes = getAllNodesInSameBranch(nodeId!)
      const effectNodes = findUsedVarNodes(varSelector, afterNodes)
      return effectNodes.length > 0
    },
    [getAllNodesInSameBranch],
  )

  const handleInputVarRename = useCallback(
    (nodeId: string, oldValeSelector: ValueSelector, newVarSelector: ValueSelector) => {
      const { getNodes, setNodes } = store.getState()
      const afterNodes = getAllNodesInSameBranch(nodeId)
      const effectNodes = findUsedVarNodes(oldValeSelector, afterNodes)
      if (effectNodes.length > 0) {
        const newNodes = getNodes().map((node) => {
          if (effectNodes.find((n) => n.id === node.id))
            return updateNodeVars(node, oldValeSelector, newVarSelector)

          return node
        })
        setNodes(newNodes)
      }
    },
    [getAllNodesInSameBranch, store],
  )

  const removeUsedVarInNodes = useCallback(
    (varSelector: ValueSelector) => {
      const nodeId = varSelector[1] // Assuming the first element is always 'VARIABLE_PREFIX'(rag)
      const { getNodes, setNodes } = store.getState()
      const afterNodes = getAllNodesInSameBranch(nodeId!)
      const effectNodes = findUsedVarNodes(varSelector, afterNodes)
      if (effectNodes.length > 0) {
        const newNodes = getNodes().map((node) => {
          if (effectNodes.find((n) => n.id === node.id))
            return updateNodeVars(node, varSelector, [])

          return node
        })
        setNodes(newNodes)
      }
    },
    [getAllNodesInSameBranch, store],
  )

  return {
    handleInputVarRename,
    isVarUsedInNodes,
    removeUsedVarInNodes,
  }
}
