"""
Transformer model serving via HuggingFace Transformers.
"""

from __future__ import annotations

from collections.abc import Iterator
from pathlib import Path
from threading import Thread

from llmpm import display
from llmpm.core import runner
from llmpm.core import server as _server


class ModelNotTextGenerationError(Exception):
    """Raised when the loaded model is not a text-generation model."""


def _check_fp8_gpu_requirement(model_dir: Path) -> None:
    """
    Exit with a clear message if the model uses FP8 quantization but no GPU
    is available.  Transformers 5.x crashes with a cryptic KeyError when it
    tries to dequantize FP8 weights on CPU.
    """
    import json  # pylint: disable=import-outside-toplevel
    config_path = model_dir / "config.json"
    if not config_path.exists():
        return
    try:
        cfg = json.loads(config_path.read_text(encoding="utf-8"))
    except Exception:  # pylint: disable=broad-except
        return
    quant = cfg.get("quantization_config", {})
    if quant.get("quant_method") != "fp8":
        return

    try:
        import torch  # type: ignore  # pylint: disable=import-outside-toplevel
        has_gpu = torch.cuda.is_available() or (
            hasattr(torch, "xpu") and torch.xpu.is_available()
        )
    except ImportError:
        has_gpu = False

    if not has_gpu:
        display.error(
            f"[bold]{model_dir.name}[/] uses FP8 quantization which requires "
            "a CUDA or XPU GPU.\n\n"
            "  No GPU was detected on this machine.  FP8 models cannot be "
            "dequantized on CPU.\n\n"
            "  [dim]Options:[/]\n"
            "    • Run on a machine with a compatible GPU\n"
            "    • Use a non-quantized or GGUF variant of this model"
        )
        raise SystemExit(1)


def load(
    repo_id: str,
    model_dir: Path,
    default_max_tokens: int,
    metadata: dict | None = None,
) -> _server.HandlerContext:
    """Load a Transformers model and return a HandlerContext."""
    transformers = runner.load_transformers()
    _check_fp8_gpu_requirement(model_dir)

    display.step(f"Loading {repo_id}…")
    try:
        pipe = transformers.pipeline(
            "text-generation",
            model=str(model_dir),
            device_map="auto",
            trust_remote_code=True,
        )
    except ValueError as exc:
        msg = str(exc)
        if "Unrecognized model" in msg or "model_type" in msg:
            raise ModelNotTextGenerationError(model_dir.name) from exc
        display.error(f"Failed to load model: {exc}")
        raise SystemExit(1) from None
    except Exception as exc:  # pylint: disable=broad-except
        display.error(f"Failed to load model: {exc}")
        raise SystemExit(1) from None

    display.ok(f"Model ready  [dim]{repo_id}[/]")

    tokenizer = pipe.tokenizer
    model     = pipe.model

    def infer(
        messages: list[dict],
        max_tokens: int,
        temperature: float,
    ) -> str:
        gen_kwargs: dict = {
            "max_new_tokens": max_tokens,
            "return_full_text": False,
        }
        if temperature > 0:
            gen_kwargs["do_sample"] = True
            gen_kwargs["temperature"] = temperature
        output = pipe(runner.build_pipe_input(pipe, messages), **gen_kwargs)
        result = output[0]["generated_text"]
        return result[-1]["content"] if isinstance(result, list) else result

    def infer_stream(
        messages: list[dict],
        max_tokens: int,
        temperature: float,
    ) -> Iterator[str]:
        try:
            from transformers import (  # type: ignore  # pylint: disable=import-outside-toplevel
                TextIteratorStreamer,
            )
        except ImportError:
            yield infer(messages, max_tokens, temperature)
            return

        if getattr(tokenizer, "chat_template", None):
            text = tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
            )
        else:
            text = runner.build_pipe_input(pipe, messages)

        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

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

        gen_kwargs: dict = {
            **model_inputs,
            "max_new_tokens": max_tokens,
            "streamer": streamer,
        }
        if temperature > 0:
            gen_kwargs["do_sample"] = True
            gen_kwargs["temperature"] = temperature
            gen_kwargs["top_p"] = 0.95

        def _generate() -> None:
            try:
                model.generate(**gen_kwargs)
            except Exception:  # pylint: disable=broad-except
                streamer.on_finalized_text("", stream_end=True)

        Thread(target=_generate, daemon=True).start()
        try:
            yield from streamer
        except Exception:  # pylint: disable=broad-except
            return

    def embed(texts: list[str]) -> list[list[float]]:
        try:
            import torch  # type: ignore  # pylint: disable=import-outside-toplevel
        except ImportError as exc:
            raise RuntimeError("torch is required for embeddings") from exc

        encoded = tokenizer(
            texts, return_tensors="pt", padding=True, truncation=True
        ).to(model.device)

        with torch.no_grad():
            outputs = model(**encoded, output_hidden_states=True)

        hidden  = outputs.hidden_states[-1]
        mask    = encoded["attention_mask"].unsqueeze(-1).float()
        vectors = (hidden * mask).sum(1) / mask.sum(1)
        return vectors.cpu().tolist()

    return _server.HandlerContext(
        model_id=repo_id,
        default_max_tokens=default_max_tokens,
        category="text-generation",
        infer=infer,
        infer_stream=infer_stream,
        embed=embed,
        metadata=metadata or {},
    )


def serve(
    repo_id: str,
    model_dir: Path,
    host: str,
    port: int,
    default_max_tokens: int,
) -> None:
    """Load a Transformers model and start the HTTP server."""
    ctx = load(repo_id, model_dir, default_max_tokens)
    display.serve_banner(host, port)
    _server.start_multi([ctx], host=host, port=port)
