from __future__ import annotations

import asyncio
import io
import threading
import time
from dataclasses import replace

import pytest
from PIL import Image

from mcp_sd15.config import ControlModelConfig, ControlNetConfig
from mcp_sd15.errors import GenerationError, InvalidRequestError
from mcp_sd15.schema import GuidedImageRequest, InpaintImageRequest, TextToImageRequest
from mcp_sd15.service import ImageGenerationService


class FakeBackend:
    def __init__(self, delay: float = 0) -> None:
        self.calls: list[str] = []
        self.delay = delay
        self.active = 0
        self.max_active = 0
        self.guard = threading.Lock()

    def _result(self, name: str) -> Image.Image:
        with self.guard:
            self.active += 1
            self.max_active = max(self.max_active, self.active)
        time.sleep(self.delay)
        with self.guard:
            self.active -= 1
        self.calls.append(name)
        return Image.new("RGB", (4, 4), "red")

    def text_to_image(self, request: object, seed: int) -> Image.Image:
        return self._result("text")

    def guided_image(self, request: object, image: Image.Image, seed: int) -> Image.Image:
        assert image.mode == "RGB"
        return self._result("guided")

    def inpaint_image(
        self, request: object, image: Image.Image, mask: Image.Image, seed: int
    ) -> Image.Image:
        assert mask.mode == "L"
        return self._result("inpaint")


async def test_routes_all_generation_types(app_config, image_base64: str) -> None:
    backend = FakeBackend()
    service = ImageGenerationService(app_config, backend)

    text = await service.generate_image(TextToImageRequest.create("cat"))
    guided = await service.generate_guided_image(
        GuidedImageRequest.create("cat", "depth", image_base64)
    )
    inpaint = await service.inpaint_image(
        InpaintImageRequest.create("cat", image_base64, image_base64)
    )

    assert backend.calls == ["text", "guided", "inpaint"]
    for result in (text, guided, inpaint):
        assert result.png_data.startswith(b"\x89PNG")
        assert isinstance(result.metadata["seed"], int)


async def test_generation_is_strictly_serialized(app_config) -> None:
    backend = FakeBackend(delay=0.03)
    service = ImageGenerationService(app_config, backend)
    await asyncio.gather(
        *(service.generate_image(TextToImageRequest.create(f"image {i}")) for i in range(4))
    )
    assert backend.max_active == 1


async def test_disabled_controlnet_fails_before_backend(app_config, image_base64: str) -> None:
    disabled = replace(
        app_config,
        controlnet=ControlNetConfig(
            depth=ControlModelConfig(False, None, "midas"),
            openpose=app_config.controlnet.openpose,
            inpaint=app_config.controlnet.inpaint,
        ),
    )
    service = ImageGenerationService(disabled, FakeBackend())
    with pytest.raises(InvalidRequestError, match="disabled"):
        await service.generate_guided_image(GuidedImageRequest.create("cat", "depth", image_base64))


async def test_backend_error_is_sanitized(app_config) -> None:
    class BrokenBackend(FakeBackend):
        def text_to_image(self, request: object, seed: int) -> Image.Image:
            raise RuntimeError("boom")

    service = ImageGenerationService(app_config, BrokenBackend())
    with pytest.raises(GenerationError, match="boom"):
        await service.generate_image(TextToImageRequest.create("cat"))


async def test_cuda_oom_is_classified(app_config) -> None:
    class OomBackend(FakeBackend):
        def text_to_image(self, request: object, seed: int) -> Image.Image:
            raise RuntimeError("CUDA out of memory while allocating")

    service = ImageGenerationService(app_config, OomBackend())
    with pytest.raises(GenerationError, match="CUDA out of memory"):
        await service.generate_image(TextToImageRequest.create("cat"))


async def test_optional_output_is_saved(app_config) -> None:
    enabled = replace(app_config, output=replace(app_config.output, enabled=True))
    service = ImageGenerationService(enabled, FakeBackend())
    await service.generate_image(TextToImageRequest.create("cat"))
    files = list(enabled.output.directory.glob("*.png"))
    assert len(files) == 1
    with Image.open(io.BytesIO(files[0].read_bytes())) as image:
        assert image.format == "PNG"