import subprocess
import sys
import tempfile
import tomllib
from dataclasses import dataclass
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile


PROJECT_ROOT = Path(__file__).resolve().parent
SOURCE_DIR = PROJECT_ROOT / "src" / "anyimage"
DIST_DIR = PROJECT_ROOT / "dist"
PYTHON_VERSION = "3.11"


@dataclass(frozen=True)
class Platform:
    pip_platform: str


PLATFORM_CONFIGS = {
    "windows-x64": Platform(pip_platform="win_amd64"),
    "macos-arm64": Platform(pip_platform="macosx_12_0_arm64"),
    "linux-x64": Platform(pip_platform="manylinux2014_x86_64"),
}
PLATFORMS = tuple(PLATFORM_CONFIGS)


def project_dependencies():
    """Return Extension wheel requirements declared by the project."""
    project = tomllib.loads(
        (PROJECT_ROOT / "pyproject.toml").read_text(encoding="utf-8")
    )["project"]
    return project.get("dependencies", [])


def download_wheels(platform, destination):
    """Download default dependencies as wheels for one target platform."""
    dependencies = project_dependencies()
    if not dependencies:
        return []
    config = PLATFORM_CONFIGS[platform]
    command = [
        sys.executable,
        "-m",
        "pip",
        "download",
        *dependencies,
        "--dest",
        str(destination),
        "--find-links",
        str(SOURCE_DIR / "wheels"),
        "--only-binary=:all:",
        f"--python-version={PYTHON_VERSION}",
        f"--platform={config.pip_platform}",
    ]
    subprocess.run(command, check=True)
    wheels = sorted(destination.glob("*.whl"))
    if not wheels:
        raise RuntimeError(f"No wheels downloaded for {platform}")
    return wheels


def manifest_for_platform(platform, wheels):
    """Return manifest text for one platform and its downloaded wheels."""
    text = (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8")
    lines = text.splitlines()
    table_start = next(
        (i for i, line in enumerate(lines) if line.startswith("[")),
        len(lines),
    )
    head = lines[:table_start]
    tail = lines[table_start:]
    head = [
        line
        for line in head
        if not line.strip().startswith(("platforms =", "wheels ="))
    ]
    insert_at = next(
        (
            i + 1
            for i, line in enumerate(head)
            if line.strip().startswith("blender_version_min")
        ),
        len(head),
    )
    head.insert(insert_at, f'platforms = ["{platform}"]')
    wheel_names = ", ".join(
        f'"./wheels/{wheel.name}"' for wheel in wheels
    )
    head.insert(insert_at + 1, f"wheels = [{wheel_names}]")
    return "\n".join(head + tail) + "\n"


def _build_extension(platform, wheels):
    manifest_text = manifest_for_platform(platform, wheels)
    manifest = tomllib.loads(manifest_text)
    output_path = DIST_DIR / (
        f"{manifest['name']}.v{manifest['version']}.{platform}.zip"
    )
    DIST_DIR.mkdir(exist_ok=True)
    with ZipFile(output_path, "w", ZIP_DEFLATED) as archive:
        for source_path in sorted(SOURCE_DIR.rglob("*")):
            if not source_path.is_file():
                continue
            if "__pycache__" in source_path.parts:
                continue
            if (
                source_path.suffix.startswith(".blend")
                and source_path.suffix != ".blend"
            ):
                continue
            relative_path = source_path.relative_to(SOURCE_DIR)
            if relative_path.parts[0] == "wheels":
                continue
            if source_path.name == "blender_manifest.toml":
                continue
            archive.write(source_path, relative_path)
        archive.writestr("blender_manifest.toml", manifest_text)
        for wheel in wheels:
            archive.write(wheel, Path("wheels") / wheel.name)
    return output_path


def build_extension(platform):
    if platform not in PLATFORMS:
        raise ValueError(f"Unsupported platform: {platform}")
    with tempfile.TemporaryDirectory() as directory:
        wheels = download_wheels(platform, Path(directory))
        return _build_extension(platform, wheels)


def build_all():
    return [build_extension(platform) for platform in PLATFORMS]


if __name__ == "__main__":
    extension_paths = build_all()
    for extension_path in extension_paths:
        size_mib = extension_path.stat().st_size / (1024 * 1024)
        print(f"Built {extension_path} ({size_mib:.2f} MiB)")