import { Knowledge } from "@/db/knowledge"
import { ChatDocuments } from "@/models/ChatTypes"
import { create } from "zustand"
import { type UploadedFile } from "@/db/dexie/types"
import { isFireFoxPrivateMode } from "@/utils/is-private-mode"
import {
  ChatActionInfo,
  ChatMessageKind,
  McpPendingApprovalRequest,
  McpToolCall
} from "@/libs/mcp/types"

type WebSearch = {
  search_engine: string
  search_url: string
  search_query: string
  search_results: {
    title: string
    link: string
  }[]
}
export type Message = {
  isBot: boolean
  name: string
  message: string
  sources: any[]
  images?: string[]
  search?: WebSearch
  reasoning_time_taken?: number
  id?: string
  messageType?: string
  modelName?: string
  modelImage?: string
  documents?: ChatDocuments
  generationInfo?: any
  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 PendingMcpApproval = McpPendingApprovalRequest & {
  approve: () => void
  reject: (reason?: string) => void
}

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"
  setChatMode: (chatMode: "normal" | "rag") => void
  isEmbedding: boolean
  setIsEmbedding: (isEmbedding: boolean) => void
  webSearch: boolean
  setWebSearch: (webSearch: boolean) => void
  pageAction: boolean
  setPageAction: (pageAction: boolean) => void
  isSearchingInternet: boolean
  setIsSearchingInternet: (isSearchingInternet: boolean) => void

  selectedSystemPrompt: string | null
  setSelectedSystemPrompt: (selectedSystemPrompt: string) => void

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

  selectedKnowledge: Knowledge | null
  setSelectedKnowledge: (selectedKnowledge: Knowledge) => void

  setSpeechToTextLanguage: (language: string) => void
  speechToTextLanguage: string

  temporaryChat: boolean
  setTemporaryChat: (temporaryChat: boolean) => void

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

  documentContext: ChatDocuments | null
  setDocumentContext: (documentContext: ChatDocuments) => void

  uploadedFiles: UploadedFile[]
  setUploadedFiles: (uploadedFiles: UploadedFile[]) => void

  contextFiles: UploadedFile[]
  setContextFiles: (contextFiles: UploadedFile[]) => void

  actionInfo: ChatActionInfo | null
  setActionInfo: (actionInfo: ChatActionInfo | null) => void

  pendingMcpApproval: PendingMcpApproval | null
  setPendingMcpApproval: (pendingMcpApproval: PendingMcpApproval | null) => void

  fileRetrievalEnabled: boolean
  setFileRetrievalEnabled: (fileRetrievalEnabled: boolean) => void
}

export const useStoreMessageOption = 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 }),
  speechToTextLanguage: "en-US",
  setSpeechToTextLanguage: (language) =>
    set({ speechToTextLanguage: language }),
  selectedModel: null,
  setSelectedModel: (selectedModel) => set({ selectedModel }),
  chatMode: "normal",
  setChatMode: (chatMode) => set({ chatMode }),
  isEmbedding: false,
  setIsEmbedding: (isEmbedding) => set({ isEmbedding }),
  webSearch: false,
  setWebSearch: (webSearch) => set({ webSearch }),
  pageAction: false,
  setPageAction: (pageAction) => set({ pageAction }),
  isSearchingInternet: false,
  setIsSearchingInternet: (isSearchingInternet) => set({ isSearchingInternet }),
  selectedSystemPrompt: null,
  setSelectedSystemPrompt: (selectedSystemPrompt) =>
    set({ selectedSystemPrompt }),
  selectedQuickPrompt: null,
  setSelectedQuickPrompt: (selectedQuickPrompt) => set({ selectedQuickPrompt }),

  selectedKnowledge: null,
  setSelectedKnowledge: (selectedKnowledge) => set({ selectedKnowledge }),

  temporaryChat: isFireFoxPrivateMode,
  setTemporaryChat: (temporaryChat) => set({ temporaryChat }),

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

  documentContext: null,
  setDocumentContext: (documentContext) => set({ documentContext }),

  uploadedFiles: [],
  setUploadedFiles: (uploadedFiles) => set({ uploadedFiles }),
  contextFiles: [],
  setContextFiles: (contextFiles) => set({ contextFiles }),

  actionInfo: null,
  setActionInfo: (actionInfo) => set({ actionInfo }),

  pendingMcpApproval: null,
  setPendingMcpApproval: (pendingMcpApproval) => set({ pendingMcpApproval }),

  fileRetrievalEnabled: false,
  setFileRetrievalEnabled: (fileRetrievalEnabled) =>
    set({ fileRetrievalEnabled })
}))
