import { create } from 'zustand'
import { AI_OPTIONS } from './config'
import aiService from './api'
import type { AIOption, AIMessage, AIData } 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: AIData) => 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.messages.filter((item) => item.role !== 'system')],
    })

    return new Promise<AIMessage>((resolve, reject) => {
      aiService.post(
        { model: ai, ...data },
        {
          signal: controller.signal,
          onUpdate: (message) => {
            console.log(message)
            set((state) => {
              const messageIndex = state.messages.findIndex((msg) => msg.id === message.id)

              if (messageIndex === -1) {
                return { message, messages: [...state.messages, message] }
              } else {
                const updatedMessages = state.messages.map((msg) =>
                  msg.id === message.id
                    ? {
                        ...msg,
                        content: msg.content + message.content,
                        reasoning_content: msg.reasoning_content + message.reasoning_content,
                      }
                    : msg
                )

                return { message: updatedMessages[messageIndex], messages: updatedMessages }
              }
            })
          },
          onSuccess: (message) => {
            resolve(message)
            set({ isLoading: false })
          },
          onError: (err) => {
            reject(err)
            set({ isLoading: false })
            set({ error: err.message })
          },
        }
      )
    })
  },
}))

export type { AIMessage }
