"""
Core Chatbot class — the main entry point for ArielAI.
"""

from __future__ import annotations

import logging
from typing import Generator, Optional, Union

from .backends import BaseBackend, create_backend
from .models import ModelPreset, get_preset, list_presets

logger = logging.getLogger(__name__)


class Chatbot:
    """
    Build and launch a Gradio-powered chatbot in minutes.

    Supports HuggingFace Transformers, Ollama, and the HuggingFace
    Inference API as backends.

    Basic usage::

        from arielai import Chatbot

        bot = Chatbot("microsoft/DialoGPT-medium")
        bot.launch()

    Using a preset::

        bot = Chatbot.from_preset("zephyr")
        bot.launch()

    Using Ollama::

        bot = Chatbot("llama3", backend="ollama")
        bot.launch()
    """

    def __init__(
        self,
        model: str = "microsoft/DialoGPT-medium",
        *,
        backend: str = "transformers",
        title: str = "ArielAI Chatbot",
        description: str = "",
        system_prompt: str = "",
        placeholder: str = "Ask me anything...",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        examples: Optional[list[str]] = None,
        theme: Optional[str] = "soft",
        avatar_images: Optional[tuple[Optional[str], Optional[str]]] = None,
        streaming: bool = True,
        device: Optional[str] = None,
        load_in_8bit: bool = False,
        load_in_4bit: bool = False,
        hf_token: Optional[str] = None,
        ollama_host: str = "http://localhost:11434",
        backend_instance: Optional[BaseBackend] = None,
    ) -> None:
        """
        Args:
            model: HuggingFace model ID, Ollama model name, or preset name.
            backend: One of "transformers", "ollama", or "inference".
            title: Chatbot window title shown in Gradio.
            description: Short description shown below the title.
            system_prompt: System/instruction prompt prepended to every conversation.
            placeholder: Input box placeholder text.
            max_new_tokens: Maximum tokens to generate per response.
            temperature: Sampling temperature (0 = deterministic, 1 = creative).
            examples: List of example user messages shown below the chat.
            theme: Gradio theme name ("soft", "default", "monochrome", "glass").
            avatar_images: Tuple of (user_avatar_path, bot_avatar_path) or None.
            streaming: Whether to stream responses token-by-token.
            device: Device for transformers ("cpu", "cuda", "mps"). Auto-detected if None.
            load_in_8bit: Enable 8-bit quantization (requires bitsandbytes + GPU).
            load_in_4bit: Enable 4-bit quantization (requires bitsandbytes + GPU).
            hf_token: HuggingFace API token for private/gated models.
            ollama_host: Ollama server URL (default: http://localhost:11434).
            backend_instance: Provide a custom BaseBackend instance directly.
        """
        self.model = model
        self.backend_name = backend
        self.title = title
        self.description = description
        self.system_prompt = system_prompt
        self.placeholder = placeholder
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.examples = examples
        self.theme = theme
        self.avatar_images = avatar_images
        self.streaming = streaming

        if backend_instance is not None:
            self._backend = backend_instance
        else:
            backend_kwargs: dict = {}
            if backend == "transformers":
                backend_kwargs = {
                    "device": device,
                    "load_in_8bit": load_in_8bit,
                    "load_in_4bit": load_in_4bit,
                }
            elif backend == "ollama":
                backend_kwargs = {"host": ollama_host}
            elif backend == "inference":
                backend_kwargs = {"token": hf_token}

            self._backend = create_backend(backend, model, **backend_kwargs)

        self._demo = None

    @classmethod
    def from_preset(cls, preset_name: str, **kwargs) -> "Chatbot":
        """
        Create a Chatbot from a named preset.

        Args:
            preset_name: Name of the preset (e.g. "zephyr", "tiny", "ollama-llama3").
            **kwargs: Override any preset defaults.

        Raises:
            ValueError: If the preset name is not found.

        Example::

            bot = Chatbot.from_preset("zephyr", title="My Assistant")
            bot.launch()
        """
        preset = get_preset(preset_name)
        if preset is None:
            available = ", ".join(list_presets().keys())
            raise ValueError(
                f"Unknown preset '{preset_name}'. Available: {available}"
            )

        init_kwargs = {
            "model": preset.model_id,
            "backend": preset.backend,
            "max_new_tokens": preset.max_new_tokens,
            "temperature": preset.temperature,
            "system_prompt": preset.system_prompt,
        }
        init_kwargs.update(kwargs)
        return cls(**init_kwargs)

    def chat(
        self, message: str, history: list[tuple[str, str]]
    ) -> Generator[str, None, None]:
        """
        Core chat handler. Used internally by Gradio but can also be called
        programmatically.

        Always a generator function for Gradio 6+ compatibility. When streaming
        is disabled it yields the full response in one step; when enabled it
        yields partial tokens as they arrive.

        Args:
            message: The user's current message.
            history: List of (user, assistant) message pairs.

        Yields:
            Partial or complete response strings.
        """
        if self.streaming:
            for chunk in self._backend.generate_stream(
                message=message,
                history=history,
                system_prompt=self.system_prompt,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            ):
                yield chunk
        else:
            yield self._backend.generate(
                message=message,
                history=history,
                system_prompt=self.system_prompt,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            )

    def build(self):
        """
        Build and return the Gradio demo without launching it.
        Useful for deploying to Hugging Face Spaces or embedding in
        existing Gradio apps.

        Returns:
            A ``gr.ChatInterface`` instance.
        """
        try:
            import gradio as gr
        except ImportError as e:
            raise ImportError(
                "gradio is required. Install it with: pip install arielai"
            ) from e

        theme = self._resolve_theme(self.theme)

        self._demo = gr.ChatInterface(
            fn=self.chat,
            title=self.title,
            description=self.description or None,
            placeholder=self.placeholder,
            examples=self.examples,
            theme=theme,
            avatar_images=self.avatar_images,
        )
        return self._demo

    def launch(
        self,
        share: bool = False,
        server_name: str = "0.0.0.0",
        server_port: Optional[int] = None,
        inbrowser: bool = True,
        **gradio_kwargs,
    ) -> None:
        """
        Build and launch the Gradio chatbot UI.

        Args:
            share: Create a public Gradio share link.
            server_name: Bind address (default "0.0.0.0" for all interfaces).
            server_port: Port number. Auto-selected if None.
            inbrowser: Open the browser automatically.
            **gradio_kwargs: Any additional kwargs passed to ``demo.launch()``.
        """
        demo = self.build()
        demo.launch(
            share=share,
            server_name=server_name,
            server_port=server_port,
            inbrowser=inbrowser,
            **gradio_kwargs,
        )

    def close(self) -> None:
        """Shut down the Gradio server and release model resources."""
        if self._demo is not None:
            self._demo.close()
        self._backend.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()

    @staticmethod
    def _resolve_theme(theme: Optional[str]):
        """Convert a theme name string to a Gradio theme object."""
        try:
            import gradio as gr

            theme_map = {
                "soft": gr.themes.Soft(),
                "default": gr.themes.Default(),
                "monochrome": gr.themes.Monochrome(),
                "glass": gr.themes.Glass(),
            }
            if theme is None:
                return gr.themes.Soft()
            return theme_map.get(theme, gr.themes.Soft())
        except Exception:
            return None
