from pathlib import Path
import tempfile
import unittest
from types import SimpleNamespace

from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.config import (
    CHECKPOINT_DIFFUSION_PREFIX,
    CHECKPOINT_TEXT_ENCODER_PREFIX,
    CHECKPOINT_VAE_PREFIX,
    GenerationDefaults,
    ModelPathConfig,
    UserConfig,
)
from diffusion_cli.paths import resolveModelFiles, resolveModelSources


class PathsTest(unittest.TestCase):
    def testResolveModelFiles(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            diffusion_model = root / "diffusion.safetensors"
            text_encoder = root / "text_encoder.safetensors"
            vae = root / "ae.safetensors"
            for path in (diffusion_model, text_encoder, vae):
                path.write_bytes(b"placeholder")

            args = SimpleNamespace(
                checkpoint=None,
                diffusion_model=None,
                text_encoder=None,
                vae=None,
            )
            user_config = UserConfig(
                models=ModelPathConfig(
                    diffusion_model=diffusion_model,
                    text_encoder=text_encoder,
                    vae=vae,
                ),
                generation=GenerationDefaults(),
            )
            files = resolveModelFiles(args, user_config)

        self.assertEqual(
            files.diffusion_model.name,
            "diffusion.safetensors",
        )
        self.assertEqual(files.text_encoder.name, "text_encoder.safetensors")
        self.assertEqual(files.vae.name, "ae.safetensors")

    def testConfigCheckpointResolvesAllComponents(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            checkpoint = Path(temp_dir) / "aio.safetensors"
            checkpoint.write_bytes(b"placeholder")
            args = SimpleNamespace(
                checkpoint=None,
                diffusion_model=None,
                text_encoder=None,
                vae=None,
            )
            user_config = UserConfig(
                models=ModelPathConfig(checkpoint=checkpoint),
                generation=GenerationDefaults(),
            )

            sources = resolveModelSources(args, user_config)

        self.assertEqual(sources.diffusion_model.path, checkpoint.resolve())
        self.assertEqual(sources.text_encoder.path, checkpoint.resolve())
        self.assertEqual(sources.vae.path, checkpoint.resolve())
        self.assertEqual(
            sources.diffusion_model.checkpoint_prefix,
            CHECKPOINT_DIFFUSION_PREFIX,
        )
        self.assertEqual(
            sources.text_encoder.checkpoint_prefix,
            CHECKPOINT_TEXT_ENCODER_PREFIX,
        )
        self.assertEqual(
            sources.vae.checkpoint_prefix,
            CHECKPOINT_VAE_PREFIX,
        )

    def testCliCheckpointOverridesConfigCheckpoint(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            config_checkpoint = root / "config.safetensors"
            cli_checkpoint = root / "cli.safetensors"
            config_checkpoint.write_bytes(b"placeholder")
            cli_checkpoint.write_bytes(b"placeholder")
            args = SimpleNamespace(
                checkpoint=cli_checkpoint,
                diffusion_model=None,
                text_encoder=None,
                vae=None,
            )
            user_config = UserConfig(
                models=ModelPathConfig(checkpoint=config_checkpoint),
                generation=GenerationDefaults(),
            )

            sources = resolveModelSources(args, user_config)

        self.assertEqual(sources.diffusion_model.path, cli_checkpoint.resolve())

    def testComponentPathOverridesCheckpointForOneRole(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            checkpoint = root / "aio.safetensors"
            vae = root / "ae.safetensors"
            checkpoint.write_bytes(b"placeholder")
            vae.write_bytes(b"placeholder")
            args = SimpleNamespace(
                checkpoint=checkpoint,
                diffusion_model=None,
                text_encoder=None,
                vae=vae,
            )
            user_config = UserConfig(
                models=ModelPathConfig(),
                generation=GenerationDefaults(),
            )

            sources = resolveModelSources(args, user_config)

        self.assertEqual(sources.vae.path, vae.resolve())
        self.assertIsNone(sources.vae.checkpoint_prefix)
        self.assertEqual(
            sources.diffusion_model.checkpoint_prefix,
            CHECKPOINT_DIFFUSION_PREFIX,
        )

    def testMissingFileFailsClearly(self):
        args = SimpleNamespace(
            checkpoint=None,
            diffusion_model=None,
            text_encoder=None,
            vae=None,
        )
        user_config = UserConfig(
            models=ModelPathConfig(),
            generation=GenerationDefaults(),
        )

        with self.assertRaises(DiffusionCliError) as context:
            resolveModelSources(args, user_config)

        self.assertIn(
            "Missing models.diffusion_model or models.checkpoint",
            str(context.exception),
        )

    def testMissingCheckpointFileFailsClearly(self):
        args = SimpleNamespace(
            checkpoint=Path("/missing/aio.safetensors"),
            diffusion_model=None,
            text_encoder=None,
            vae=None,
        )
        user_config = UserConfig(
            models=ModelPathConfig(),
            generation=GenerationDefaults(),
        )

        with self.assertRaises(DiffusionCliError) as context:
            resolveModelSources(args, user_config)

        self.assertIn("Missing checkpoint", str(context.exception))


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