from pathlib import Path
import tempfile
import unittest

import torch
from safetensors.torch import save_file

from diffusion_cli.config import (
    CHECKPOINT_DIFFUSION_PREFIX,
    DIFFUSION_ROLE,
    ModelSource,
)
from diffusion_cli.model_inspect import guessArchitecture, inspectModelSource


class ModelInspectTest(unittest.TestCase):
    def testPrefixedDiffusionKeysAreDetected(self):
        keys = ["model.diffusion_model.layers.0.attention.qkv.weight"]

        self.assertEqual(guessArchitecture(keys), "z_image_diffusion")

    def testInspectCheckpointSourceUsesStrippedKeys(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),
                    "model.diffusion_model.layers.0.attention.qkv.weight":
                        torch.ones(1, 1),
                    "unrelated.weight": torch.ones(1, 1),
                },
                path,
            )
            source = ModelSource(
                path=path,
                role=DIFFUSION_ROLE,
                checkpoint_prefix=CHECKPOINT_DIFFUSION_PREFIX,
            )

            summary = inspectModelSource(source)

        self.assertEqual(summary.tensor_count, 2)
        self.assertEqual(summary.source_prefix, CHECKPOINT_DIFFUSION_PREFIX)
        self.assertEqual(summary.top_level_counts, {"x_embedder": 1,
                                                    "layers": 1})
        self.assertEqual(summary.architecture_guess, "z_image_diffusion")


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