import ast
import importlib.util
import io
import tarfile
import tempfile
import tomllib
import unittest
from pathlib import Path
from zipfile import ZipFile


PROJECT_ROOT = Path(__file__).resolve().parents[1]


def load_pack():
    specification = importlib.util.spec_from_file_location(
        "snapform_pack",
        PROJECT_ROOT / "pack.py",
    )
    module = importlib.util.module_from_spec(specification)
    specification.loader.exec_module(module)
    return module


class PackTest(unittest.TestCase):
    def test_builds_platform_specific_extension_zips(self):
        pack = load_pack()
        with tempfile.TemporaryDirectory() as directory:
            original_dist = pack.DIST_DIR
            pack.DIST_DIR = Path(directory)
            try:
                output_paths = pack.build_all()
            finally:
                pack.DIST_DIR = original_dist

            version = tomllib.loads(
                (pack.SOURCE_DIR / "blender_manifest.toml").read_text(
                    encoding="utf-8"
                )
            )["version"]
            self.assertEqual(
                {path.name for path in output_paths},
                {
                    f"SnapForm.v{version}.{platform}.zip"
                    for platform in pack.PLATFORMS
                },
            )
            for output_path in output_paths:
                platform = output_path.stem.split(".")[-1]
                with ZipFile(output_path) as archive:
                    names = set(archive.namelist())
                    manifest = archive.read("blender_manifest.toml").decode(
                        "utf-8"
                    )

                self.assertIn("__init__.py", names)
                self.assertIn("blender_manifest.toml", names)
                self.assertIn("image_depth.py", names)
                self.assertIn("worker.py", names)
                self.assertIn("ai_host.py", names)
                self.assertNotIn("snapform/__init__.py", names)
                self.assertNotIn("importer.py", names)
                self.assertIn("vendor/depth-anything-3.tar.gz", names)
                self.assertIn("vendor/ben2.tar.gz", names)
                self.assertIn(
                    f"vendor/{pack.PLATFORMS[platform]}",
                    names,
                )
                self.assertFalse(
                    any(name.startswith("sources/") for name in names)
                )
                bundled_uv = {
                    name
                    for name in names
                    if name.startswith("vendor/uv-")
                }
                self.assertEqual(
                    bundled_uv,
                    {f"vendor/{pack.PLATFORMS[platform]}"},
                )
                self.assertIn(f'platforms = ["{platform}"]', manifest)
                self.assertNotIn(
                    "windows-x64"
                    if platform != "windows-x64"
                    else "macos-arm64",
                    manifest.split("platforms =", 1)[-1],
                )

    def test_rejects_unknown_platform(self):
        pack = load_pack()
        with self.assertRaisesRegex(ValueError, "Unsupported platform"):
            pack.build_extension("windows-arm64")

    def test_platforms_match_manifest(self):
        pack = load_pack()
        manifest = tomllib.loads(
            (pack.SOURCE_DIR / "blender_manifest.toml").read_text(
                encoding="utf-8"
            )
        )
        self.assertEqual(set(pack.PLATFORMS), set(manifest["platforms"]))
        self.assertEqual(
            set(pack.PLATFORMS),
            {"windows-x64", "macos-arm64", "linux-x64"},
        )

    def test_bundled_depth_anything_3_has_no_xformers_dependency(self):
        pack = load_pack()
        with tarfile.open(
            fileobj=io.BytesIO(
                pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
            ),
            mode="r:gz",
        ) as archive:
            raw_pyproject = archive.extractfile(
                "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
                "pyproject.toml"
            ).read()
        pyproject = raw_pyproject.decode("utf-8")
        self.assertNotIn("xformers", pyproject)
        self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))

    def test_bundled_depth_anything_3_exports_are_lazy(self):
        pack = load_pack()
        with tarfile.open(
            fileobj=io.BytesIO(
                pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
            ),
            mode="r:gz",
        ) as archive:
            source = archive.extractfile(
                "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
                "src/depth_anything_3/utils/export/__init__.py"
            ).read().decode("utf-8")
        tree = ast.parse(source)
        module_imports = [
            ast.get_source_segment(source, node)
            for node in tree.body
            if isinstance(node, (ast.Import, ast.ImportFrom))
        ]
        self.assertEqual(
            module_imports,
            ["from depth_anything_3.specs import Prediction"],
        )

    def test_official_depth_anything_3_baseline_is_pristine(self):
        pack = load_pack()
        archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz"
        with tarfile.open(archive_path, "r:gz") as archive:
            raw_pyproject = archive.extractfile(
                "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
                "pyproject.toml"
            ).read()
            export_source = archive.extractfile(
                "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
                "src/depth_anything_3/utils/export/__init__.py"
            ).read().decode("utf-8")
        pyproject = raw_pyproject.decode("utf-8")
        self.assertIn("xformers", pyproject)
        self.assertFalse(raw_pyproject.startswith(b"\xef\xbb\xbf"))
        self.assertIn(
            "from depth_anything_3.utils.export.gs import",
            export_source,
        )

    def test_bundled_depth_anything_3_qr_falls_back_on_mps(self):
        pack = load_pack()
        with tarfile.open(
            fileobj=io.BytesIO(
                pack.bundled_artifact_bytes("depth-anything-3.tar.gz")
            ),
            mode="r:gz",
        ) as archive:
            source = archive.extractfile(
                "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
                "src/depth_anything_3/utils/ray_utils.py"
            ).read().decode("utf-8")
        self.assertIn('use_cpu = device.type == "mps"', source)
        self.assertIn("A_tilde = A_tilde.cpu()", source)
        self.assertIn("Q_tilde = Q_tilde.to(device=device)", source)

    def test_official_depth_anything_3_ray_utils_is_pristine(self):
        pack = load_pack()
        archive_path = pack.VENDOR_DIR / "depth-anything-3.tar.gz"
        with tarfile.open(archive_path, "r:gz") as archive:
            source = archive.extractfile(
                "Depth-Anything-3-3d835ec1a5802d64a8b8b15f817a1ab54809bfe4/"
                "src/depth_anything_3/utils/ray_utils.py"
            ).read().decode("utf-8")
        self.assertIn(
            "Q_tilde, R_tilde = torch.linalg.qr(A_tilde)",
            source,
        )
        self.assertNotIn("use_cpu", source)


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