"""Optional diffusion backends for image artifact generation."""
from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any

import httpx


@dataclass(frozen=True)
class DiffusionImageResult:
    """Single generated image payload."""

    base64_data: str
    mime_type: str
    metadata: dict[str, Any]


class DiffusionBackend:
    """Abstract diffusion image backend."""

    async def generate_image(
        self,
        *,
        prompt: str,
        negative_prompt: str | None = None,
        width: int = 1024,
        height: int = 1024,
        steps: int = 28,
        guidance_scale: float = 7.0,
        seed: int | None = None,
        model: str | None = None,
        **kwargs: Any,
    ) -> DiffusionImageResult:
        raise NotImplementedError


class Automatic1111DiffusionBackend(DiffusionBackend):
    """Stable Diffusion WebUI / Forge txt2img backend."""

    def __init__(
        self,
        *,
        base_url: str | None = None,
        api_key: str | None = None,
        timeout: float = 180.0,
        client: httpx.AsyncClient | None = None,
    ) -> None:
        self.base_url = (base_url or os.getenv("LOGOS_ADK_DIFFUSION_BASE_URL") or "http://127.0.0.1:7860").rstrip("/")
        self.api_key = api_key or os.getenv("LOGOS_ADK_DIFFUSION_API_KEY")
        self.timeout = float(timeout)
        self._client = client

    async def generate_image(
        self,
        *,
        prompt: str,
        negative_prompt: str | None = None,
        width: int = 1024,
        height: int = 1024,
        steps: int = 28,
        guidance_scale: float = 7.0,
        seed: int | None = None,
        model: str | None = None,
        sampler_name: str | None = None,
        cfg_scale: float | None = None,
        **kwargs: Any,
    ) -> DiffusionImageResult:
        if not isinstance(prompt, str) or not prompt.strip():
            raise ValueError("prompt must be a non-empty string")
        payload: dict[str, Any] = {
            "prompt": prompt.strip(),
            "negative_prompt": negative_prompt or "",
            "width": max(64, int(width)),
            "height": max(64, int(height)),
            "steps": max(1, int(steps)),
            "cfg_scale": float(cfg_scale if cfg_scale is not None else guidance_scale),
            "sampler_name": sampler_name or "DPM++ 2M Karras",
        }
        if seed is not None:
            payload["seed"] = int(seed)
        if model is not None:
            payload["override_settings"] = {"sd_model_checkpoint": model}
        payload.update(kwargs)

        headers: dict[str, str] = {}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        owns_client = self._client is None
        client = self._client or httpx.AsyncClient(timeout=self.timeout)
        try:
            response = await client.post(
                f"{self.base_url}/sdapi/v1/txt2img",
                json=payload,
                headers=headers or None,
            )
            response.raise_for_status()
        except httpx.HTTPError as exc:
            raise RuntimeError(f"Diffusion backend request failed: {exc}") from exc
        finally:
            if owns_client:
                await client.aclose()

        data = response.json()
        images = data.get("images")
        if not isinstance(images, list) or not images or not isinstance(images[0], str) or not images[0].strip():
            raise RuntimeError("Diffusion backend returned no image payload")
        return DiffusionImageResult(
            base64_data=images[0],
            mime_type="image/png",
            metadata={
                "backend": "automatic1111",
                "base_url": self.base_url,
                "prompt": prompt.strip(),
                "negative_prompt": negative_prompt or "",
                "width": payload["width"],
                "height": payload["height"],
                "steps": payload["steps"],
                "cfg_scale": payload["cfg_scale"],
                "seed": payload.get("seed"),
                "model": model,
                "sampler_name": payload["sampler_name"],
                "info": data.get("info"),
                "parameters": payload,
            },
        )


def build_diffusion_backend(
    *,
    backend: str | None = None,
    base_url: str | None = None,
    api_key: str | None = None,
    timeout: float = 180.0,
    client: httpx.AsyncClient | None = None,
) -> DiffusionBackend:
    """Build a diffusion backend from explicit args or environment."""
    selected = (backend or os.getenv("LOGOS_ADK_DIFFUSION_BACKEND") or "automatic1111").strip().lower()
    if selected in {"automatic1111", "stable-diffusion-webui", "forge"}:
        return Automatic1111DiffusionBackend(
            base_url=base_url,
            api_key=api_key,
            timeout=timeout,
            client=client,
        )
    raise ValueError(f"Unsupported diffusion backend '{selected}'")


__all__ = [
    "Automatic1111DiffusionBackend",
    "DiffusionBackend",
    "DiffusionImageResult",
    "build_diffusion_backend",
]
