from types import SimpleNamespace
import unittest
from unittest.mock import patch

from diffusion_cli.config import ImageGenerationRequest, UserConfig
from diffusion_cli.generation_service import GenerationService


class FakeComponent:
    load_count = 0
    to_device_count = 0
    to_cpu_count = 0

    def __init__(self, *_args, **_kwargs):
        type(self).load_count += 1

    def toDevice(self, *_args):
        type(self).to_device_count += 1

    def toCpu(self):
        type(self).to_cpu_count += 1


class FakeTextEncoder(FakeComponent):
    def encodePrompts(self, *_args):
        return "conditioning"


class FakeModel(FakeComponent):
    pass


class FakeVae(FakeComponent):
    def decode(self, latent):
        return latent


class GenerationServiceTest(unittest.TestCase):
    def setUp(self):
        for component in (FakeTextEncoder, FakeModel, FakeVae):
            component.load_count = 0
            component.to_device_count = 0
            component.to_cpu_count = 0

    def _runTwoRequests(self, model_residency):
        config = SimpleNamespace(
            prompt="a mug",
            negative_prompt="bad",
            seed=123,
            batch_size=1,
            height=64,
            width=64,
            steps=1,
            cfg=1.0,
            device="cuda",
            dtype="dtype",
            dtype_name="fp32",
            output="output.png",
            output_extension="webp",
            output_quality=80,
            tokenizer_path="tokenizer",
        )
        sources = SimpleNamespace(
            text_encoder="text",
            diffusion_model="diffusion",
            vae="vae",
        )
        request = ImageGenerationRequest(prompt="a mug")
        service = GenerationService(
            UserConfig(models=None, generation=None),
            model_residency=model_residency,
        )

        with (
            patch(
                "diffusion_cli.generation_service."
                "buildGenerationConfigFromRequest",
                return_value=config,
            ),
            patch(
                "diffusion_cli.generation_service."
                "resolveModelSourcesFromConfig",
                return_value=sources,
            ),
            patch(
                "diffusion_cli.generation_service.componentDtype",
                return_value="dtype",
            ),
            patch(
                "diffusion_cli.generation_service.ZImageTextEncoder",
                FakeTextEncoder,
            ),
            patch("diffusion_cli.generation_service.ZImageModel", FakeModel),
            patch("diffusion_cli.generation_service.ZImageVae", FakeVae),
            patch(
                "diffusion_cli.generation_service.sampleLatents",
                return_value="images",
            ),
            patch(
                "diffusion_cli.generation_service.encodeImages",
                return_value=[b"webp"],
            ) as encode_images,
        ):
            first = service.generateImages(request)
            service.generateImages(request)

        self.assertEqual(first[0].format, "webp")
        encode_images.assert_called_with("images", "webp", 80)

    def testCpuCacheLoadsComponentsOnceAcrossRequests(self):
        self._runTwoRequests("cpu-cache")

        self.assertEqual(FakeTextEncoder.load_count, 1)
        self.assertEqual(FakeModel.load_count, 1)
        self.assertEqual(FakeVae.load_count, 1)

    def testCpuCacheMovesComponentsForEachRequest(self):
        self._runTwoRequests("cpu-cache")

        self.assertEqual(FakeTextEncoder.to_device_count, 2)
        self.assertEqual(FakeModel.to_device_count, 2)
        self.assertEqual(FakeVae.to_device_count, 2)
        self.assertEqual(FakeTextEncoder.to_cpu_count, 2)
        self.assertEqual(FakeModel.to_cpu_count, 2)
        self.assertEqual(FakeVae.to_cpu_count, 2)

    def testStagedReloadsComponentsForEachRequest(self):
        self._runTwoRequests("staged")

        self.assertEqual(FakeTextEncoder.load_count, 2)
        self.assertEqual(FakeModel.load_count, 2)
        self.assertEqual(FakeVae.load_count, 2)

    def testGenerateToFilesPassesOutputFormatToImageIo(self):
        config = SimpleNamespace(
            output="output.png",
            output_extension="jpg",
            output_quality=85,
        )
        service = GenerationService(UserConfig(models=None, generation=None))

        with (
            patch(
                "diffusion_cli.generation_service."
                "buildGenerationConfigFromRequest",
                return_value=config,
            ),
            patch.object(service, "_generateTensor", return_value="images"),
            patch(
                "diffusion_cli.generation_service.saveImages",
                return_value=["output.jpg"],
            ) as save_images,
        ):
            paths = service.generateToFiles(ImageGenerationRequest("a mug"))

        self.assertEqual(paths, ["output.jpg"])
        save_images.assert_called_with("images", "output.png", "jpg", 85)


if __name__ == "__main__":
    unittest.main()