import { useState, useEffect, useCallback, useRef } from 'react'
import { AI_OPTIONS } from './config'
import aiService from './api'
import type { AIMessage, AIData, InitOptions } from './types'

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)

  // 使用 useRef 存储最新消息，避免闭包问题
  const latestMessageRef = useRef<AIMessage | null>(null)
  // 使用 useRef 存储消息列表，避免异步更新问题
  const messagesRef = useRef<AIMessage[]>([])

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

  // 更新消息的辅助函数
  const updateMessage = useCallback((response: any, isNew: boolean) => {
    if (isNew) {
      // 新消息
      const newMessage: AIMessage = {
        id: response.id,
        role: 'assistant',
        content: response.content,
        json: response.json,
      }

      latestMessageRef.current = newMessage
      messagesRef.current = [...messagesRef.current, newMessage]

      // 更新状态
      setMessage(newMessage)
      setMessages(messagesRef.current)
    } else {
      // 更新现有消息
      messagesRef.current = messagesRef.current.map((msg) => {
        if (msg.id === response.id) {
          const updatedMessage = {
            ...msg,
            content: msg.content + response.content,
            json: msg.json || response.json,
          }
          latestMessageRef.current = updatedMessage
          setMessage(updatedMessage)
          return updatedMessage
        }
        return msg
      })

      setMessages([...messagesRef.current])
    }
  }, [])

  const post = useCallback(
    async (data: AIData, options?: InitOptions) => {
      return new Promise<AIMessage>((resolve, reject) => {
        try {
          setIsLoading(true)

          // 添加用户消息
          const userMessages = data.messages.filter((item) => item.role !== 'system')
          messagesRef.current = [...messagesRef.current, ...userMessages]
          setMessages((prev) => [...prev, ...userMessages])

          aiService.post(
            { ...data, model: ai },
            {
              onMessage: (response) => {
                const messageIndex = messagesRef.current.findIndex((msg) => msg.id === response.id)
                updateMessage(response, messageIndex === -1)
              },
              onerror: (error) => {
                console.error('AI service error:', error)
                setIsLoading(false)
                reject(error)
              },
              onclose: () => {
                setIsLoading(false)
                if (latestMessageRef.current) {
                  resolve(latestMessageRef.current)
                } else {
                  reject(new Error('No message received'))
                }
              },
              ...globalOptions,
              ...options,
            }
          )
        } catch (error) {
          console.error('Post error:', error)
          setIsLoading(false)
          reject(error)
        }
      })
    },
    [ai, updateMessage, globalOptions]
  )

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