"""
Diffusion model serving via the diffusers library.

Exposes:
  POST /v1/images/generations   → text-to-image (OpenAI-compatible)
"""

from __future__ import annotations

import base64
import io
import warnings
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 diffusion model and return a HandlerContext."""
    try:
        import torch  # type: ignore  # pylint: disable=import-outside-toplevel
    except ImportError:
        display.error(
            "[bold red]torch[/] is not installed.\n\n"
            "  Install with:\n\n"
            "    [bold cyan]pip install torch[/]"
        )
        raise SystemExit(1) from None

    try:
        from diffusers import (  # type: ignore  # pylint: disable=import-outside-toplevel
            AutoPipelineForImage2Image,
            AutoPipelineForText2Image,
            DiffusionPipeline,
        )
    except ImportError:
        display.error(
            "[bold red]diffusers[/] is not installed.\n\n"
            "  Install with:\n\n"
            "    [bold cyan]pip install diffusers transformers accelerate[/]"
        )
        raise SystemExit(1) from None

    if torch.cuda.is_available():
        device, dtype = "cuda", torch.float16
    elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        device, dtype = "mps", torch.float16
    else:
        device, dtype = "cpu", torch.float32

    display.step(f"Loading {repo_id}…")
    try:
        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message=".*device_type of 'cuda'.*CUDA is not available.*",
            )
            try:
                pipe = AutoPipelineForText2Image.from_pretrained(
                    str(model_dir),
                    torch_dtype=dtype,
                )
            except ValueError:
                pipe = DiffusionPipeline.from_pretrained(
                    str(model_dir),
                    torch_dtype=dtype,
                )
            pipe = pipe.to(device)
    except Exception as exc:  # pylint: disable=broad-except
        display.error(f"Failed to load model: {exc}")
        raise SystemExit(1) from None

    img2img_pipe = None
    try:
        img2img_pipe = AutoPipelineForImage2Image.from_pipe(pipe)
    except Exception:  # pylint: disable=broad-except
        try:
            with warnings.catch_warnings():
                warnings.filterwarnings(
                    "ignore",
                    message=".*device_type of 'cuda'.*CUDA is not available.*",
                )
                img2img_pipe = AutoPipelineForImage2Image.from_pretrained(
                    str(model_dir),
                    torch_dtype=dtype,
                )
                img2img_pipe = img2img_pipe.to(device)
        except Exception:  # pylint: disable=broad-except
            pass

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

    def infer_image(
        prompt: str, n: int, size: str | None = None
    ) -> list[str]:
        kwargs: dict = {"prompt": prompt}
        if size:
            w, h = (int(x) for x in size.split("x"))
            kwargs["width"] = w
            kwargs["height"] = h
        images = []
        for _ in range(n):
            result = pipe(**kwargs)
            img = result.images[0]
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            images.append(base64.b64encode(buf.getvalue()).decode())
        return images

    def infer_image_edit(
        prompt: str, image_bytes: bytes, n: int, size: str | None = None
    ) -> list[str]:
        if img2img_pipe is None:
            raise RuntimeError(
                "This model does not support image-to-image generation."
            )
        from PIL import Image  # pylint: disable=import-outside-toplevel
        src = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        kwargs: dict = {"prompt": prompt, "image": src}
        if size:
            w, h = (int(x) for x in size.split("x"))
            kwargs["width"] = w
            kwargs["height"] = h
        images = []
        for _ in range(n):
            result = img2img_pipe(**kwargs)
            img = result.images[0]
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            images.append(base64.b64encode(buf.getvalue()).decode())
        return images

    def infer(
        messages: list[dict],
        _max_tokens: int,
        _temperature: float,
    ) -> str:
        prompt = next(
            (m["content"] for m in reversed(messages) if m.get("role") == "user"),
            "",
        )
        b64 = infer_image(prompt, 1)[0]
        return f"![generated image](data:image/png;base64,{b64})"

    return _server.HandlerContext(
        model_id=repo_id,
        default_max_tokens=default_max_tokens,
        category="text-to-image",
        infer=infer,
        infer_image=infer_image,
        infer_image_edit=infer_image_edit if img2img_pipe is not None else None,
        metadata=metadata or {},
    )


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