import { useCallback, useEffect, useRef } from 'react'

import { streamSSE } from '../../lib/sse'
import { useChatStreamStore } from './store'

export function useChatStream(onDone: () => Promise<void>) {
  const controllerRef = useRef<AbortController | null>(null)
  const store = useChatStreamStore()

  const abort = useCallback(() => controllerRef.current?.abort(), [])
  const cancel = useCallback(() => {
    controllerRef.current?.abort()
    useChatStreamStore.getState().fail('已停止生成，可重新发送问题。')
  }, [])

  const send = useCallback(async (sessionId: string, query: string) => {
    controllerRef.current?.abort()
    const controller = new AbortController()
    controllerRef.current = controller
    store.start(query)
    try {
      await streamSSE('/api/v1/chat/stream', { session_id: sessionId, query }, (event) => {
        if (event.event === 'status') store.setPhase(event.data.phase)
        if (event.event === 'sources') store.setSources(event.data.sources)
        if (event.event === 'token') store.appendText(event.data.content)
        if (event.event === 'error') store.fail(event.data.message)
      }, controller.signal)
      if (!controller.signal.aborted && useChatStreamStore.getState().phase !== 'error') await onDone()
    } catch (reason) {
      if (controller.signal.aborted) return
      store.fail(reason instanceof Error ? reason.message : '流式回答失败')
    } finally {
      if (controllerRef.current === controller) controllerRef.current = null
    }
  }, [onDone, store])

  useEffect(() => abort, [abort])
  return { send, cancel }
}
