"""
Vision model serving via HuggingFace Transformers.

Exposes:
  POST /v1/chat/completions   → accepts image URLs in message content
                                (OpenAI vision-compatible format)
"""

from __future__ import annotations

import base64
import io
import re
import urllib.request
from pathlib import Path

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


def load(
    repo_id: str,
    model_dir: Path,
    default_max_tokens: int,
    metadata: dict | None = None,
) -> _server.HandlerContext:
    """Load a vision model and return a HandlerContext."""
    try:
        from transformers import (  # type: ignore  # pylint: disable=import-outside-toplevel
            pipeline as hf_pipeline,
        )
    except ImportError:
        display.error(
            "[bold red]transformers[/] is not installed.\n\n"
            "  Install with:\n\n"
            "    [bold cyan]pip install transformers torch accelerate[/]"
        )
        raise SystemExit(1) from None

    display.step(f"Loading {repo_id}…")
    try:
        pipe = hf_pipeline(
            "image-to-text",
            model=str(model_dir),
            device_map="auto",
            trust_remote_code=True,
        )
    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}[/]")

    def _load_image(url: str):  # type: ignore[return]
        try:
            from PIL import Image  # type: ignore  # pylint: disable=import-outside-toplevel
        except ImportError as exc:
            raise RuntimeError(
                "Pillow is required for vision models: pip install Pillow"
            ) from exc

        if url.startswith("data:"):
            b64 = re.sub(r"^data:image/[^;]+;base64,", "", url)
            return Image.open(io.BytesIO(base64.b64decode(b64)))

        with urllib.request.urlopen(url) as resp:  # pylint: disable=consider-using-with
            return Image.open(io.BytesIO(resp.read()))

    def _parse_messages(messages: list[dict]) -> tuple[object | None, str]:
        image = None
        text  = ""
        for msg in messages:
            content = msg.get("content", "")
            if isinstance(content, str):
                text = content
            elif isinstance(content, list):
                for part in content:
                    if part.get("type") == "text":
                        text = part.get("text", "")
                    elif part.get("type") == "image_url":
                        url = part.get("image_url", {}).get("url", "")
                        if url:
                            image = _load_image(url)
        return image, text

    def infer(
        messages: list[dict],
        max_tokens: int,
        temperature: float,
    ) -> str:
        image, text = _parse_messages(messages)
        inp    = image if image is not None else (text or "")
        result = pipe(inp, max_new_tokens=max_tokens)
        return result[0]["generated_text"]

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


def serve(
    repo_id: str,
    model_dir: Path,
    host: str,
    port: int,
    default_max_tokens: int,
) -> None:
    """Load a vision 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)
