import json
import logging
import textwrap
import uuid

import google.generativeai as genai
import gradio as gr
import PIL
import requests

from modules.presets import i18n

from ..index_func import construct_index
from ..utils import count_token
from .base_model import BaseLLMModel


class GoogleGeminiClient(BaseLLMModel):
    def __init__(self, model_name, api_key, user_name="") -> None:
        super().__init__(model_name=model_name, user=user_name)
        self.api_key = api_key
        if "vision" in model_name.lower():
            self.multimodal = True
        else:
            self.multimodal = False
        self.image_paths = []

    def _get_gemini_style_input(self):
        self.history.extend([{"role": "image", "content": i} for i in self.image_paths])
        self.image_paths = []
        messages = []
        for item in self.history:
            if item["role"] == "image":
                messages.append(PIL.Image.open(item["content"]))
            else:
                messages.append(item["content"])
        return messages

    def to_markdown(self, text):
        text = text.replace("•", "  *")
        return textwrap.indent(text, "> ", predicate=lambda _: True)

    def handle_file_upload(self, files, chatbot, language):
        if files:
            if self.multimodal:
                for file in files:
                    if file.name:
                        self.image_paths.append(file.name)
                        chatbot = chatbot + [((file.name,), None)]
                return None, chatbot, None
            else:
                construct_index(self.api_key, file_src=files)
                status = i18n("索引构建完成")
                return gr.update(), chatbot, status

    def get_answer_at_once(self):
        genai.configure(api_key=self.api_key)
        messages = self._get_gemini_style_input()
        model = genai.GenerativeModel(self.model_name)
        response = model.generate_content(messages)
        try:
            return self.to_markdown(response.text), len(response.text)
        except ValueError:
            return (
                i18n("由于下面的原因，Google 拒绝返回 Gemini 的回答：\n\n")
                + str(response.prompt_feedback),
                0,
            )

    def get_answer_stream_iter(self):
        genai.configure(api_key=self.api_key)
        messages = self._get_gemini_style_input()
        model = genai.GenerativeModel(self.model_name)
        response = model.generate_content(messages, stream=True)
        partial_text = ""
        for i in response:
            response = i.text
            partial_text += response
            yield partial_text
        self.all_token_counts[-1] = count_token(partial_text)
        yield partial_text