import type {
  AgentLogItemWithChildren,
  IterationDurationMap,
  LoopDurationMap,
  LoopVariableMap,
  NodeTracing,
} from '@/types/workflow'
import { useBoolean } from 'ahooks'
import {
  useCallback,
  useRef,
  useState,
} from 'react'

export const useLogs = () => {
  const [showRetryDetail, {
    setTrue: setShowRetryDetailTrue,
    setFalse: setShowRetryDetailFalse,
  }] = useBoolean(false)
  const [retryResultList, setRetryResultList] = useState<NodeTracing[]>([])
  const handleShowRetryResultList = useCallback((detail: NodeTracing[]) => {
    setShowRetryDetailTrue()
    setRetryResultList(detail)
  }, [setShowRetryDetailTrue, setRetryResultList])

  const [showIteratingDetail, {
    setTrue: setShowIteratingDetailTrue,
    setFalse: setShowIteratingDetailFalse,
  }] = useBoolean(false)
  const [iterationResultList, setIterationResultList] = useState<NodeTracing[][]>([])
  const [iterationResultDurationMap, setIterationResultDurationMap] = useState<IterationDurationMap>({})
  const handleShowIterationResultList = useCallback((detail: NodeTracing[][], iterDurationMap: IterationDurationMap) => {
    setShowIteratingDetailTrue()
    setIterationResultList(detail)
    setIterationResultDurationMap(iterDurationMap)
  }, [setShowIteratingDetailTrue, setIterationResultList, setIterationResultDurationMap])

  const [showLoopingDetail, {
    setTrue: setShowLoopingDetailTrue,
    setFalse: setShowLoopingDetailFalse,
  }] = useBoolean(false)
  const [loopResultList, setLoopResultList] = useState<NodeTracing[][]>([])
  const [loopResultDurationMap, setLoopResultDurationMap] = useState<LoopDurationMap>({})
  const [loopResultVariableMap, setLoopResultVariableMap] = useState<Record<string, any>>({})
  const handleShowLoopResultList = useCallback((detail: NodeTracing[][], loopDurationMap: LoopDurationMap, loopVariableMap: LoopVariableMap) => {
    setShowLoopingDetailTrue()
    setLoopResultList(detail)
    setLoopResultDurationMap(loopDurationMap)
    setLoopResultVariableMap(loopVariableMap)
  }, [setShowLoopingDetailTrue, setLoopResultList, setLoopResultDurationMap])

  const [agentOrToolLogItemStack, setAgentOrToolLogItemStack] = useState<AgentLogItemWithChildren[]>([])
  const agentOrToolLogItemStackRef = useRef(agentOrToolLogItemStack)
  const [agentOrToolLogListMap, setAgentOrToolLogListMap] = useState<Record<string, AgentLogItemWithChildren[]>>({})
  const agentOrToolLogListMapRef = useRef(agentOrToolLogListMap)
  const handleShowAgentOrToolLog = useCallback((detail?: AgentLogItemWithChildren) => {
    if (!detail) {
      setAgentOrToolLogItemStack([])
      agentOrToolLogItemStackRef.current = []
      return
    }
    const { message_id: id, children } = detail
    let currentAgentOrToolLogItemStack = agentOrToolLogItemStackRef.current.slice()
    const index = currentAgentOrToolLogItemStack.findIndex(logItem => logItem.message_id === id)

    if (index > -1)
      currentAgentOrToolLogItemStack = currentAgentOrToolLogItemStack.slice(0, index + 1)
    else
      currentAgentOrToolLogItemStack = [...currentAgentOrToolLogItemStack.slice(), detail]

    setAgentOrToolLogItemStack(currentAgentOrToolLogItemStack)
    agentOrToolLogItemStackRef.current = currentAgentOrToolLogItemStack

    if (children) {
      setAgentOrToolLogListMap({
        ...agentOrToolLogListMapRef.current,
        [id]: children,
      })
    }
  }, [setAgentOrToolLogItemStack, setAgentOrToolLogListMap])

  return {
    showSpecialResultPanel: showRetryDetail || showIteratingDetail || showLoopingDetail || !!agentOrToolLogItemStack.length,
    showRetryDetail,
    setShowRetryDetailTrue,
    setShowRetryDetailFalse,
    retryResultList,
    setRetryResultList,
    handleShowRetryResultList,

    showIteratingDetail,
    setShowIteratingDetailTrue,
    setShowIteratingDetailFalse,
    iterationResultList,
    setIterationResultList,
    iterationResultDurationMap,
    setIterationResultDurationMap,
    handleShowIterationResultList,

    showLoopingDetail,
    setShowLoopingDetailTrue,
    setShowLoopingDetailFalse,
    loopResultList,
    setLoopResultList,
    loopResultDurationMap,
    setLoopResultDurationMap,
    loopResultVariableMap,
    setLoopResultVariableMap,
    handleShowLoopResultList,

    agentOrToolLogItemStack,
    agentOrToolLogListMap,
    handleShowAgentOrToolLog,
  }
}