import { cleanUrl } from "~/libs/clean-url"
import {
  geWebSearchFollowUpPrompt,
  getOllamaURL
} from "~/services/ollama"
import { type ChatHistory, type Message } from "~/store/option"
import { generateID } from "@/db/dexie/helpers"
import { getSystemPromptForWeb, isQueryHaveWebsite } from "~/web/web"
import { generateHistory } from "@/utils/generate-history"
import { pageAssistModel } from "@/models"
import { humanMessageFormatter } from "@/utils/human-message"
import {
  isReasoningEnded,
  isReasoningStarted,
  mergeReasoningContent,
  removeReasoning
} from "@/libs/reasoning"
import { getModelNicknameByID } from "@/db/dexie/nickname"
import { systemPromptFormatter } from "@/utils/system-message"

export const searchChatMode = async (
  message: string,
  image: string,
  isRegenerate: boolean,
  messages: Message[],
  history: ChatHistory,
  signal: AbortSignal,
  {
    selectedModel,
    useOCR,
    setMessages,
    setIsSearchingInternet,
    saveMessageOnSuccess,
    saveMessageOnError,
    setHistory,
    setIsProcessing,
    setStreaming,
    setAbortController,
    historyId,
    setHistoryId,
    images
  }: {
    selectedModel: string
    useOCR: boolean
    setMessages: (messages: Message[] | ((prev: Message[]) => Message[])) => void
    setIsSearchingInternet: (value: boolean) => void
    saveMessageOnSuccess: (data: any) => Promise<string | null>
    saveMessageOnError: (data: any) => Promise<string | null>
    setHistory: (history: ChatHistory) => void
    setIsProcessing: (value: boolean) => void
    setStreaming: (value: boolean) => void
    setAbortController: (controller: AbortController | null) => void
    historyId: string | null
    setHistoryId: (id: string) => void
    images?: string[]
  }
) => {
  console.log("Using searchChatMode")
  const url = await getOllamaURL()

  // Process images array for base64 formatting
  const processedImages = (images || []).map((img) => {
    if (img.length > 0 && !img.startsWith("data:image")) {
      return `data:image/jpeg;base64,${img.split(",")[1]}`
    }
    return img
  })

  // Keep backward compatibility with single image
  if (image.length > 0) {
    image = `data:image/jpeg;base64,${image.split(",")[1]}`
  }

  const ollama = await pageAssistModel({
    model: selectedModel!,
    baseUrl: cleanUrl(url)
  })

  let newMessage: Message[] = []
  let generateMessageId = generateID()

  const modelInfo = await getModelNicknameByID(selectedModel)
  if (!isRegenerate) {
    // Use images array if available, otherwise fall back to single image
    const userImages = processedImages.length > 0 ? processedImages : (image ? [image] : [])

    newMessage = [
      ...messages,
      {
        isBot: false,
        name: "You",
        message,
        sources: [],
        images: userImages
      },
      {
        isBot: true,
        name: selectedModel,
        message: "▋",
        sources: [],
        id: generateMessageId,
        modelImage: modelInfo?.model_avatar,
        modelName: modelInfo?.model_name || selectedModel
      }
    ]
  } else {
    newMessage = [
      ...messages,
      {
        isBot: true,
        name: selectedModel,
        message: "▋",
        sources: [],
        id: generateMessageId,
        modelImage: modelInfo?.model_avatar,
        modelName: modelInfo?.model_name || selectedModel
      }
    ]
  }
  setMessages(newMessage)
  let fullText = ""
  let contentToSave = ""
  let timetaken = 0

  try {
    setIsSearchingInternet(true)

    let query = message

    let questionPrompt = await geWebSearchFollowUpPrompt()
    const lastTenMessages = newMessage.slice(-10)
    lastTenMessages.pop()
    const chat_history = lastTenMessages
      .map((message) => {
        return `${message.isBot ? "Assistant: " : "Human: "}${message.message}`
      })
      .join("\n")
    const promptForQuestion = questionPrompt
      .replaceAll("{chat_history}", chat_history)
      .replaceAll("{question}", message)
    const questionModel = await pageAssistModel({
      model: selectedModel!,
      baseUrl: cleanUrl(url)
    })

    // Build content array with text and images for question message
    const questionContentArray: any[] = [
      {
        text: promptForQuestion,
        type: "text"
      }
    ]

    // Add all images to content array (use processedImages if available, otherwise single image)
    const imagesToUse = processedImages.length > 0 ? processedImages : (image.length > 0 ? [image] : [])

    imagesToUse.forEach((img) => {
      questionContentArray.push({
        image_url: img,
        type: "image_url"
      })
    })

    let questionMessage = await humanMessageFormatter({
      content: questionContentArray,
      model: selectedModel,
      useOCR: useOCR
    })
    try {
      const isWebQuery = await isQueryHaveWebsite(query)
      if (!isWebQuery) {
        const response = await questionModel.invoke([questionMessage])
        query = response?.content?.toString() || message
        query = removeReasoning(query)
      }
    } catch (error) {
      console.error("Error in questionModel.invoke:", error)
    }
    // }

    const { prompt, source } = await getSystemPromptForWeb(query)
    setIsSearchingInternet(false)

    // Build content array with text and all images
    const contentArray: any[] = [
      {
        text: message,
        type: "text"
      }
    ]

    imagesToUse.forEach((img) => {
      contentArray.push({
        image_url: img,
        type: "image_url"
      })
    })

    let humanMessage = await humanMessageFormatter({
      content: contentArray,
      model: selectedModel,
      useOCR: useOCR
    })

    const applicationChatHistory = generateHistory(history, selectedModel)

    if (prompt) {
      applicationChatHistory.unshift(
        await systemPromptFormatter({
          content: prompt
        })
      )
    }

    let generationInfo: any | undefined = undefined

    const chunks = await ollama.stream(
      [...applicationChatHistory, humanMessage],
      {
        signal: signal,
        callbacks: [
          {
            handleLLMEnd(output: any): any {
              try {
                generationInfo = output?.generations?.[0][0]?.generationInfo
              } catch (e) {
                console.error("handleLLMEnd error", e)
              }
            }
          }
        ]
      }
    )
    let count = 0
    let reasoningStartTime: Date | undefined = undefined
    let reasoningEndTime: Date | undefined = undefined
    let apiReasoning = false
    for await (const chunk of chunks) {
      if (chunk?.additional_kwargs?.reasoning_content) {
        const reasoningContent = mergeReasoningContent(
          fullText,
          chunk?.additional_kwargs?.reasoning_content || ""
        )
        contentToSave = reasoningContent
        fullText = reasoningContent
        apiReasoning = true
      }

      if (apiReasoning && chunk?.content) {
        fullText += "</think>"
        contentToSave += "</think>"
        apiReasoning = false
      }

      contentToSave += chunk?.content
      fullText += chunk?.content
      if (count === 0) {
        setIsProcessing(true)
      }
      if (isReasoningStarted(fullText) && !reasoningStartTime) {
        reasoningStartTime = new Date()
      }

      if (
        reasoningStartTime &&
        !reasoningEndTime &&
        isReasoningEnded(fullText)
      ) {
        reasoningEndTime = new Date()
        const reasoningTime =
          reasoningEndTime.getTime() - reasoningStartTime.getTime()
        timetaken = reasoningTime
      }
      setMessages((prev) => {
        return prev.map((message) => {
          if (message.id === generateMessageId) {
            return {
              ...message,
              message: fullText + "▋",
              reasoning_time_taken: timetaken
            }
          }
          return message
        })
      })
      count++
    }
    // update the message with the full text
    setMessages((prev) => {
      return prev.map((message) => {
        if (message.id === generateMessageId) {
          return {
            ...message,
            message: fullText,
            sources: source,
            generationInfo,
            reasoning_time_taken: timetaken
          }
        }
        return message
      })
    })

    const imagesToSave = processedImages.length > 0 ? processedImages : (image ? [image] : [])

    setHistory([
      ...history,
      {
        role: "user",
        content: message,
        image: imagesToSave.length > 0 ? imagesToSave[0] : undefined,
        images: imagesToSave.length > 0 ? imagesToSave : undefined
      },
      {
        role: "assistant",
        content: fullText
      }
    ])

    await saveMessageOnSuccess({
      historyId,
      setHistoryId,
      isRegenerate,
      selectedModel: selectedModel,
      message,
      image: imagesToSave.length > 0 ? imagesToSave[0] : "",
      images: imagesToSave,
      fullText,
      source,
      generationInfo,
      reasoning_time_taken: timetaken
    })

    setIsProcessing(false)
    setStreaming(false)
  } catch (e) {
    const imagesToSave = processedImages.length > 0 ? processedImages : (image ? [image] : [])

    const errorSave = await saveMessageOnError({
      e,
      botMessage: fullText,
      history,
      historyId,
      image: imagesToSave.length > 0 ? imagesToSave[0] : "",
      images: imagesToSave,
      selectedModel,
      setHistory,
      setHistoryId,
      userMessage: message,
      isRegenerating: isRegenerate
    })

    if (!errorSave) {
      throw e // Re-throw to be handled by the calling function
    }
    setIsProcessing(false)
    setStreaming(false)
  } finally {
    setAbortController(null)
  }
}
