from __future__ import annotations

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

from .models import DEFAULT_MODELS, OLLAMA_PREFERRED

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# History normalisation — handles Gradio <5 (tuples) and Gradio 5+ (dicts)
# ---------------------------------------------------------------------------

def _normalize_history(history: list) -> list[tuple[str, str]]:
    if not history:
        return []
    pairs: list[tuple[str, str]] = []
    first = history[0]
    if isinstance(first, dict):
        i = 0
        while i + 1 < len(history):
            u, a = history[i], history[i + 1]
            if u.get("role") == "user" and a.get("role") == "assistant":
                pairs.append((u.get("content", ""), a.get("content", "")))
                i += 2
            else:
                i += 1
    else:
        for item in history:
            if isinstance(item, (list, tuple)) and len(item) == 2:
                pairs.append((item[0] or "", item[1] or ""))
    return pairs


# ---------------------------------------------------------------------------
# Auto-detection
# ---------------------------------------------------------------------------

def _ollama_running() -> bool:
    try:
        s = socket.create_connection(("localhost", 11434), timeout=0.5)
        s.close()
        return True
    except OSError:
        return False


def _list_ollama_models() -> list[str]:
    try:
        import urllib.request, json
        with urllib.request.urlopen("http://localhost:11434/api/tags", timeout=2) as r:
            data = json.loads(r.read())
            return [m["name"] for m in data.get("models", [])]
    except Exception:
        return []


def auto_select_backend() -> tuple[str, str]:
    """
    Try Ollama first (best quality, fully local, free).
    Falls back to local Transformers if Ollama isn't running.
    Returns (backend_name, model_id).
    """
    if _ollama_running():
        available = _list_ollama_models()
        if available:
            for preferred in OLLAMA_PREFERRED:
                for m in available:
                    if preferred in m:
                        logger.info(f"Auto-selected Ollama model: {m}")
                        return "ollama", m
            logger.info(f"Auto-selected Ollama model: {available[0]}")
            return "ollama", available[0]
        logger.info("Ollama running but no models pulled — defaulting to llama3.2:3b")
        return "ollama", "llama3.2:3b"

    logger.info("Ollama not detected — using local transformers model")
    return "transformers", DEFAULT_MODELS["transformers"]


# ---------------------------------------------------------------------------
# Base
# ---------------------------------------------------------------------------

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

    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]:
        yield self.generate(message, history,
                            system_prompt=system_prompt,
                            max_new_tokens=max_new_tokens,
                            temperature=temperature)

    def close(self) -> None:
        pass


# ---------------------------------------------------------------------------
# OpenAI
# ---------------------------------------------------------------------------

class OpenAIBackend(BaseBackend):
    """
    OpenAI backend — supports GPT-4o, GPT-4o-mini, GPT-3.5-turbo, etc.

    Usage::

        from arielai import Chatbot
        Chatbot(openai_api_key="sk-...").launch()
        Chatbot("gpt-4o", backend="openai", openai_api_key="sk-...").launch()
    """

    def __init__(self, model_id: str = "gpt-4o-mini", api_key: Optional[str] = None) -> None:
        self.model_id = model_id
        self._api_key = api_key
        self._client = None

    def _get_client(self):
        if self._client is None:
            try:
                from openai import OpenAI
                self._client = OpenAI(api_key=self._api_key)
            except ImportError as exc:
                raise ImportError(
                    "OpenAI package not found. Install with: pip install openai"
                ) from exc
        return self._client

    def _build_messages(self, message: str, history: list[tuple[str, str]],
                         system_prompt: str) -> list[dict]:
        msgs = []
        if system_prompt:
            msgs.append({"role": "system", "content": system_prompt})
        for h, a in history:
            msgs.append({"role": "user", "content": h})
            msgs.append({"role": "assistant", "content": a})
        msgs.append({"role": "user", "content": message})
        return msgs

    def generate(self, message, history, system_prompt="",
                 max_new_tokens=1024, temperature=0.7) -> str:
        resp = self._get_client().chat.completions.create(
            model=self.model_id,
            messages=self._build_messages(message, history, system_prompt),
            max_tokens=max_new_tokens,
            temperature=temperature,
        )
        return resp.choices[0].message.content

    def generate_stream(self, message, history, system_prompt="",
                        max_new_tokens=1024, temperature=0.7) -> Generator[str, None, None]:
        accumulated = ""
        for chunk in self._get_client().chat.completions.create(
            model=self.model_id,
            messages=self._build_messages(message, history, system_prompt),
            max_tokens=max_new_tokens,
            temperature=temperature,
            stream=True,
        ):
            delta = chunk.choices[0].delta.content or ""
            accumulated += delta
            yield accumulated


# ---------------------------------------------------------------------------
# Transformers
# ---------------------------------------------------------------------------

class TransformersBackend(BaseBackend):
    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._device = device
        self._load_in_8bit = load_in_8bit
        self._load_in_4bit = load_in_4bit
        self._torch_dtype = torch_dtype
        self._pipe = None
        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 exc:
            raise ImportError(
                "transformers is required. Install with: pip install arielai"
            ) from exc

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

        kwargs: dict = {}
        if self._load_in_8bit or self._load_in_4bit:
            try:
                kwargs["quantization_config"] = _build_quant_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)
        arch_str = " ".join(getattr(config, "architectures", []) or []).lower()

        if "blenderbot" in arch_str:
            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 ready: {self.model_id}")

    def _build_prompt(self, message: str, history: list[tuple[str, str]],
                      system_prompt: str) -> str:
        msgs = []
        if system_prompt:
            msgs.append({"role": "system", "content": system_prompt})
        for h, a in history:
            msgs.append({"role": "user", "content": h})
            msgs.append({"role": "assistant", "content": a})
        msgs.append({"role": "user", "content": message})
        tok = self._pipe.tokenizer
        if hasattr(tok, "apply_chat_template") and tok.chat_template:
            return tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        return _simple_prompt(msgs)

    def _gen_kwargs(self, max_new_tokens: int, temperature: float) -> dict:
        kwargs: dict = {"max_new_tokens": max_new_tokens}
        if temperature > 0:
            kwargs["do_sample"] = True
            kwargs["temperature"] = temperature
        else:
            kwargs["do_sample"] = False
        return kwargs

    def generate(self, message, history, system_prompt="",
                 max_new_tokens=512, temperature=0.7) -> str:
        self._load()
        if self._is_conversational:
            return self._gen_conv(message, history, system_prompt, max_new_tokens)
        prompt = self._build_prompt(message, history, system_prompt)
        result = self._pipe(prompt, return_full_text=False,
                            **self._gen_kwargs(max_new_tokens, temperature))
        return result[0]["generated_text"].strip()

    def _gen_conv(self, message, history, system_prompt, max_new_tokens) -> str:
        from transformers import Conversation
        conv = Conversation()
        if system_prompt:
            conv.add_message({"role": "system", "content": system_prompt})
        for h, a in history:
            conv.add_message({"role": "user", "content": h})
            conv.add_message({"role": "assistant", "content": a})
        conv.add_message({"role": "user", "content": message})
        result = self._pipe(conv, max_new_tokens=max_new_tokens)
        return result.messages[-1]["content"]

    def generate_stream(self, message, history, system_prompt="",
                        max_new_tokens=512, temperature=0.7) -> Generator[str, None, None]:
        self._load()
        if self._is_conversational:
            yield self._gen_conv(message, history, system_prompt, max_new_tokens)
            return
        try:
            from transformers import TextIteratorStreamer
            from threading import Thread

            prompt = self._build_prompt(message, history, system_prompt)
            tokenizer = self._pipe.tokenizer
            model = self._pipe.model
            inputs = tokenizer(prompt, return_tensors="pt")
            device = next(model.parameters()).device
            inputs = {k: v.to(device) for k, v in inputs.items()}

            streamer = TextIteratorStreamer(
                tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=30)

            gen_kwargs = {**inputs, "streamer": streamer,
                          **self._gen_kwargs(max_new_tokens, temperature)}
            thread = Thread(target=model.generate, kwargs=gen_kwargs, daemon=True)
            thread.start()

            accumulated = ""
            for chunk in streamer:
                accumulated += chunk
                yield accumulated
            thread.join(timeout=60)

        except Exception as exc:
            logger.warning(f"Streaming failed ({exc}), falling back to batch generation")
            yield self.generate(message, history,
                                system_prompt=system_prompt,
                                max_new_tokens=max_new_tokens,
                                temperature=temperature)

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


# ---------------------------------------------------------------------------
# Ollama
# ---------------------------------------------------------------------------

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

    def _client_inst(self):
        if self._client is None:
            try:
                import ollama
                self._client = ollama.Client(host=self.host)
            except ImportError as exc:
                raise ImportError(
                    "Install Ollama support: pip install ollama"
                ) from exc
        return self._client

    def _messages(self, message, history, system_prompt):
        msgs = []
        if system_prompt:
            msgs.append({"role": "system", "content": system_prompt})
        for h, a in history:
            msgs.append({"role": "user", "content": h})
            msgs.append({"role": "assistant", "content": a})
        msgs.append({"role": "user", "content": message})
        return msgs

    def generate(self, message, history, system_prompt="",
                 max_new_tokens=1024, temperature=0.7) -> str:
        resp = self._client_inst().chat(
            model=self.model_id,
            messages=self._messages(message, history, system_prompt),
            options={"temperature": temperature, "num_predict": max_new_tokens},
        )
        return resp["message"]["content"]

    def generate_stream(self, message, history, system_prompt="",
                        max_new_tokens=1024, temperature=0.7) -> Generator[str, None, None]:
        accumulated = ""
        for chunk in self._client_inst().chat(
            model=self.model_id,
            messages=self._messages(message, history, system_prompt),
            options={"temperature": temperature, "num_predict": max_new_tokens},
            stream=True,
        ):
            accumulated += chunk["message"]["content"]
            yield accumulated


# ---------------------------------------------------------------------------
# HuggingFace Inference API
# ---------------------------------------------------------------------------

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

    def _client_inst(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 exc:
                raise ImportError("pip install huggingface-hub") from exc
        return self._client

    def _messages(self, message, history, system_prompt):
        msgs = []
        if system_prompt:
            msgs.append({"role": "system", "content": system_prompt})
        for h, a in history:
            msgs.append({"role": "user", "content": h})
            msgs.append({"role": "assistant", "content": a})
        msgs.append({"role": "user", "content": message})
        return msgs

    def generate(self, message, history, system_prompt="",
                 max_new_tokens=512, temperature=0.7) -> str:
        resp = self._client_inst().chat_completion(
            messages=self._messages(message, history, system_prompt),
            max_tokens=max_new_tokens, temperature=temperature)
        return resp.choices[0].message.content

    def generate_stream(self, message, history, system_prompt="",
                        max_new_tokens=512, temperature=0.7) -> Generator[str, None, None]:
        accumulated = ""
        for chunk in self._client_inst().chat_completion(
            messages=self._messages(message, history, system_prompt),
            max_tokens=max_new_tokens, temperature=temperature, stream=True,
        ):
            accumulated += chunk.choices[0].delta.content or ""
            yield accumulated


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------

def create_backend(backend: str, model_id: str, **kwargs) -> BaseBackend:
    mapping = {
        "transformers": TransformersBackend,
        "ollama": OllamaBackend,
        "openai": OpenAIBackend,
        "inference": HFInferenceBackend,
    }
    cls = mapping.get(backend)
    if cls is None:
        raise ValueError(f"Unknown backend '{backend}'. Choose from: {list(mapping)}")
    valid = {k: v for k, v in kwargs.items() if k in cls.__init__.__code__.co_varnames}
    return cls(model_id=model_id, **valid)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _simple_prompt(messages: list[dict]) -> str:
    parts = []
    for m in messages:
        if m["role"] == "system":
            parts.append(f"System: {m['content']}")
        elif m["role"] == "user":
            parts.append(f"User: {m['content']}")
        else:
            parts.append(f"Assistant: {m['content']}")
    parts.append("Assistant:")
    return "\n".join(parts)


def _build_quant_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)
