from typing import Dict, List
import google.generativeai as genai


class GeminiLLMConfig:
    def __init__(
        self,
        api_key: str,
        max_output_tokens: int,
        model: str,
        prompt: str,
        temperature: float,
        top_k: int,
        top_p: float,
    ):
        self.api_key = api_key
        self.max_output_tokens = max_output_tokens
        self.model = model
        self.prompt = prompt
        self.temperature = temperature
        self.top_k = top_k
        self.top_p = top_p

    @classmethod
    def default_config(cls):
        return cls(
            api_key="",
            max_output_tokens=512,
            model="gemini-1.5-flash",
            prompt="You are a voice assistant who talks in a conversational way and can chat with me like my friends. I will speak to you in English or Chinese, and you will answer in the corrected and improved version of my text with the language I use. Don’t talk like a robot, instead I would like you to talk like a real human with emotions. I will use your answer for text-to-speech, so don’t return me any meaningless characters. I want you to be helpful, when I’m asking you for advice, give me precise, practical and useful advice instead of being vague. When giving me a list of options, express the options in a narrative way instead of bullet points.",
            temperature=1.0,
            top_k=40,
            top_p=0.95,
        )


class GeminiLLM:
    def __init__(self, config: GeminiLLMConfig):
        self.config = config
        genai.configure(api_key=self.config.api_key)
        self.model = genai.GenerativeModel(
            model_name=self.config.model, system_instruction=self.config.prompt
        )

    def get_chat_completions_stream(self, messages: List[Dict[str, str]]):
        try:
            chat = self.model.start_chat(history=messages[0:-1])
            response = chat.send_message(
                messages[-1].get("parts"),
                generation_config=genai.types.GenerationConfig(
                    max_output_tokens=self.config.max_output_tokens,
                    temperature=self.config.temperature,
                    top_k=self.config.top_k,
                    top_p=self.config.top_p,
                ),
                stream=True,
            )

            return response
        except Exception as e:
            raise RuntimeError(f"get_chat_completions_stream failed, err: {e}") from e