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

interface AIState {
  ai: string
  options: AIOption[]
  message: AIMessage | null
  messages: AIMessage[]
  isLoading: boolean
  collapsed: boolean
  error: string | null
  controller: AbortController | null
  setAI: (ai: string) => void
  toggleCollapsed: () => void
  post: (data: AIRequestData) => Promise<AIMessage>
  clearError: () => void
}

export const useAIStore = create<AIState>((set, get) => ({
  ai: localStorage.getItem('ai') || 'Pro/deepseek-ai/DeepSeek-R1',
  options: AI_OPTIONS,
  message: null,
  messages: [],
  isLoading: false,
  collapsed: false,
  error: null,
  controller: null,
  setAI: (ai) => {
    localStorage.setItem('ai', ai)
    set({ ai })
  },
  toggleCollapsed: () => {
    set((state) => ({ collapsed: !state.collapsed }))
  },
  clearError: () => {
    set({ error: null })
  },
  post: async (data) => {
    const { ai, messages } = get()
    const controller = new AbortController()
    set({
      controller,
      collapsed: true,
      isLoading: true,
      error: null,
      messages: [
        ...messages,
        ...(data as ChatRequestData).messages.filter((item: AIMessage) => item.role !== 'system'),
      ],
    })
    const callback = (message: AIMessage) => {
      set((state) => {
        const messageIndex = state.messages.findIndex((item) => item.id === message.id)
        const messages =
          messageIndex === -1
            ? [...state.messages, message]
            : state.messages.map((msg) => (msg.id === message.id ? message : msg))
        return { message, messages }
      })
    }
    return new Promise<AIMessage>((resolve, reject) => {
      aiService.post(
        { model: ai, ...data },
        {
          signal: controller.signal,
          onUpdate: (message) => {
            callback(message)
          },
          onSuccess: (message) => {
            callback(message)
            resolve(message)
            set({ isLoading: false, controller: null })
          },
          onError: (err) => {
            reject(err)
            set({ isLoading: false, error: err.message, controller: null })
          },
        }
      )
    })
  },
}))

export type { AIMessage }
