from __future__ import annotations

import logging
import os
from typing import Generator, Optional

from .backends import (
    BaseBackend, TransformersBackend, OllamaBackend,
    HFInferenceBackend, OpenAIBackend,
    create_backend, auto_select_backend, _normalize_history,
)
from .models import ModelPreset, get_preset, list_presets, DEFAULT_MODELS

logger = logging.getLogger(__name__)


class Chatbot:
    """
    Build and launch an AI chatbot with one line of code.

    Auto-selects the best available free backend:
    - OpenAI if ``openai_api_key`` is provided (GPT-4o-mini by default)
    - Ollama if running locally — highest quality, fully free
    - Transformers — runs on any machine, no API key needed

    Examples::

        from arielai import Chatbot

        Chatbot().launch()                                          # auto-detect
        Chatbot(openai_api_key="sk-...").launch()                   # OpenAI
        Chatbot("gpt-4o", backend="openai", openai_api_key="sk-...").launch()
        Chatbot("llama3.2:3b", backend="ollama").launch()
        Chatbot.from_preset("gpt4o-mini", openai_api_key="sk-...").launch()
    """

    def __init__(
        self,
        model: Optional[str] = None,
        *,
        backend: Optional[str] = None,
        openai_api_key: Optional[str] = None,
        system_prompt: str = "You are a helpful AI assistant. Be concise and clear.",
        title: str = "ArielAI",
        description: str = "",
        placeholder: str = "Type a message...",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
        examples: Optional[list[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:
        # Resolve API key from env if not passed directly
        api_key = openai_api_key or os.environ.get("OPENAI_API_KEY")

        if backend_instance is not None:
            self._backend = backend_instance
            self.backend_name = "custom"
            self.model = model or "custom"
        elif api_key and backend is None and model is None:
            # OpenAI key provided — default to GPT-4o-mini
            self.backend_name = "openai"
            self.model = DEFAULT_MODELS["openai"]
            self._backend = OpenAIBackend(model_id=self.model, api_key=api_key)
        elif api_key and (backend == "openai" or model in ("gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo")):
            self.backend_name = "openai"
            self.model = model or DEFAULT_MODELS["openai"]
            self._backend = OpenAIBackend(model_id=self.model, api_key=api_key)
        else:
            if backend is None and model is None:
                self.backend_name, self.model = auto_select_backend()
            else:
                self.backend_name = backend or "transformers"
                self.model = model or DEFAULT_MODELS.get(self.backend_name, model)

            bkwargs: dict = {}
            if self.backend_name == "transformers":
                bkwargs = dict(device=device, load_in_8bit=load_in_8bit,
                               load_in_4bit=load_in_4bit)
            elif self.backend_name == "ollama":
                bkwargs = dict(host=ollama_host)
            elif self.backend_name == "openai":
                bkwargs = dict(api_key=api_key)
            elif self.backend_name == "inference":
                bkwargs = dict(token=hf_token)

            self._backend = create_backend(self.backend_name, self.model, **bkwargs)

        self.system_prompt = system_prompt
        self.title = title
        self.description = description
        self.placeholder = placeholder
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.examples = examples
        self.streaming = streaming
        self._demo = None

        logger.info(f"Chatbot ready — backend={self.backend_name} model={self.model}")

    @classmethod
    def from_preset(cls, preset_name: str, **kwargs) -> "Chatbot":
        """Create a Chatbot from a named preset (e.g. 'gpt4o-mini', 'llama3', 'qwen')."""
        preset = get_preset(preset_name)
        if preset is None:
            available = ", ".join(list_presets())
            raise ValueError(f"Unknown preset '{preset_name}'. Available: {available}")
        init_kwargs = dict(
            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) -> Generator[str, None, None]:
        """Generator chat handler compatible with Gradio ChatInterface."""
        normalised = _normalize_history(history)
        if self.streaming:
            yield from self._backend.generate_stream(
                message=message,
                history=normalised,
                system_prompt=self.system_prompt,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            )
        else:
            yield self._backend.generate(
                message=message,
                history=normalised,
                system_prompt=self.system_prompt,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            )

    def build(self):
        """Build and return the Gradio app without launching."""
        from .ui import ChatUI
        self._demo = ChatUI(self).build()
        return self._demo

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

    def close(self) -> None:
        if self._demo is not None:
            self._demo.close()
        self._backend.close()

    def __enter__(self): return self
    def __exit__(self, *a): self.close()

    def __repr__(self) -> str:
        return f"Chatbot(model={self.model!r}, backend={self.backend_name!r})"
