from google import genai
from .memory import Memory
from .tools import Tool, ToolRegistry, ToolResult
from .stats import Stats
from .models import GeminiModel
from ._version import __version__
from .version_checker import check_latest_version
from .logger import (
    get_logger,
    enable_logging,
    disable_logging,
    enable_file_logging,
    disable_file_logging,
)
from .personas import PERSONAS
from .exceptions import InvalidAPIKeyException, ChatException, PersonaException
from .validators import (
    validate_api_key,
    validate_filepath,
    validate_max_output_tokens,
    validate_message,
    validate_prompt,
    validate_temperature,
    validate_max_messages,
    validate_language,
    validate_model,
)

logger = get_logger(f"dracula.{__name__}")


class Dracula:
    """
    The main class for interacting with Google Gemini using the Dracula library.

    Dracula provides a simple and elegant interface for chatting with Gemini,
    managing conversation memory, controlling AI behavior, and tracking usage stats.

    Args:
        api_key (str): Your Google Gemini API key.
        model (str): The Gemini model to use. Defaults to 'gemini-2.0-flash'.
        max_messages (int): Maximum number of messages to keep in memory. Defaults to 10.
        prompt (str): System prompt that defines the AI's personality. Defaults to 'You are a helpful assistant.'
        temperature (float): Controls response creativity between 0.0 and 2.0. Defaults to 1.0.
        max_output_tokens (int): Maximum length of responses in tokens. Defaults to 8192.
        stats_filepath (str): Path to save usage stats. Defaults to 'dracula_stats.json'.
        language (str): Language for responses. Defaults to 'English'.

    Example:
        >>> from dracula import Dracula
        >>> ai = Dracula(api_key="your-api-key", language="Turkish")
        >>> print(ai.chat("Hello!"))
    """

    def __init__(
        self,
        api_key: str,
        model: GeminiModel | str = GeminiModel.FLASH,
        max_messages: int = 10,
        prompt: str = "You are a helpful asisstant.",
        temperature: float = 1.0,
        max_output_tokens: int = 8192,
        stats_filepath: str = "dracula_stats.json",
        language: str = "English",
        logging: bool = False,
        log_file: str = None,
        log_level: str = "Debug",
        log_max_bytes: int = 5 * 1024 * 1024,
        log_backup_count: int = 5,
        tools: list = None,
    ):
        validate_api_key(api_key)
        validate_model(model)
        validate_temperature(temperature)
        validate_max_output_tokens(max_output_tokens)
        validate_max_messages(max_messages)
        validate_prompt(prompt)
        validate_filepath(stats_filepath)

        self.tool_registry = ToolRegistry(tools or [])

        if logging:
            enable_logging(log_level)
        else:
            disable_logging()

        if log_file:
            enable_file_logging(
                log_file, max_bytes=log_max_bytes, backup_count=log_backup_count
            )

        try:
            self.client = genai.Client(api_key=api_key)
            self.model_name = model.value if isinstance(model, GeminiModel) else model

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

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

        check_latest_version(__version__)
        logger.debug(
            f"Dracula initialized with model={model}, language={language}, temperature={temperature}"
        )

        self._create_chat_session()

    def __enter__(self):
        """
        Enter the context manager.

        Returns:
            Dracula: The current instance.

        Example:
            >>> with Dracula(api_key="your-api-key") as ai:
            ...     ai.chat("Hello!")
        """

        return self

    def __exit__(self, exc_type, exc, tb):
        """
        Exit the context manager and clean up.

        Automatically clears memory and resets stats when the with block ends.

        Returns:
            bool: False so exceptions are not suppressed.
        """

        self.clear_memory()
        self.reset_stats()
        return False

    def chat(self, message: str, auto_call: bool = True) -> str | ToolResult:
        """
        Send a message to Gemini and get a response.

        If tools are registered, Gemini may decide to call one of them.
        Set auto_call=True to let Dracula handle the tool call automatically,
        or auto_call=False to handle it yourself.

        Args:
            message (str): The message to send to Gemini.
            auto_call (bool): Whether to automatically call tools. Defaults to True.

        Returns:
            str: Gemini's response (when auto_call=True or no tool call needed)
            ToolResult: Tool call information (when auto_call=False)

        Example:
            >>> response = ai.chat("What's the weather in Istanbul?")
            >>> print(response)

            >>> result = ai.chat("What's the weather?", auto_call=False)
            >>> if result.requires_tool_call:
            ...     print(result.tool_name)
            ...     print(result.tool_args)
        """

        validate_message(message)

        try:
            self.stats.record_message(message)
            logger.debug(f"Sending message: {message[:50]}...")

            response = self._chat_session.send_message(message)

            # Loop to handle multiple tool calls
            while True:
                tool_call_found = False

                for part in response.candidates[0].content.parts:
                    if hasattr(part, "function_call") and part.function_call:
                        tool_call_found = True
                        tool_name = part.function_call.name
                        tool_args = dict(part.function_call.args)

                        logger.debug(f"Tool call requested: {tool_name}({tool_args})")

                        if not auto_call:
                            return ToolResult(
                                requires_tool_call=True,
                                tool_name=tool_name,
                                tool_args=tool_args,
                                text=None,
                            )

                        # Auto call the tool
                        tool = self.tool_registry.get(tool_name)
                        tool_result = tool.call(**tool_args)
                        logger.debug(f"Tool '{tool_name}' returned: {tool_result}")

                        # Send tool result back and get next response
                        response = self._chat_session.send_message(
                            genai.types.Part(
                                function_response=genai.types.FunctionResponse(
                                    name=tool_name,
                                    response={"result": str(tool_result)},
                                )
                            )
                        )
                        break  # Check the new response from the top of the loop

                # No tool call found — we have a text response
                if not tool_call_found:
                    reply = response.text
                    if not reply:
                        raise ChatException("Received an empty response from Gemini.")

                    self.memory.add_message("user", message)
                    self.memory.add_message("model", reply)
                    self.stats.record_response(reply)
                    logger.debug(f"Received response: {reply[:50]}...")
                    return reply

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

    def stream(self, message: str):
        """
        Send a message and receive the response word by word as it is generated.

        Instead of waiting for the full response, this method yields small chunks
        of text as they arrive, creating a real-time streaming effect similar to ChatGPT.

        Args:
            message (str): The message to send to Gemini.

        Yields:
            str: Small chunks of the response text.

        Raises:
            ValidationException: If the message is empty.
            ChatException: If something goes wrong during the request.

        Example:
            >>> for chunk in ai.stream("Tell me a long story."):
            ...     print(chunk, end="", flush=True)
        """

        validate_message(message)

        try:
            self.stats.record_message(message)

            full_reply = ""
            for chunk in self._chat_session.send_message_stream(message):
                if chunk.text:
                    full_reply += chunk.text
                    yield chunk.text

            self.memory.add_message("user", message)
            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 the current usage statistics.

        Returns:
            dict: A dictionary containing total_messages, total_responses,
                total_characters_sent, and total_characters_received.

        Example:
            >>> print(ai.get_stats())
        """

        return self.stats.get_stats()

    def reset_stats(self):
        """
        Reset all usage statistics back to zero.

        Example:
            >>> ai.reset_stats()
        """

        self.stats.reset()

    def print_history(self):
        """
        Print the conversation history in a clean, human-readable format.

        Example:
            >>> ai.print_history()
        """

        history = self.memory.get_history()

        if not history:
            print("No conversaiton history yet.")
            return

        print("\n" + "=" * 150)
        print("         🧛 DRACULA CONVERSATION HISTORY")
        print("=" * 150)

        for i, msg in enumerate(history, 1):
            if msg["role"] == "user":
                print(f"\n[{i}] 👤 You:")
            else:
                print(f"\n[{i}] 🤖 Gemini:")

            print(f"    {msg['content']}")

        print("\n" + "═" * 150 + "\n")

    def save_history(self, filepath: str):
        """
        Save the conversation history to a JSON file.

        Args:
            filepath (str): Path to the file where history will be saved.

        Raises:
            ValidationException: If the filepath is empty.
            ChatException: If the file cannot be saved.

        Example:
            >>> ai.save_history("conversation.json")
        """

        validate_filepath(filepath)
        self.memory.save(filepath)

    def load_history(self, filepath: str):
        """
        Load conversation history from a JSON file.

        Args:
            filepath (str): Path to the file to load history from.

        Raises:
            ValidationException: If the filepath is empty.
            ChatException: If the file cannot be loaded or does not exist.

        Example:
            >>> ai.load_history("conversation.json")
        """

        validate_filepath(filepath)
        self.memory.load(filepath)

    def clear_memory(self):
        """
        Clear the conversation history.

        After calling this, Gemini will have no memory of previous messages.

        Example:
            >>> ai.clear_memory()
        """

        self.memory.clear()
        self._create_chat_session()
        logger.info("Memory cleared.")

    def set_model(self, model: GeminiModel | str) -> "Dracula":
        """
        Change the Gemini model.

        Args:
            model (GeminiModel | str): The model to use.
                                    Use GeminiModel enum for reliability.

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> from dracula import GeminiModel
            >>> ai.set_model(GeminiModel.PRO)
        """

        validate_model(model)
        self.model_name = model.value if isinstance(model, GeminiModel) else model
        self._create_chat_session()
        logger.info(f"Model changed to '{self.model_name}'")
        return self

    def set_prompt(self, prompt: str) -> "Dracula":
        """
        Change the system prompt and clear conversation memory.

        Args:
            prompt (str): The new system prompt.

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> ai.set_prompt("You are a pirate.")
        """

        validate_prompt(prompt)
        self.prompt = prompt
        self.memory.clear()
        logger.info("Prompt updated.")
        self._create_chat_session()
        return self

    def set_temperature(self, temperature: float) -> "Dracula":
        """
        Change the temperature setting.

        Temperature controls how creative and random responses are.
        Lower values (close to 0.0) produce focused, predictable responses.
        Higher values (close to 2.0) produce creative, varied responses.

        Args:
            temperature (float): A value between 0.0 and 2.0.

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> ai.set_temperature(0.5)
        """

        validate_temperature(temperature)
        self.temperature = temperature
        return self

    def set_max_output_tokens(self, max_output_tokens: int) -> "Dracula":
        """
        Change the maximum response length in tokens.

        Args:
            max_output_tokens (int): Maximum number of tokens in the response.

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> ai.set_max_output_tokens(512)
        """

        validate_max_output_tokens(max_output_tokens)
        self.max_output_tokens = max_output_tokens
        logger.info(f"Max output tokens changed to '{max_output_tokens}'")
        return self

    def set_language(self, language: str) -> "Dracula":
        """
        Change the response language and clear conversation memory.

        Forces Gemini to always respond in the specified language
        regardless of what language the user writes in.

        Args:
            language (str): The language for responses (e.g. 'Turkish', 'Spanish').

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> ai.set_language("Turkish")
        """

        validate_language(language)
        self.language = language
        self.memory.clear()
        self._create_chat_session()
        logger.info(f"Language changed to '{language}'")
        return self

    def set_persona(self, persona: str) -> "Dracula":
        """
        Switch to a built-in persona instantly.

        Each persona comes with a predefined prompt, temperature, and language.
        Switching personas automatically clears the conversation memory.

        Args:
            persona (str): The name of the persona to use.

        Returns:
            Dracula: The current instance for method chaining.

        Raises:
            PersonaException: If the persona name is not recognized.

        Example:
            >>> ai.set_persona("pirate")
        """

        if persona not in PERSONAS:
            available = ", ".join(PERSONAS.keys())

            raise PersonaException(
                f"Unknown persona '{persona}'. Available personas: {available}"
            )

        selected = PERSONAS[persona]
        self.set_prompt(selected["prompt"])
        self.set_temperature(selected["temperature"])
        self.set_language(selected["language"])

        logger.info(f"Persona changed to '{persona}'")

        return self

    def set_log_level(self, level: str) -> "Dracula":
        """
        Change the logging level at any time.

        Args:
            level (str): Logging level — DEBUG, INFO, WARNING, ERROR, or CRITICAL.

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> ai.set_log_level("WARNING")
        """
        enable_logging(level)
        return self

    def list_personas(self):
        """
        Return a list of all available built-in persona names.

        Returns:
            list: A list of persona name strings.

        Example:
            >>> print(ai.list_personas())
            ['assistant', 'pirate', 'chef', 'shakespeare', 'scientist', 'comedian']
        """

        return list(PERSONAS.keys())

    def list_available_models(self) -> list:
        """
        Fetch all currently available Gemini models from the API in real time.

        Use this to discover new models that may not yet be in the
        GeminiModel enum. Unlike GeminiModel, this method makes an
        API request so use it sparingly.

        Returns:
            list: A list of available model name strings.

        Example:
            >>> print(ai.list_available_models())
        """
        try:
            models = self.client.models.list()
            return [model.name for model in models]
        except Exception as e:
            raise ChatException(f"Failed to fetch models: {str(e)}")

    def get_history(self):
        """
        Return the full conversation history as a list of dictionaries.

        Returns:
            list: A list of message dictionaries with 'role' and 'content' keys.

        Example:
            >>> print(ai.get_history())
        """

        return self.memory.get_history()

    def add_tool(self, tool: Tool) -> "Dracula":
        """
        Register a new tool after initialization.

        Args:
            tool (Tool): The tool to register.

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> ai.add_tool(get_weather)
        """

        self.tool_registry.register(tool)
        logger.info(f"Tool '{tool.name}' registered.")
        return self

    def list_tools(self) -> list:
        """
        Return a list of all registered tool names.

        Returns:
            list: List of tool name strings.

        Example:
            >>> print(ai.list_tools())
        """

        return self.tool_registry.list_tools()

    def _build_system_instruction(self) -> str:
        """Build the system instruction with language."""
        if self.language.lower() != "auto":
            return f"{self.prompt} Always respond in {self.language}."
        return self.prompt

    def _build_config(self) -> genai.types.GenerateContentConfig:
        """Build the GenerateContentConfig."""
        config = genai.types.GenerateContentConfig(
            system_instruction=self._build_system_instruction(),
            temperature=self.temperature,
            max_output_tokens=self.max_output_tokens,
        )
        if self.tool_registry.has_tools():
            config.tools = [
                genai.types.Tool(
                    function_declarations=self.tool_registry.get_gemini_schemas()
                )
            ]
        return config

    def _create_chat_session(self):
        """Create a new Gemini chat session."""
        self._chat_session = self.client.chats.create(
            model=self.model_name, config=self._build_config()
        )
