import importlib.util
import tempfile
import tomllib
import unittest
from pathlib import Path
from unittest.mock import patch
from zipfile import ZipFile


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


def load_pack():
    specification = importlib.util.spec_from_file_location(
        "anyimage_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:
            directory_path = Path(directory)
            wheel_path = directory_path / "blendjob-0.1.13-py3-none-any.whl"
            wheel_path.write_bytes(b"wheel")
            original_dist = pack.DIST_DIR
            pack.DIST_DIR = directory_path / "dist"
            try:
                with patch.object(
                    pack,
                    "download_wheels",
                    return_value=[wheel_path],
                ):
                    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"AnyImage.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("operators/depth_plane/operator.py", names)
                self.assertIn("operators/image.py", names)
                self.assertIn("operators/balloon/operator.py", names)
                self.assertIn("operators/surface.py", names)
                self.assertIn(
                    "assets/O_AnyImage.blend",
                    names,
                )
                self.assertIn("server/media/input.py", names)
                self.assertIn("server/geometry/projection.py", names)
                self.assertIn("server/geometry/grid_mesh.py", names)
                self.assertIn("server/geometry/surface.py", names)
                self.assertIn("server/jobs/background.py", names)
                self.assertIn("server/models/ben2.py", names)
                self.assertIn("server/models/da3.py", names)
                self.assertIn("server/models/moge2.py", names)
                self.assertIn("server/jobs/depth.py", names)
                self.assertIn("server/jobs/surface.py", names)
                self.assertIn("server/jobs/debug_da3.py", names)
                self.assertIn("server/jobs/debug_moge2.py", names)
                self.assertIn("server/jobs/debug_ben2.py", names)
                self.assertNotIn("server/processing.py", names)
                self.assertNotIn("server/progress.py", names)
                self.assertIn("model_catalog.py", names)
                self.assertIn("server/model_download.py", names)
                self.assertIn("server/models/onnx_ben2.py", names)
                self.assertIn("server/models/onnx_da3.py", names)
                self.assertIn("server/models/onnx_moge2.py", names)
                self.assertNotIn("server/da3.py", names)
                self.assertNotIn("server/moge2.py", names)
                self.assertIn("server/__init__.py", names)
                self.assertIn(
                    "wheels/blendjob-0.1.13-py3-none-any.whl",
                    names,
                )
                self.assertFalse(
                    any(name.startswith("blendjob/") for name in names)
                )
                self.assertNotIn("install_env.py", names)
                self.assertNotIn("anyimage/__init__.py", names)
                self.assertNotIn("importer.py", names)
                self.assertFalse(
                    any(name.startswith("sources/") for name in names)
                )
                self.assertFalse(
                    any(name.startswith("vendor/") for name in names)
                )
                self.assertIn(f'platforms = ["{platform}"]', manifest)
                self.assertNotIn(
                    "windows-x64"
                    if platform != "windows-x64"
                    else "macos-arm64",
                    manifest.split("platforms =", 1)[-1],
                )

    def test_reads_extension_dependencies_from_pyproject(self):
        pack = load_pack()
        self.assertEqual(pack.project_dependencies(), ["blendjob==0.1.13"])

    def test_downloads_declared_wheels_for_target_platform(self):
        pack = load_pack()
        with tempfile.TemporaryDirectory() as directory:
            destination = Path(directory)

            def fake_run(command, check):
                self.assertTrue(check)
                (destination / "blendjob-0.1.13-py3-none-any.whl").write_bytes(
                    b"wheel"
                )

            with patch.object(pack.subprocess, "run", side_effect=fake_run) as run:
                wheels = pack.download_wheels("linux-x64", destination)

        command = run.call_args.args[0]
        self.assertIn("blendjob==0.1.13", command)
        self.assertIn("--python-version=3.11", command)
        self.assertIn("--platform=manylinux2014_x86_64", command)
        self.assertEqual(
            [wheel.name for wheel in wheels],
            ["blendjob-0.1.13-py3-none-any.whl"],
        )
    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"},
        )

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