import { useState, useEffect, useCallback } from 'react'
import aiService, { AIMessage, AIData, AI_OPTIONS, InitOptions } from '@/utils/ai'

export function useAI(globalOptions?: InitOptions) {
  const [ai, setAI] = useState<string>(localStorage.getItem('ai') || 'qwen')
  const [messages, setMessages] = useState<AIMessage[]>([])
  const [message, setMessage] = useState<AIMessage | null>(null) // 存储最新的消息
  const [isLoading, setIsLoading] = useState<boolean>(false)

  useEffect(() => {
    localStorage.setItem('ai', ai)
  }, [ai])

  const post = useCallback(
    async (data: AIData, options?: InitOptions) => {
      setIsLoading(true)

      // 先添加用户消息（去掉 system 角色的消息）
      const userMessages = data.messages.filter((item) => item.role !== 'system')
      setMessages((prevMessages) => [...prevMessages, ...userMessages])

      let latestMessage: AIMessage | null = null

      const onComplete = options?.onComplete || globalOptions?.onComplete

      await aiService.post(ai, data, {
        onMessage: (response) => {
          setMessages((prevMessages) => {
            const messageIndex = prevMessages.findIndex((msg) => msg.id === response.id)
            if (messageIndex === -1) {
              latestMessage = { id: response.id, role: 'assistant', content: response.content }
              setMessage(latestMessage) // 更新最新消息
              return [...prevMessages, latestMessage]
            } else {
              return prevMessages.map((msg) => {
                if (msg.id === response.id) {
                  const updatedMessage = { ...msg, content: msg.content + response.content }
                  latestMessage = updatedMessage
                  setMessage(updatedMessage) // 更新最新消息
                  return updatedMessage
                }
                return msg
              })
            }
          })
        },
        onerror: () => {
          setIsLoading(false)
        },
        onclose: () => {
          setIsLoading(false)
          if (latestMessage && onComplete) onComplete(latestMessage)
        },
        ...globalOptions,
        ...options,
      })
    },
    [ai]
  )

  return { ai, setAI, options: AI_OPTIONS, post, messages, message, isLoading }
}
