import json
import asyncio
from google import genai
from pathlib import Path
from .memory import AsyncMemory
from .tools import Tool, ToolRegistry, ToolResult
from .stats import AsyncStats
from .cache import AsyncCache
from .pricing import estimate_cost
from .retry import async_retry
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,
    ValidationException,
    BudgetExceededException,
)
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'.
        session_id (str): Optional identifier for this conversation session. Allows
            multiple isolated sessions to share one database. Auto-generated if not provided.
        cache_ttl (int): Seconds to cache identical responses. 0 disables caching (default).
        token_budget (int): Maximum cumulative tokens before BudgetExceededException is raised.
        auto_compress (bool): Automatically summarize old history instead of hard-trimming.
        compress_ratio (float): Fraction of max_messages at which compression triggers.
        compress_keep_turns (int): Recent messages preserved verbatim after compression.
        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,
        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,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        db_path: str | Path = None,
        stats_db_path: str | Path = None,
        session_id: str = None,
        cache_ttl: int = 0,
        cache_db_path: str | Path = None,
        token_budget: int = None,
        auto_compress: bool = False,
        compress_ratio: float = 0.8,
        compress_keep_turns: int = 4,
    ):
        """Initializes the AsyncDracula client.

        Args:
            api_key (str): Your Google Gemini API key.
            model (GeminiModel | str): The model to use for generations.
            max_messages (int): Limit for message history retention.
            prompt (str): The initial system instruction.
            db_path (str): Filepath for the SQLite memory database.
        """

        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._api_key = api_key
        self.tool_registry = ToolRegistry(tools or [])
        self._token_budget = token_budget
        self._auto_compress = auto_compress
        self._compress_ratio = compress_ratio
        self._compress_keep_turns = compress_keep_turns
        self._before_hooks: list = []
        self._after_hooks: list = []

        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 = AsyncMemory(db_path=db_path, max_messages=max_messages, session_id=session_id)
        self.stats = AsyncStats(db_path=stats_db_path)
        self._cache = AsyncCache(db_path=cache_db_path, ttl=cache_ttl)
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        check_latest_version(__version__)

        self._chat_session = None
        self._initialized = False

        logger.debug(
            f"AsyncDracula initialized with model={self.model_name}, "
            f"language={language}, temperature={temperature}"
        )

    async def _ensure_initialized(self):
        """Ensures that the memory database and chat session are ready.

        Since async operations cannot occur in __init__, this method handles
        the deferred setup of the SQLite schema and the initial session.
        """

        if not self._initialized:
            await self.memory.initialize()
            await self.stats.initialize()
            await self._cache.initialize()
            await self._create_chat_session()
            self._initialized = True

    async def _prepare_file_part(self, filepath: str | Path) -> genai.types.Part:
        """
        Asynchronously prepares a file part for the Gemini API.
        Uses thread pooling to prevent blocking the event loop during disk I/O.
        """
        path = Path(filepath)
        if not path.exists():
            raise ValidationException(f"File not found: {filepath}")

        import mimetypes
        import asyncio

        mime_type, _ = mimetypes.guess_type(path)
        if not mime_type:
            mime_type = "application/octet-stream"

        data = await asyncio.to_thread(path.read_bytes)

        return genai.types.Part.from_bytes(data=data, mime_type=mime_type)

    # ------------------------------------------------------------------
    # Hooks
    # ------------------------------------------------------------------

    def before_chat(self, func):
        """Register a hook that runs before every chat() call."""
        self._before_hooks.append(func)
        return func

    def after_chat(self, func):
        """Register a hook that runs after every chat() call."""
        self._after_hooks.append(func)
        return func

    async def _run_before_hooks(self, message: str) -> str:
        for hook in self._before_hooks:
            if asyncio.iscoroutinefunction(hook):
                result = await hook(message)
            else:
                result = hook(message)
            if result is not None:
                message = result
        return message

    async def _run_after_hooks(self, message: str, reply: str) -> str:
        for hook in self._after_hooks:
            if asyncio.iscoroutinefunction(hook):
                result = await hook(message, reply)
            else:
                result = hook(message, reply)
            if result is not None:
                reply = result
        return reply

    # ------------------------------------------------------------------
    # Context compression
    # ------------------------------------------------------------------

    async def _maybe_compress(self):
        history = await self.memory.get_history()
        threshold = int(self.memory.max_messages * self._compress_ratio)
        if len(history) >= threshold:
            await self._compress_history(history)

    async def _compress_history(self, history: list):
        keep = self._compress_keep_turns
        to_summarize = history[:-keep] if keep > 0 else history
        if len(to_summarize) < 2:
            return
        conv_text = "\n".join(
            f"{m['role'].upper()}: {m['content']}" for m in to_summarize
        )
        prompt = (
            "Summarize the following conversation in 2-3 sentences, "
            "preserving all key facts and decisions:\n\n" + conv_text
        )
        try:
            resp = await self.client.aio.models.generate_content(
                model=self.model_name, contents=prompt
            )
            summary = resp.text or "No summary available."
            await self.memory.compress_to_summary(summary, keep_last_n=keep)
            await self._create_chat_session()
            logger.info(
                f"Async context compressed: {len(to_summarize)} messages -> 1 summary."
            )
        except Exception as e:
            logger.warning(f"Async context compression failed, skipping: {e}")

    # ------------------------------------------------------------------
    # chat()
    # ------------------------------------------------------------------

    async def chat(
        self,
        message: str,
        filepath: str | Path = None,
        auto_call: bool = True,
        schema=None,
    ) -> str | dict | ToolResult:
        """Sends a message and returns the model's response asynchronously.

        Args:
            message (str): The user input text.
            filepath (str | Path): Optional file to attach.
            auto_call (bool): Handle tool calls automatically.
            schema: Optional Pydantic BaseModel subclass or JSON-schema dict
                for structured output.

        Returns:
            str | dict | BaseModel: The response in the requested format.
        """
        await self._ensure_initialized()
        validate_message(message)

        # Budget check
        if self._token_budget is not None:
            s = await self.stats.get_stats()
            used = s.get("total_input_tokens", 0) + s.get("total_output_tokens", 0)
            if used >= self._token_budget:
                raise BudgetExceededException(
                    f"Token budget of {self._token_budget:,} exceeded "
                    f"(used: {used:,} tokens)."
                )

        message = await self._run_before_hooks(message)

        if schema is not None:
            return await self._chat_with_schema(message, schema)

        cached = await self._cache.get(self.memory.session_id, message)
        if cached is not None:
            logger.debug("Cache hit — returning cached response.")
            return await self._run_after_hooks(message, cached)

        if self._auto_compress:
            await self._maybe_compress()

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

            if filepath:
                file_part = await self._prepare_file_part(filepath)
                payload = [message, file_part]
            else:
                payload = message

            response = await self._chat_session.send_message(payload)

            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 = 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.")

                    if hasattr(response, "usage_metadata") and response.usage_metadata:
                        await self.stats.record_tokens(
                            getattr(response.usage_metadata, "prompt_token_count", 0) or 0,
                            getattr(response.usage_metadata, "candidates_token_count", 0) or 0,
                        )

                    history_msg = (
                        f"{message} [Attached: {Path(filepath).name}]"
                        if filepath
                        else message
                    )
                    await self.memory.add_message("user", history_msg)
                    await self.memory.add_message("model", reply)
                    await self.stats.record_response(reply)

                    await self._cache.set(self.memory.session_id, message, reply)

                    reply = await self._run_after_hooks(message, reply)
                    logger.debug(f"Received response: {reply[:50]}...")
                    return reply

        return await async_retry(
            _do_chat, max_retries=self.max_retries, retry_delay=self.retry_delay
        )

    async def _chat_with_schema(self, message: str, schema) -> str | dict:
        """Internal: one-shot generate_content call with structured output."""
        try:
            pydantic_model = None
            if isinstance(schema, type) and hasattr(schema, "model_json_schema"):
                pydantic_model = schema
                json_schema = schema.model_json_schema()
            elif isinstance(schema, dict):
                json_schema = schema
            else:
                raise ValidationException(
                    "schema must be a Pydantic BaseModel subclass or a JSON-schema dict."
                )

            history = await self.memory.get_history()
            contents = [
                genai.types.Content(
                    role=m["role"], parts=[genai.types.Part(text=m["content"])]
                )
                for m in history
            ]
            contents.append(
                genai.types.Content(
                    role="user", parts=[genai.types.Part(text=message)]
                )
            )

            response = await self.client.aio.models.generate_content(
                model=self.model_name,
                contents=contents,
                config=genai.types.GenerateContentConfig(
                    system_instruction=self._build_system_instruction(),
                    temperature=self.temperature,
                    max_output_tokens=self.max_output_tokens,
                    response_mime_type="application/json",
                    response_schema=json_schema,
                ),
            )

            reply = response.text
            if not reply:
                raise ChatException("Received an empty structured response from Gemini.")

            if hasattr(response, "usage_metadata") and response.usage_metadata:
                await self.stats.record_tokens(
                    getattr(response.usage_metadata, "prompt_token_count", 0) or 0,
                    getattr(response.usage_metadata, "candidates_token_count", 0) or 0,
                )

            await self.memory.add_message("user", message)
            await self.memory.add_message("model", reply)
            await self.stats.record_message(message)
            await self.stats.record_response(reply)

            if pydantic_model is not None:
                return pydantic_model.model_validate_json(reply)
            return json.loads(reply)

        except (ValidationException, ChatException):
            raise
        except Exception as e:
            raise ChatException(f"Structured output request failed: {str(e)}")

    async def stream(self, message: str):
        """Asynchronously streams the response word by word with memory integration.

        This method ensures the database is initialized, records the outgoing
        user message, and yields chunks of the response in real-time. Once
        the stream completes, the full response is persisted to the SQLite
        backend to maintain conversational context.

        Args:
            message (str): The user input text to send to Gemini.

        Yields:
            str: Chunks of text as they are generated by the model.

        Raises:
            ChatException: If the stream is interrupted or the database
                operations fail.
        """
        await self._ensure_initialized()
        validate_message(message)

        # Budget check
        if self._token_budget is not None:
            s = await self.stats.get_stats()
            used = s.get("total_input_tokens", 0) + s.get("total_output_tokens", 0)
            if used >= self._token_budget:
                raise BudgetExceededException(
                    f"Token budget of {self._token_budget:,} exceeded "
                    f"(used: {used:,} tokens)."
                )

        message = await self._run_before_hooks(message)

        if self._auto_compress:
            await self._maybe_compress()

        await self.stats.record_message(message)
        last_exception = None

        for attempt in range(self.max_retries + 1):
            full_reply = ""
            try:
                async for chunk in await self._chat_session.send_message_stream(
                    message
                ):
                    if chunk.text:
                        full_reply += chunk.text
                        yield chunk.text

                await self.memory.add_message("user", message)
                await self.memory.add_message("model", full_reply)

                await self.stats.record_response(full_reply)
                await self._run_after_hooks(message, full_reply)
                return

            except Exception as e:
                last_exception = e

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

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

                import asyncio

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

        raise last_exception

    async 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
        await self.memory.clear()
        await 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

    async 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
        await self.memory.clear()
        await self._create_chat_session()
        logger.info(f"Language changed to '{language}'")
        return self

    async 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
        await self._create_chat_session()
        logger.info(f"Model changed to '{self.model_name}'")
        return self

    async 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]
        await self.set_prompt(selected["prompt"])
        self.set_temperature(selected["temperature"])
        await 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 set_retry(self, max_retries: int, retry_delay: float = 1.0) -> "AsyncDracula":
        """
        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 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)}")

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

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

    async def estimated_cost(self) -> float:
        """Return the estimated API cost in USD based on tokens used so far.

        Returns:
            float: Estimated cost in USD (rounded to 6 decimal places).
        """
        s = await self.stats.get_stats()
        return estimate_cost(
            self.model_name,
            s.get("total_input_tokens", 0),
            s.get("total_output_tokens", 0),
        )

    async def fork(self, session_id: str = None) -> "AsyncDracula":
        """Create a new AsyncDracula instance with a copy of the current conversation.

        Args:
            session_id (str): Optional session ID for the fork. Auto-generated if not provided.

        Returns:
            AsyncDracula: A new instance pre-loaded with the current history.
        """
        branch = AsyncDracula(
            api_key=self._api_key,
            model=self.model_name,
            max_messages=self.memory.max_messages,
            prompt=self.prompt,
            temperature=self.temperature,
            max_output_tokens=self.max_output_tokens,
            language=self.language,
            db_path=self.memory.db_path,
            stats_db_path=self.stats.db_path,
            session_id=session_id,
            cache_ttl=self._cache.ttl,
            token_budget=self._token_budget,
            auto_compress=self._auto_compress,
            compress_ratio=self._compress_ratio,
            compress_keep_turns=self._compress_keep_turns,
        )
        await branch._ensure_initialized()
        for msg in await self.memory.get_history():
            await branch.memory.add_message(msg["role"], msg["content"])
        await branch._create_chat_session()
        logger.info(f"Forked async session into '{branch.memory.session_id}'.")
        return branch

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

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

    async def export_to_markdown(self, filepath: str):
        """
        Asynchronously export the conversation history to a formatted Markdown file.

        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:
            >>> await ai.export_to_markdown("my_chat.md")
        """
        validate_filepath(filepath)
        await self._ensure_initialized()
        history = await 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"Async history successfully exported to markdown: {filepath}")

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

    async def clear_memory(self):
        """Permanently clears conversational history from the SQLite database.

        This action also resets the current active chat session.
        """
        await self._ensure_initialized()
        await self.memory.clear()
        await self._create_chat_session()
        logger.info("Asynchronous memory and session cleared.")

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

    async def print_history(self):
        """Print the conversation history in a clean, readable format."""
        history = await 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:
        await self.clear_memory()
        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

    async def _create_chat_session(self):
        """Creates or restores a Gemini chat session with stored history.

        Fetches history from the asynchronous SQLite backend to populate
        the initial context of the Gemini session.
        """
        history_data = await self.memory.get_history()

        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()
                )
            ]

        self._chat_session = self.client.aio.chats.create(
            model=self.model_name,
            config=config,
            history=[
                genai.types.Content(
                    role=msg["role"], parts=[genai.types.Part(text=msg["content"])]
                )
                for msg in history_data
            ],
        )

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

    async def count_tokens(self, text: str) -> int:
        """
        Asynchronously count the number of tokens in a given text.

        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 = await ai.count_tokens("Hello, Gemini!")
        """

        validate_message(text)
        try:
            response = await self.client.aio.models.count_tokens(
                model=self.model_name,
                contents=text,
            )
            return response.total_tokens
        except Exception as e:
            raise ChatException(f"Failed to count async tokens: {str(e)}")

    async def count_history_tokens(self) -> int:
        """
        Asynchronously 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 = await ai.count_history_tokens()
        """
        try:
            await self._ensure_initialized()
            history = await self.memory.get_history()

            # If there is no history, we only count the system prompt
            if not history:
                prompt_text = self._build_system_instruction()
                return await self.count_tokens(prompt_text)

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

            response = await self.client.aio.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 async history tokens: {str(e)}")
