import unittest
from unittest.mock import patch

from diffusion_cli.cli import buildParser, main


class CliTest(unittest.TestCase):
    def testParserLeavesConfigurableDefaultsUnset(self):
        args = buildParser().parse_args(["--prompt", "test"])

        self.assertEqual(args.prompt, "test")
        self.assertIsNone(args.width)
        self.assertIsNone(args.height)
        self.assertIsNone(args.batch_size)
        self.assertIsNone(args.steps)
        self.assertIsNone(args.cfg)
        self.assertIsNone(args.output)
        self.assertIsNone(args.output_extension)
        self.assertIsNone(args.output_quality)
        self.assertIsNone(args.device)
        self.assertIsNone(args.checkpoint)
        self.assertIsNone(args.diffusion_model)
        self.assertIsNone(args.text_encoder)
        self.assertIsNone(args.vae)
        self.assertIsNone(args.tokenizer_path)
        self.assertIsNone(args.dtype)

    def testParserAcceptsExplicitConfigurableValues(self):
        args = buildParser().parse_args(
            [
                "--prompt",
                "test",
                "--steps",
                "12",
                "--cfg",
                "1.5",
                "--dtype",
                "fp16",
                "--checkpoint",
                "aio.safetensors",
                "--output-extension",
                ".jpeg",
                "--output-quality",
                "85",
            ]
        )

        self.assertEqual(args.steps, 12)
        self.assertEqual(args.cfg, 1.5)
        self.assertEqual(args.dtype, "fp16")
        self.assertEqual(args.checkpoint.name, "aio.safetensors")
        self.assertEqual(args.output_extension, ".jpeg")
        self.assertEqual(args.output_quality, 85)

    def testModelRootIsNoLongerAccepted(self):
        with patch("sys.stderr"):
            with self.assertRaises(SystemExit):
                buildParser().parse_args(["--model-root", "/tmp/models"])

    def testMissingPromptFailsForGeneration(self):
        with patch("sys.stderr"):
            self.assertEqual(main([]), 2)

    def testServeParserRequiresApiProfile(self):
        with patch("sys.stderr"):
            with self.assertRaises(SystemExit):
                buildParser().parse_args(["serve"])

    def testServeParserDefaultsToLocalCpuCacheServer(self):
        args = buildParser().parse_args(
            ["serve", "--api-profile", "sillytavern-sdcpp"]
        )

        self.assertEqual(args.command, "serve")
        self.assertEqual(args.api_profile, "sillytavern-sdcpp")
        self.assertEqual(args.host, "127.0.0.1")
        self.assertEqual(args.port, 7860)
        self.assertEqual(args.model_residency, "cpu-cache")


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