import hashlib
import tomllib
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile


PROJECT_ROOT = Path(__file__).resolve().parent
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
VENDOR_DIR = PROJECT_ROOT / "vendor"
DIST_DIR = PROJECT_ROOT / "dist"
LOCK_PATH = PROJECT_ROOT / "vendor.lock.toml"


def load_vendor_lock():
    """Return the locked third-party artifact metadata."""
    with LOCK_PATH.open("rb") as stream:
        lock = tomllib.load(stream)
    if lock.get("schema_version") != 1:
        raise RuntimeError("Unsupported vendor.lock.toml schema")
    return lock


VENDOR_LOCK = load_vendor_lock()
PLATFORMS = {
    platform: artifact["filename"]
    for platform, artifact in VENDOR_LOCK["platforms"].items()
}


def file_sha256(path):
    """Return the SHA-256 digest for a file."""
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def validate_artifact(path, expected_sha256):
    """Require an artifact to exist and match its locked digest."""
    if not path.is_file():
        raise RuntimeError(
            f"Missing vendor artifact: {path}. "
            "Run: uv run python prepare_vendor.py"
        )
    actual = file_sha256(path)
    if actual != expected_sha256:
        raise RuntimeError(
            f"Vendor artifact hash mismatch: {path}. "
            "Run: uv run python prepare_vendor.py"
        )


def manifest_for_platform(platform):
    """Return the manifest text declaring only the given platform."""
    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 =")]
    # 插入顶层区域，避免 platforms 被解析为后续 table 的子键。
    # 放在 blender_version_min 后面，使扩展元数据保持易读。
    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}"]')
    lines = head + tail
    return "\n".join(lines) + "\n"


def bundled_vendor_files(platform):
    return {
        f"vendor/{PLATFORMS[platform]}": VENDOR_DIR / PLATFORMS[platform]
    }


def validate_vendor_files(platform):
    """Validate common archives and the selected uv archive."""
    uv_artifact = VENDOR_LOCK["platforms"][platform]
    validate_artifact(
        VENDOR_DIR / uv_artifact["filename"],
        uv_artifact["sha256"],
    )


def build_extension(platform):
    if platform not in PLATFORMS:
        raise ValueError(f"Unsupported platform: {platform}")
    manifest = tomllib.loads(manifest_for_platform(platform))
    output_path = DIST_DIR / (
        f"{manifest['name']}.v{manifest['version']}.{platform}.zip"
    )
    DIST_DIR.mkdir(exist_ok=True)
    vendor_files = bundled_vendor_files(platform)
    validate_vendor_files(platform)
    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
            relative_path = source_path.relative_to(SOURCE_DIR)
            if relative_path.parts[0] == "vendor":
                continue
            if source_path.name == "blender_manifest.toml":
                continue
            archive.write(source_path, relative_path)
        archive.writestr(
            "blender_manifest.toml",
            manifest_for_platform(platform),
        )
        for archive_name, source_path in vendor_files.items():
            archive.writestr(
                archive_name,
                source_path.read_bytes(),
            )
    return output_path


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)")