import { create } from 'zustand'
import { AI_OPTIONS } from './config'
import aiService from './api'
import type { AIOption, AIMessage, AIData, InitOptions } from './types'

interface AIState {
  ai: string
  options: AIOption[]
  message: AIMessage | null
  messages: AIMessage[]
  isLoading: boolean
  collapsed: boolean
  setAI: (ai: string) => void
  toggleCollapsed: () => void
  post: (data: AIData, options?: InitOptions) => Promise<AIMessage>
}

export const useAIStore = create<AIState>((set, get) => ({
  ai: localStorage.getItem('ai') || 'qwen',
  options: AI_OPTIONS,
  message: null,
  messages: [],
  isLoading: false,
  collapsed: false,
  setAI: (ai) => {
    localStorage.setItem('ai', ai)
    set({ ai })
  },
  toggleCollapsed: () => {
    set((state) => ({ collapsed: !state.collapsed }))
  },
  post: async (data, options) => {
    const { ai, messages } = get()

    // 处理用户消息（去掉 system 角色的消息）
    const userMessages = data.messages.filter((item) => item.role !== 'system')

    set({
      collapsed: true,
      isLoading: true,
      messages: [...messages, ...userMessages],
    })

    return new Promise<AIMessage>((resolve, reject) => {
      try {
        aiService.post(
          { ...data, model: ai },
          {
            onMessage: (response) => {
              set((state) => {
                const messageIndex = state.messages.findIndex((msg) => msg.id === response.id)

                if (messageIndex === -1) {
                  // 新的 AI 回复
                  const newMessage: AIMessage = {
                    id: response.id,
                    role: 'assistant',
                    content: response.content,
                    json: response.json,
                  }
                  return {
                    message: newMessage, // 存储最新的 AI 消息
                    messages: [...state.messages, newMessage], // 追加到历史消息
                  }
                } else {
                  // 追加内容到已有的消息
                  const updatedMessages = state.messages.map((msg) =>
                    msg.id === response.id
                      ? { ...msg, content: msg.content + response.content, json: msg.json || response.json }
                      : msg
                  )

                  return {
                    message: updatedMessages[messageIndex], // 更新最新的 AI 消息
                    messages: updatedMessages,
                  }
                }
              })
            },
            onerror: (err) => {
              console.error('AI 请求失败:', err)
              set({ isLoading: false })
              reject(err)
            },
            onclose: () => {
              set({ isLoading: false })
              const { message } = get()
              if (message) {
                resolve(message)
              } else {
                reject(new Error('No message received'))
              }
            },
            ...options,
          }
        )
      } catch (err) {
        console.error('AI 请求失败:', err)
        set({ isLoading: false })
        reject(err)
      }
    })
  },
}))

export type { AIMessage }
