from google import genai
from .memory import Memory
from .stats import Stats
from .exceptions import InvalidAPIKeyException, ChatException
from .validators import (
    validate_api_key,
    validate_filepath,
    validate_max_output_tokens,
    validate_message,
    validate_prompt,
    validate_temperature,
    validate_max_messages,
)


class Dracula:
    def __init__(
        self,
        api_key: str,
        model: str = "gemini-3-flash-preview",
        max_messages: int = 10,
        prompt: str = "You are a helpful asisstant.",
        temperature: float = 1.0,
        max_output_tokens: int = 1024,
        stats_filepath: str = "dracula_stats.json",
    ):
        validate_api_key(api_key)
        validate_temperature(temperature)
        validate_max_output_tokens(max_output_tokens)
        validate_max_messages(max_messages)
        validate_prompt(prompt)
        validate_filepath(stats_filepath)

        try:
            self.client = genai.Client(api_key=api_key)
            self.model_name = model

        except Exception:
            raise InvalidAPIKeyException("Invalid API key or model name.")

        self.prompt = prompt
        self.temperature = temperature
        self.max_output_tokens = max_output_tokens
        self.memory = Memory(max_messages=max_messages)
        self.stats = Stats(filepath=stats_filepath)

    def chat(self, message: str) -> str:
        validate_message(message)

        try:
            self.memory.add_message("user", message)
            self.stats.record_message(message)

            history = [
                genai.types.Content(
                    role=msg["role"], parts=[genai.types.Part(text=msg["content"])]
                )
                for msg in self.memory.get_history()
            ]

            response = self.client.models.generate_content(
                model=self.model_name,
                contents=history,
                config=genai.types.GenerateContentConfig(
                    system_instruction=self.prompt,
                    temperature=self.temperature,
                    max_output_tokens=self.max_output_tokens,
                ),
            )

            reply = response.text

            self.memory.add_message("model", reply)
            self.stats.record_response(reply)
            return reply

        except Exception as e:
            raise ChatException(f"Something went wrong: {str(e)}")

    def stream(self, message: str):
        validate_message(message)

        try:
            self.memory.add_message("user", message)
            self.stats.record_message(message)

            history = [
                genai.types.Content(
                    role=msg["role"], parts=[genai.types.Part(text=msg["content"])]
                )
                for msg in self.memory.get_history()
            ]

            response = self.client.models.generate_content_stream(
                model=self.model_name,
                contents=history,
                config=genai.types.GenerateContentConfig(
                    system_instruction=self.prompt,
                    temperature=self.temperature,
                    max_output_tokens=self.max_output_tokens,
                ),
            )

            full_reply = ""
            for chunk in response:
                text = chunk.text
                full_reply += text
                yield text

            self.memory.add_message("model", full_reply)
            self.stats.record_response(full_reply)

        except Exception as e:
            raise ChatException(f"Something went wrong: {str(e)}")

    def get_stats(self) -> dict:
        return self.stats.get_stats()

    def reset_stats(self):
        self.stats.reset()

    def save_history(self, filepath: str):
        validate_filepath(filepath)
        self.memory.save(filepath)

    def load_history(self, filepath: str):
        validate_filepath(filepath)
        self.memory.load(filepath)

    def clear_memory(self):
        self.memory.clear()

    def set_prompt(self, prompt: str):
        validate_prompt(prompt)
        self.prompt = prompt

    def set_temperature(self, temperature: float):
        validate_temperature(temperature)
        self.temperature = temperature

    def set_max_output_tokens(self, max_output_tokens: int):
        validate_max_output_tokens(max_output_tokens)
        self.max_output_tokens = max_output_tokens

    def get_history(self):
        return self.memory.get_history()
