"""Serialized generation service and persistence policy."""

from __future__ import annotations

import asyncio
import logging
import secrets
import time
import uuid
from datetime import UTC, datetime
from typing import Protocol

from PIL import Image

from .config import AppConfig
from .errors import GenerationError, InvalidRequestError, SD15Error
from .images import decode_image, encode_png
from .schema import (
    GenerationResult,
    GuidedImageRequest,
    InpaintImageRequest,
    TextToImageRequest,
)

logger = logging.getLogger(__name__)


class GenerationBackend(Protocol):
    """Synchronous GPU boundary; the service runs it in a worker thread."""

    def text_to_image(self, request: TextToImageRequest, seed: int) -> Image.Image: ...

    def guided_image(
        self, request: GuidedImageRequest, image: Image.Image, seed: int
    ) -> Image.Image: ...

    def inpaint_image(
        self,
        request: InpaintImageRequest,
        image: Image.Image,
        mask: Image.Image,
        seed: int,
    ) -> Image.Image: ...


class ImageGenerationService:
    """Validate/decode requests and serialize access to the single GPU."""

    def __init__(self, config: AppConfig, backend: GenerationBackend) -> None:
        self.config = config
        self.backend = backend
        self._generation_lock = asyncio.Lock()

    async def generate_image(self, request: TextToImageRequest) -> GenerationResult:
        return await self._run("text_to_image", request)

    async def generate_guided_image(self, request: GuidedImageRequest) -> GenerationResult:
        control = self.config.controlnet.for_mode(request.control_mode)
        if not control.enabled:
            raise InvalidRequestError(f"ControlNet mode is disabled: {request.control_mode}")
        image = decode_image(request.input_image)
        return await self._run("guided_image", request, image)

    async def inpaint_image(self, request: InpaintImageRequest) -> GenerationResult:
        if not self.config.controlnet.inpaint.enabled:
            raise InvalidRequestError("ControlNet mode is disabled: inpaint")
        image = decode_image(request.input_image)
        mask = decode_image(request.mask_image, mask=True)
        return await self._run("inpaint_image", request, image, mask)

    async def _run(
        self,
        method: str,
        request: TextToImageRequest | GuidedImageRequest | InpaintImageRequest,
        *images: Image.Image,
    ) -> GenerationResult:
        seed = secrets.randbelow(2**63)
        started = time.monotonic()
        async with self._generation_lock:
            logger.info("generation started mode=%s size=%s seed=%d", method, request.size, seed)
            try:
                function = getattr(self.backend, method)
                output = await asyncio.to_thread(function, request, *images, seed)
                if not isinstance(output, Image.Image):
                    raise GenerationError("generation backend returned a non-image result")
            except SD15Error:
                raise
            except Exception as exc:
                logger.exception("generation failed mode=%s", method)
                if "out of memory" in str(exc).lower():
                    self._recover_cuda()
                    raise GenerationError("CUDA out of memory") from exc
                raise GenerationError(f"generation failed: {exc}") from exc

            metadata: dict[str, str | int | float] = {
                "seed": seed,
                "size": str(request.size),
                "mode": method,
                "timestamp": datetime.now(UTC).isoformat(),
                "elapsed_seconds": round(time.monotonic() - started, 3),
            }
            png_data = encode_png(output, metadata)
            if self.config.output.enabled:
                self._save(png_data)
            logger.info("generation completed mode=%s seed=%d", method, seed)
            return GenerationResult(png_data, metadata)

    def _save(self, png_data: bytes) -> None:
        directory = self.config.output.directory
        try:
            directory.mkdir(parents=True, exist_ok=True)
            destination = directory / f"{uuid.uuid4()}.png"
            destination.write_bytes(png_data)
        except OSError as exc:
            raise GenerationError(f"could not save generated image: {exc}") from exc

    @staticmethod
    def _recover_cuda() -> None:
        try:
            import torch

            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        except ImportError:
            pass