"""SDXL Turbo wrapper for horror image generation."""

import io

import torch

try:
    from diffusers import AutoPipelineForText2Image
except ImportError:  # pragma: no cover - exercised only in test environments
    AutoPipelineForText2Image = None  # Tests patch this attribute directly.

from server.config import config
from server.prompts import NEGATIVE_PROMPT


class AssetGenerator:
    """Generates horror images via SDXL Turbo."""

    def __init__(self, device: str | None = None, model_id: str | None = None):
        self.device = device or config.device
        self.model_id = model_id or config.models.sdxl_model_id

        if AutoPipelineForText2Image is None:
            raise RuntimeError(
                "diffusers is not installed; install diffusers to use AssetGenerator"
            )

        use_fp16 = self.device == "cuda"
        self._pipe = AutoPipelineForText2Image.from_pretrained(
            self.model_id,
            torch_dtype=torch.float16 if use_fp16 else torch.float32,
            variant="fp16" if use_fp16 else None,
        )
        if self.device == "cuda":
            self._pipe = self._pipe.to("cuda")

    def generate(self, prompt: str, seed: int | None = None) -> bytes:
        """Generate a 512x512 PNG image from a horror prompt. Returns PNG bytes."""
        if seed is None:
            seed = torch.randint(0, 2**32, (1,)).item()

        generator = torch.Generator(device=self.device).manual_seed(seed)

        result = self._pipe(
            prompt=prompt,
            negative_prompt=NEGATIVE_PROMPT,
            num_inference_steps=config.models.sdxl_steps,
            guidance_scale=config.models.sdxl_guidance_scale,
            width=config.models.sdxl_width,
            height=config.models.sdxl_height,
            generator=generator,
        )

        image = result.images[0]
        buf = io.BytesIO()
        image.save(buf, format="PNG")
        return buf.getvalue()