import { create } from "zustand"
import { ChatMessageKind, McpToolCall } from "@/libs/mcp/types"

export type Message = {
  isBot: boolean
  name: string
  message: string
  sources: any[]
  images?: string[]
  modelName?: string
  modelImage?: string
  id?: string
  messageType?: string
  generationInfo?: any
  reasoning_time_taken?: number
  messageKind?: ChatMessageKind
  toolCalls?: McpToolCall[]
  toolCallId?: string
  toolName?: string
  toolServerName?: string
  toolError?: boolean
}

export type ChatHistory = {
  role: "user" | "assistant" | "system" | "tool"
  content: string
  image?: string
  images?: string[]
  messageType?: string
  messageKind?: ChatMessageKind
  toolCalls?: McpToolCall[]
  toolCallId?: string
  toolName?: string
  toolServerName?: string
  toolError?: boolean
}[]

type State = {
  messages: Message[]
  setMessages: (messages: Message[]) => void
  history: ChatHistory
  setHistory: (history: ChatHistory) => void
  streaming: boolean
  setStreaming: (streaming: boolean) => void
  isFirstMessage: boolean
  setIsFirstMessage: (isFirstMessage: boolean) => void
  historyId: string | null
  setHistoryId: (history_id: string | null) => void
  isLoading: boolean
  setIsLoading: (isLoading: boolean) => void
  isProcessing: boolean
  setIsProcessing: (isProcessing: boolean) => void
  selectedModel: string | null
  setSelectedModel: (selectedModel: string) => void
  chatMode: "normal" | "rag" | "vision"
  setChatMode: (chatMode: "normal" | "rag" | "vision") => void
  isEmbedding: boolean
  setIsEmbedding: (isEmbedding: boolean) => void
  speechToTextLanguage: string
  setSpeechToTextLanguage: (speechToTextLanguage: string) => void
  currentURL: string
  setCurrentURL: (currentURL: string) => void
  selectedSystemPrompt: string | null
  setSelectedSystemPrompt: (selectedSystemPrompt: string) => void

  selectedQuickPrompt: string | null
  setSelectedQuickPrompt: (selectedQuickPrompt: string) => void

  useOCR: boolean
  setUseOCR: (useOCR: boolean) => void
}

export const useStoreMessage = create<State>((set) => ({
  messages: [],
  setMessages: (messages) => set({ messages }),
  history: [],
  setHistory: (history) => set({ history }),
  streaming: false,
  setStreaming: (streaming) => set({ streaming }),
  isFirstMessage: true,
  setIsFirstMessage: (isFirstMessage) => set({ isFirstMessage }),
  historyId: null,
  setHistoryId: (historyId) => set({ historyId }),
  isLoading: false,
  setIsLoading: (isLoading) => set({ isLoading }),
  isProcessing: false,
  setIsProcessing: (isProcessing) => set({ isProcessing }),
  defaultSpeechToTextLanguage: "en-US",
  selectedModel: null,
  setSelectedModel: (selectedModel) => set({ selectedModel }),
  chatMode: "normal",
  setChatMode: (chatMode) => set({ chatMode }),
  isEmbedding: false,
  setIsEmbedding: (isEmbedding) => set({ isEmbedding }),
  speechToTextLanguage: "en-US",
  setSpeechToTextLanguage: (speechToTextLanguage) =>
    set({ speechToTextLanguage }),
  currentURL: "",
  setCurrentURL: (currentURL) => set({ currentURL }),

  selectedSystemPrompt: null,
  setSelectedSystemPrompt: (selectedSystemPrompt) =>
    set({ selectedSystemPrompt }),
  selectedQuickPrompt: null,
  setSelectedQuickPrompt: (selectedQuickPrompt) => set({ selectedQuickPrompt }),

  useOCR: false,
  setUseOCR: (useOCR) => set({ useOCR })
}))
