from google import genai
from pathlib import Path
from .memory import Memory
from .tools import Tool, ToolRegistry, ToolResult
from .stats import Stats
from .retry import retry
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 assistant.",
        temperature: float = 1.0,
        max_output_tokens: int = 8192,
        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,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        db_path: str | Path = None,
        stats_db_path: str | Path = 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_language(language)

        self.tool_registry = ToolRegistry(tools or [])
        self.max_retries = max_retries
        self.retry_delay = retry_delay

        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(db_path=db_path, max_messages=max_messages)
        self.stats = Stats(db_path=stats_db_path)

        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()
        return False

    def chat(self, message: str, auto_call: bool = True) -> str | ToolResult:
        """
        Send a message to Gemini and get a response.
        Automatically retries on failure with exponential backoff.

        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.
            ToolResult: Tool call information (when auto_call=False).
        """
        validate_message(message)

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

            response = self._chat_session.send_message(message)

            while True:
                tool_call_found = False

                if not response.candidates or not response.candidates[0].content:
                    raise ChatException(
                        "Response was blocked or returned empty (likely due to safety filters)."
                    )

                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,
                            )

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

                        response = self._chat_session.send_message(
                            genai.types.Part(
                                function_response=genai.types.FunctionResponse(
                                    name=tool_name,
                                    response={"result": str(tool_result)},
                                )
                            )
                        )
                        break

                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)
                    self._full_history = self.memory.get_history()
                    logger.debug(f"Received response: {reply[:50]}...")
                    return reply

        return retry(
            _do_chat, max_retries=self.max_retries, retry_delay=self.retry_delay
        )

    def stream(self, message: str):
        """
        Stream the response word by word with proper real-time yielding.
        Automatically retries on initial connection failure.
        """
        validate_message(message)
        self.stats.record_message(message)

        last_exception = None

        for attempt in range(self.max_retries + 1):
            full_reply = ""
            try:
                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)
                self._full_history = self.memory.get_history()
                return

            except Exception as e:
                last_exception = e

                if full_reply:
                    logger.error(f"Stream interrupted after partial response: {str(e)}")
                    raise ChatException(f"Stream interrupted: {str(e)}")

                if attempt == self.max_retries:
                    raise ChatException(
                        f"Stream failed after {self.max_retries} retries. Error: {str(e)}"
                    )

                import time

                delay = self.retry_delay * (2**attempt)
                logger.warning(
                    f"Stream attempt {attempt + 1} failed: {str(e)}. Retrying in {delay}s..."
                )
                time.sleep(delay)

        raise last_exception

    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)
        self._create_chat_session()

    def export_to_markdown(self, filepath: str | Path):
        """
        Export the conversation history to a beautifully formatted Markdown file.

        Useful for reading the chat history in editors like VS Code, Obsidian,
        or GitHub, especially when the conversation contains code blocks.

        Args:
            filepath (str): The destination path for the .md file.

        Raises:
            ValidationException: If the filepath is invalid.
            ChatException: If the file cannot be written.

        Example:
            >>> ai.export_to_markdown("my_chat.md")
        """

        validate_filepath(filepath)
        history = self.memory.get_history()

        try:
            from datetime import datetime

            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            md_content = "# Dracula Conversation History\n\n"
            md_content += f"> Exported on: {now}\n\n---\n\n"

            if not history:
                md_content += "*No conversation history yet.*\n"

            else:
                for msg in history:
                    role_icon = (
                        "👤 **You**" if msg["role"] == "user" else "🤖 **Gemini**"
                    )
                    md_content += f"### {role_icon}\n\n"
                    md_content += f"{msg['content']}\n\n---\n\n"

            Path(filepath).write_text(md_content, encoding="utf-8")
            logger.info(f"History successfully exported to markdown: {filepath}")

        except Exception as e:
            raise ChatException(f"Failed to export markdown: {str(e)}")

    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}'")

    def set_retry(self, max_retries: int, retry_delay: float = 1.0) -> "Dracula":
        """
        Change the retry settings.

        Args:
            max_retries (int): Maximum number of retry attempts.
            retry_delay (float): Initial delay in seconds between retries.

        Returns:
            Dracula: The current instance for method chaining.

        Example:
            >>> ai.set_retry(max_retries=5, retry_delay=2.0)
        """
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        logger.info(
            f"Retry settings updated: max_retries={max_retries}, retry_delay={retry_delay}"
        )
        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(),
            history=[
                genai.types.Content(
                    role=msg["role"], parts=[genai.types.Part(text=msg["content"])]
                )
                for msg in self.memory.get_history()
            ],
        )

    def count_tokens(self, text: str) -> int:
        """
        Count the number of tokens in a given text using the current model.

        Args:
            text (str): The text to calculate tokens for.

        Returns:
            int: The total number of tokens.

        Raises:
            ValidationException: If the text is invalid.
            ChatException: If the API fails to count tokens.

        Example:
            >>> tokens = ai.count_tokens("Hello, Gemini!")
            >>> print(f"This text uses {tokens} tokens.")
        """

        validate_message(text)

        try:
            response = self.client.models.count_tokens(
                model=self.model_name,
                contents=text,
            )

            return response.total_tokens

        except Exception as e:
            raise ChatException(f"Failed to count tokens: {str(e)}")

    def count_history_tokens(self) -> int:
        """
        Count the total number of tokens used by the current conversation history,
        including the system prompt and registered tools.

        Returns:
            int: The total number of tokens in the active session.

        Raises:
            ChatException: If the API fails to count tokens.

        Example:
            >>> total_usage = ai.count_history_tokens()
            >>> print(f"Current session uses {total_usage} tokens.")
        """

        try:
            history = self.memory.get_history()

            if not history:
                prompt_text = self._build_system_instruction()
                return self.count_history_tokens(prompt_text)

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

            response = self.client.models.count_tokens(
                model=self.model_name, contents=contents, config=self._build_config()
            )
            return response.total_tokens

        except Exception as e:
            raise ChatException(f"Failed to count history tokens: {str(e)}")
