"""
Backend implementations for different model providers.
Each backend provides a generate() method that returns a response string.
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from typing import Generator, Optional

logger = logging.getLogger(__name__)


class BaseBackend(ABC):
    """Abstract base class for all ArielAI backends."""

    @abstractmethod
    def generate(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str = "",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
    ) -> str:
        """Generate a response. Returns the full response string."""

    def generate_stream(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str = "",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
    ) -> Generator[str, None, None]:
        """
        Stream a response token by token. Default implementation just returns
        the full response at once. Override for true streaming.
        """
        yield self.generate(
            message,
            history,
            system_prompt=system_prompt,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
        )

    def close(self) -> None:
        """Release resources."""


class TransformersBackend(BaseBackend):
    """
    Backend using HuggingFace Transformers pipelines.
    Supports text-generation and conversational models.
    """

    def __init__(
        self,
        model_id: str,
        device: Optional[str] = None,
        load_in_8bit: bool = False,
        load_in_4bit: bool = False,
        torch_dtype: Optional[str] = None,
    ) -> None:
        self.model_id = model_id
        self._pipe = None
        self._device = device
        self._load_in_8bit = load_in_8bit
        self._load_in_4bit = load_in_4bit
        self._torch_dtype = torch_dtype
        self._is_conversational = False

    def _load(self) -> None:
        if self._pipe is not None:
            return

        try:
            from transformers import pipeline, AutoTokenizer, AutoConfig
        except ImportError as e:
            raise ImportError(
                "transformers is required for the 'transformers' backend. "
                "Install it with: pip install arielai[torch]"
            ) from e

        logger.info(f"Loading model: {self.model_id}")

        kwargs: dict = {}
        if self._load_in_8bit or self._load_in_4bit:
            try:
                import torch

                kwargs["quantization_config"] = _build_quantization_config(
                    self._load_in_8bit, self._load_in_4bit
                )
                kwargs["device_map"] = "auto"
            except ImportError:
                logger.warning("bitsandbytes not installed; skipping quantization.")

        if self._torch_dtype:
            import torch
            dtype_map = {
                "float16": torch.float16,
                "bfloat16": torch.bfloat16,
                "float32": torch.float32,
            }
            kwargs["torch_dtype"] = dtype_map.get(self._torch_dtype, torch.float16)

        config = AutoConfig.from_pretrained(self.model_id)
        architectures = getattr(config, "architectures", []) or []
        arch_str = " ".join(architectures).lower()

        if "blenderbot" in arch_str or "conversational" in arch_str.lower():
            self._is_conversational = True
            self._pipe = pipeline(
                "conversational",
                model=self.model_id,
                device=self._device,
                **kwargs,
            )
        else:
            self._is_conversational = False
            tokenizer = AutoTokenizer.from_pretrained(self.model_id)
            self._pipe = pipeline(
                "text-generation",
                model=self.model_id,
                tokenizer=tokenizer,
                device_map=self._device or "auto",
                **kwargs,
            )

        logger.info(f"Model loaded: {self.model_id}")

    def generate(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str = "",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
    ) -> str:
        self._load()

        if self._is_conversational:
            return self._generate_conversational(
                message, history, system_prompt, max_new_tokens, temperature
            )
        return self._generate_text(
            message, history, system_prompt, max_new_tokens, temperature
        )

    def _generate_conversational(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str,
        max_new_tokens: int,
        temperature: float,
    ) -> str:
        from transformers import Conversation

        conv = Conversation()
        if system_prompt:
            conv.add_message({"role": "system", "content": system_prompt})
        for human, assistant in history:
            conv.add_message({"role": "user", "content": human})
            conv.add_message({"role": "assistant", "content": assistant})
        conv.add_message({"role": "user", "content": message})

        result = self._pipe(
            conv,
            max_new_tokens=max_new_tokens,
        )
        return result.messages[-1]["content"]

    def _generate_text(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str,
        max_new_tokens: int,
        temperature: float,
    ) -> str:
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        for human, assistant in history:
            messages.append({"role": "user", "content": human})
            messages.append({"role": "assistant", "content": assistant})
        messages.append({"role": "user", "content": message})

        tokenizer = self._pipe.tokenizer
        if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
            prompt = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        else:
            prompt = _build_simple_prompt(messages)

        result = self._pipe(
            prompt,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            do_sample=temperature > 0,
            return_full_text=False,
        )
        return result[0]["generated_text"].strip()

    def close(self) -> None:
        self._pipe = None


class OllamaBackend(BaseBackend):
    """
    Backend using Ollama for local model serving.
    Requires Ollama to be installed and running: https://ollama.ai
    """

    def __init__(self, model_id: str, host: str = "http://localhost:11434") -> None:
        self.model_id = model_id
        self.host = host
        self._client = None

    def _get_client(self):
        if self._client is None:
            try:
                import ollama

                self._client = ollama.Client(host=self.host)
            except ImportError as e:
                raise ImportError(
                    "ollama is required for the 'ollama' backend. "
                    "Install it with: pip install arielai[ollama]"
                ) from e
        return self._client

    def generate(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str = "",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
    ) -> str:
        client = self._get_client()
        messages = _build_ollama_messages(message, history, system_prompt)

        response = client.chat(
            model=self.model_id,
            messages=messages,
            options={"temperature": temperature, "num_predict": max_new_tokens},
        )
        return response["message"]["content"]

    def generate_stream(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str = "",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
    ) -> Generator[str, None, None]:
        client = self._get_client()
        messages = _build_ollama_messages(message, history, system_prompt)

        accumulated = ""
        stream = client.chat(
            model=self.model_id,
            messages=messages,
            options={"temperature": temperature, "num_predict": max_new_tokens},
            stream=True,
        )
        for chunk in stream:
            token = chunk["message"]["content"]
            accumulated += token
            yield accumulated


class HFInferenceBackend(BaseBackend):
    """
    Backend using HuggingFace Inference API (cloud-based, no GPU required).
    Requires a HuggingFace API token for private/gated models.
    """

    def __init__(self, model_id: str, token: Optional[str] = None) -> None:
        self.model_id = model_id
        self._token = token
        self._client = None

    def _get_client(self):
        if self._client is None:
            try:
                from huggingface_hub import InferenceClient

                self._client = InferenceClient(
                    model=self.model_id, token=self._token
                )
            except ImportError as e:
                raise ImportError(
                    "huggingface-hub is required for the 'inference' backend. "
                    "Install it with: pip install arielai[inference]"
                ) from e
        return self._client

    def generate(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str = "",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
    ) -> str:
        client = self._get_client()
        messages = _build_ollama_messages(message, history, system_prompt)

        response = client.chat_completion(
            messages=messages,
            max_tokens=max_new_tokens,
            temperature=temperature,
        )
        return response.choices[0].message.content

    def generate_stream(
        self,
        message: str,
        history: list[tuple[str, str]],
        system_prompt: str = "",
        max_new_tokens: int = 512,
        temperature: float = 0.7,
    ) -> Generator[str, None, None]:
        client = self._get_client()
        messages = _build_ollama_messages(message, history, system_prompt)

        accumulated = ""
        for chunk in client.chat_completion(
            messages=messages,
            max_tokens=max_new_tokens,
            temperature=temperature,
            stream=True,
        ):
            delta = chunk.choices[0].delta.content or ""
            accumulated += delta
            yield accumulated


def create_backend(
    backend: str,
    model_id: str,
    **kwargs,
) -> BaseBackend:
    """Factory function to create a backend by name."""
    backend_map = {
        "transformers": TransformersBackend,
        "ollama": OllamaBackend,
        "inference": HFInferenceBackend,
    }

    cls = backend_map.get(backend)
    if cls is None:
        raise ValueError(
            f"Unknown backend '{backend}'. "
            f"Choose from: {list(backend_map.keys())}"
        )

    valid_kwargs = {k: v for k, v in kwargs.items() if k in cls.__init__.__code__.co_varnames}
    return cls(model_id=model_id, **valid_kwargs)


def _build_ollama_messages(
    message: str,
    history: list[tuple[str, str]],
    system_prompt: str,
) -> list[dict]:
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": message})
    return messages


def _build_simple_prompt(messages: list[dict]) -> str:
    parts = []
    for msg in messages:
        role = msg["role"].capitalize()
        parts.append(f"{role}: {msg['content']}")
    parts.append("Assistant:")
    return "\n".join(parts)


def _build_quantization_config(load_in_8bit: bool, load_in_4bit: bool):
    from transformers import BitsAndBytesConfig

    if load_in_4bit:
        return BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype="float16")
    return BitsAndBytesConfig(load_in_8bit=True)
