import base64
import json
import unittest

from diffusion_cli.api_profiles import API_PROFILES
from diffusion_cli.generation_service import GeneratedImage
from diffusion_cli.server import ServerConfig, createApp


class FakeGenerationService:
    def __init__(self):
        self.requests = []

    def generateImages(self, request):
        self.requests.append(request)
        return [GeneratedImage(b"\x89PNG\r\n\x1a\nimage", "png", 123)]


class ApiProfilesTest(unittest.TestCase):
    def makeClient(self):
        service = FakeGenerationService()
        app = createApp(
            ServerConfig("sillytavern-sdcpp", "127.0.0.1", 7860, "staged"),
            service,
        )
        return app.test_client(), service

    def testRegistryListsSillyTavernProfile(self):
        self.assertIn("sillytavern-sdcpp", API_PROFILES)

    def testModelsEndpointReturnsOpenAiModelList(self):
        client, _service = self.makeClient()

        response = client.get("/v1/models")

        self.assertEqual(response.status_code, 200)
        self.assertEqual(response.json["data"][0]["id"], "z-image-local")
        self.assertEqual(response.json["data"][0]["object"], "model")

    def testImageGenerationsOptionsReturnsNoContent(self):
        client, _service = self.makeClient()

        response = client.open("/v1/images/generations", method="OPTIONS")

        self.assertEqual(response.status_code, 204)
        self.assertEqual(response.headers["Allow"], "OPTIONS, POST")

    def testTxt2ImgMinimalRequestReturnsBareBase64Image(self):
        client, service = self.makeClient()

        response = client.post("/sdapi/v1/txt2img", json={"prompt": "a mug"})

        self.assertEqual(response.status_code, 200)
        self.assertEqual(service.requests[0].prompt, "a mug")
        self.assertEqual(response.json["parameters"], {"prompt": "a mug"})
        image_text = response.json["images"][0]
        self.assertNotIn("data:image", image_text)
        self.assertEqual(
            base64.b64decode(image_text),
            b"\x89PNG\r\n\x1a\nimage",
        )
        self.assertIsInstance(response.json["info"], str)
        self.assertEqual(json.loads(response.json["info"])["seed"], 123)

    def testTxt2ImgMapsFullSillyTavernPayload(self):
        client, service = self.makeClient()

        response = client.post(
            "/sdapi/v1/txt2img",
            json={
                "prompt": "a mug",
                "negative_prompt": "bad",
                "width": 512,
                "height": 768,
                "steps": 12,
                "cfg_scale": 1.5,
                "seed": 42,
                "batch_size": 2,
                "output_extension": "jpeg",
                "output_quality": 80,
                "sampler_name": "euler",
                "scheduler": "normal",
                "clip_skip": 1,
                "model": "ignored",
            },
        )

        self.assertEqual(response.status_code, 200)
        request = service.requests[0]
        self.assertEqual(request.negative_prompt, "bad")
        self.assertEqual(request.width, 512)
        self.assertEqual(request.height, 768)
        self.assertEqual(request.steps, 12)
        self.assertEqual(request.cfg, 1.5)
        self.assertEqual(request.seed, 42)
        self.assertEqual(request.batch_size, 2)
        self.assertEqual(request.output_extension, "jpg")
        self.assertEqual(request.output_quality, 80)

    def testTxt2ImgSeedMinusOneRequestsRandomSeed(self):
        client, service = self.makeClient()

        response = client.post(
            "/sdapi/v1/txt2img",
            json={"prompt": "a mug", "seed": -1},
        )

        self.assertEqual(response.status_code, 200)
        self.assertIsNone(service.requests[0].seed)

    def testTxt2ImgEmptyPromptReturnsBadRequest(self):
        client, _service = self.makeClient()

        response = client.post("/sdapi/v1/txt2img", json={"prompt": ""})

        self.assertEqual(response.status_code, 400)
        self.assertEqual(response.json["error"], "bad_request")

    def testTxt2ImgUnsupportedInitImagesReturnsBadRequest(self):
        client, _service = self.makeClient()

        response = client.post(
            "/sdapi/v1/txt2img",
            json={"prompt": "a mug", "init_images": ["abc"]},
        )

        self.assertEqual(response.status_code, 400)
        self.assertIn("init_images", response.json["message"])

    def testTxt2ImgInvalidWidthReturnsBadRequest(self):
        client, _service = self.makeClient()

        response = client.post(
            "/sdapi/v1/txt2img",
            json={"prompt": "a mug", "width": 513},
        )

        self.assertEqual(response.status_code, 400)

    def testSillyTavernEndpointSequence(self):
        client, _service = self.makeClient()

        options = client.open("/v1/images/generations", method="OPTIONS")
        models = client.get("/v1/models")
        txt2img = client.post("/sdapi/v1/txt2img", json={"prompt": "a mug"})

        self.assertEqual(options.status_code, 204)
        self.assertEqual(models.status_code, 200)
        self.assertEqual(txt2img.status_code, 200)


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