from pathlib import Path
import tempfile
import unittest

import torch
from safetensors.torch import save_file

from diffusion_cli.checkpoint import inspectSourceTorchDtype, loadStateDict
from diffusion_cli.config import (
    CHECKPOINT_DIFFUSION_PREFIX,
    DIFFUSION_ROLE,
    TEXT_ENCODER_ROLE,
    ModelSource,
)
from diffusion_cli.errors import DiffusionCliError


class CheckpointTest(unittest.TestCase):
    def testLoadsOnlyMatchingCheckpointPrefix(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "aio.safetensors"
            save_file(
                {
                    "model.diffusion_model.cap_embedder.1.weight":
                        torch.ones(1, 1),
                    "model.diffusion_model.x_embedder.weight":
                        torch.ones(1, 1) * 2,
                    "text_encoders.qwen3_4b.transformer.model."
                    "embed_tokens.weight": torch.ones(1, 1) * 3,
                    "unrelated.weight": torch.ones(1, 1) * 4,
                },
                path,
            )
            source = ModelSource(
                path=path,
                role=DIFFUSION_ROLE,
                checkpoint_prefix=CHECKPOINT_DIFFUSION_PREFIX,
            )

            state_dict = loadStateDict(source)

        self.assertEqual(
            sorted(state_dict),
            ["cap_embedder.1.weight", "x_embedder.weight"],
        )
        self.assertTrue(
            torch.equal(
                state_dict["x_embedder.weight"],
                torch.ones(1, 1) * 2,
            )
        )

    def testMissingCheckpointPrefixFailsClearly(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "aio.safetensors"
            save_file({"unrelated.weight": torch.ones(1, 1)}, path)
            source = ModelSource(
                path=path,
                role=TEXT_ENCODER_ROLE,
                checkpoint_prefix="missing.",
            )

            with self.assertRaises(DiffusionCliError) as context:
                loadStateDict(source)

        self.assertIn(
            "Checkpoint does not contain text_encoder",
            str(context.exception),
        )

    def testInspectsSingleCheckpointSourceDtype(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "aio.safetensors"
            save_file(
                {
                    "model.diffusion_model.x_embedder.weight":
                        torch.ones(1, 1, dtype=torch.float16),
                    "text_encoders.qwen3_4b.transformer.model."
                    "embed_tokens.weight": torch.ones(
                        1,
                        1,
                        dtype=torch.bfloat16,
                    ),
                },
                path,
            )
            source = ModelSource(
                path=path,
                role=DIFFUSION_ROLE,
                checkpoint_prefix=CHECKPOINT_DIFFUSION_PREFIX,
            )

            dtype = inspectSourceTorchDtype(source)

        self.assertEqual(dtype, torch.float16)

    def testMixedSourceDtypeIsNotInferred(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "mixed.safetensors"
            save_file(
                {
                    "a.weight": torch.ones(1, 1, dtype=torch.float16),
                    "b.weight": torch.ones(1, 1, dtype=torch.bfloat16),
                },
                path,
            )
            source = ModelSource(path=path, role=DIFFUSION_ROLE)

            dtype = inspectSourceTorchDtype(source)

        self.assertIsNone(dtype)


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