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

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


class AsyncDracula:
    """
    Async version of Dracula for use in async applications.

    AsyncDracula provides the same features as Dracula but with
    full async support, making it perfect for Discord bots, FastAPI
    web apps, Telegram bots, and any other async Python application.

    Args:
        api_key (str): Your Google Gemini API key.
        model (GeminiModel | str): The Gemini model to use. Defaults to GeminiModel.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.
        temperature (float): Controls response creativity between 0.0 and 2.0.
        max_output_tokens (int): Maximum length of responses in tokens.
        stats_filepath (str): Path to save usage stats.
        language (str): Language for responses. Defaults to 'English'.
        logging (bool): Enable or disable logging. Defaults to False.
        log_level (str): Logging level. Defaults to 'DEBUG'.
        log_file (str): Path to log file. Defaults to None.
        log_max_bytes (int): Maximum log file size in bytes. Defaults to 5MB.
        log_backup_count (int): Number of backup log files to keep. Defaults to 5.

    Example:
        >>> import asyncio
        >>> from dracula import AsyncDracula
        >>> async def main():
        ...     async with AsyncDracula(api_key="your-api-key") as ai:
        ...         response = await ai.chat("Hello!")
        ...         print(response)
        >>> asyncio.run(main())
    """

    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,
        stats_filepath: str = "dracula_stats.json",
        language: str = "English",
        logging: bool = False,
        log_level: str = "DEBUG",
        log_file: str = None,
        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)
        validate_language(language)

        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"AsyncDracula initialized with model={self.model_name}, "
            f"language={language}, temperature={temperature}"
        )
        self._create_chat_session()

    async def chat(self, message: str, auto_call: bool = True) -> str | ToolResult:
        """
        Asynchronously 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 = await ai.chat("What's the weather in Istanbul?")
            >>> print(response)
        """
        validate_message(message)

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

            response = await 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 asynchronously
                        tool = self.tool_registry.get(tool_name)
                        tool_result = await tool.async_call(**tool_args)
                        logger.debug(f"Tool '{tool_name}' returned: {tool_result}")

                        # Send tool result back
                        response = await 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)
                    logger.debug(f"Received response: {reply[:50]}...")
                    return reply

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

    async def stream(self, message: str):
        """
        Asynchronously send a message and stream the response word by word.

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

        Yields:
            str: Small chunks of the response text.
        """
        validate_message(message)

        try:
            self.stats.record_message(message)

            full_reply = ""
            async for chunk in await 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 set_prompt(self, prompt: str) -> "AsyncDracula":
        """
        Change the system prompt and clear conversation memory.

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

        Returns:
            AsyncDracula: The current instance for method chaining.
        """
        validate_prompt(prompt)
        self.prompt = prompt
        self.memory.clear()
        self._create_chat_session()
        logger.info("Prompt updated.")
        return self

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

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

        Returns:
            AsyncDracula: The current instance for method chaining.
        """
        validate_temperature(temperature)
        self.temperature = temperature
        return self

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

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

        Returns:
            AsyncDracula: The current instance for method chaining.
        """
        validate_max_output_tokens(max_output_tokens)
        self.max_output_tokens = max_output_tokens
        return self

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

        Args:
            language (str): The language for responses.

        Returns:
            AsyncDracula: The current instance for method chaining.
        """
        validate_language(language)
        self.language = language
        self.memory.clear()
        self._create_chat_session()
        logger.info(f"Language changed to '{language}'")
        return self

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

        Args:
            model (GeminiModel | str): The model to use.

        Returns:
            AsyncDracula: The current instance for method chaining.
        """
        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_persona(self, persona: str) -> "AsyncDracula":
        """
        Switch to a built-in persona instantly.

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

        Returns:
            AsyncDracula: The current instance for method chaining.

        Raises:
            PersonaException: If the persona name is not recognized.
        """
        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) -> "AsyncDracula":
        """
        Change the logging level.

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

        Returns:
            AsyncDracula: The current instance for method chaining.
        """
        enable_logging(level)
        return self

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

        Returns:
            list: A list of persona name strings.
        """
        return list(PERSONAS.keys())

    async def list_available_models(self) -> list:
        """
        Asynchronously fetch all currently available Gemini models from the API.

        Returns:
            list: A list of available model name strings.
        """
        try:
            models = await self.client.aio.models.list()
            return [model.name for model in models]
        except Exception as e:
            raise ChatException(f"Failed to fetch models: {str(e)}")

    def get_stats(self) -> dict:
        """Return the current usage statistics."""
        return self.stats.get_stats()

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

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

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

    def clear_memory(self):
        """Clear the conversation history."""
        self.memory.clear()
        self._create_chat_session()
        logger.info("Memory cleared.")

    def get_history(self) -> list:
        """Return the full conversation history."""
        return self.memory.get_history()

    def print_history(self):
        """Print the conversation history in a clean, readable format."""
        history = self.memory.get_history()

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

        print("\n" + "═" * 50)
        print("         🧛 DRACULA CONVERSATION HISTORY")
        print("═" * 50)

        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" + "═" * 50 + "\n")

    async def close(self):
        """Close the async client and release all resources."""
        try:
            if hasattr(self.client, "_api_client"):
                if hasattr(self.client._api_client, "_session"):
                    session = self.client._api_client._session
                    if session and not session.closed:
                        await session.close()
        except Exception:
            pass

    async def __aenter__(self) -> "AsyncDracula":
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
        self.clear_memory()
        self.reset_stats()
        await self.close()
        return False

    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 async chat session."""
        self._chat_session = self.client.aio.chats.create(
            model=self.model_name, config=self._build_config()
        )

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

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

        Returns:
            AsyncDracula: The current instance for method chaining.

        Example:
            >>> ai.add_tool(get_weather)
        """
        self.tool_registry.register(tool)
        self._create_chat_session()
        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()
