import io
import tarfile
import tomllib
from pathlib import Path
from zipfile import ZIP_DEFLATED, ZipFile


PROJECT_ROOT = Path(__file__).resolve().parent
SOURCE_DIR = PROJECT_ROOT / "src" / "snapform"
SOURCES_DIR = PROJECT_ROOT / "sources"
DIST_DIR = PROJECT_ROOT / "dist"
PATCHES_DIR = SOURCES_DIR / "patches"

PLATFORMS = {
    "windows-x64": "uv-x86_64-pc-windows-msvc.zip",
    "macos-x64": "uv-x86_64-apple-darwin.tar.gz",
    "macos-arm64": "uv-aarch64-apple-darwin.tar.gz",
    "linux-x64": "uv-x86_64-unknown-linux-gnu.tar.gz",
}
BUNDLED_ARCHIVES = {
    "sources/depth-anything-3.tar.gz": "depth-anything-3.tar.gz",
    "sources/ben2.tar.gz": "ben2.tar.gz",
}
OVERLAYS = {
    "depth-anything-3.tar.gz": "depth-anything-3",
}


def bundled_source_bytes(filename):
    """返回内置源码包内容；patches/<overlay>/ 目录中的文件会覆盖官方包对应路径。"""
    archive_path = SOURCES_DIR / filename
    overlay_dir = PATCHES_DIR / OVERLAYS.get(filename, filename)
    if not overlay_dir.is_dir():
        return archive_path.read_bytes()
    if not filename.endswith(".tar.gz"):
        raise ValueError(f"Unsupported overlay archive: {filename}")
    return _apply_tar_overlay(archive_path, overlay_dir)


def _apply_tar_overlay(archive_path, overlay_dir):
    with tarfile.open(archive_path, "r:gz") as source:
        members = source.getmembers()
        top_level = members[0].name.split("/", 1)[0]
        buffer = io.BytesIO()
        with tarfile.open(fileobj=buffer, mode="w:gz") as output:
            for member in members:
                if member.isdir() or member.islnk() or member.issym():
                    output.addfile(member)
                    continue
                relative = member.name[len(top_level) + 1:]
                overlay_file = overlay_dir / relative
                if overlay_file.is_file():
                    info = tarfile.TarInfo(member.name)
                    info.size = overlay_file.stat().st_size
                    info.mode = member.mode
                    output.addfile(info, overlay_file.open("rb"))
                else:
                    with source.extractfile(member) as extracted:
                        output.addfile(member, extracted)
    return buffer.getvalue()


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 =")]
    # Insert into the top-level section (before any `[table]` header) so the
    # field is part of the extension metadata, not a sub-key of a table.
    # Place it after the `blender_version_min` line, alongside the other
    # top-level keys (before `license`), keeping the file readable.
    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_source_files(platform):
    entries = {
        name: SOURCES_DIR / filename
        for name, filename in BUNDLED_ARCHIVES.items()
    }
    entries[f"sources/{PLATFORMS[platform]}"] = (
        SOURCES_DIR / PLATFORMS[platform]
    )
    return entries


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)
    sources = bundled_source_files(platform)
    missing = [str(path) for path in sources.values() if not path.is_file()]
    if missing:
        raise RuntimeError(
            "Missing bundled source files: " + ", ".join(missing)
        )
    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.name == "blender_manifest.toml":
                continue
            archive.write(source_path, source_path.relative_to(SOURCE_DIR))
        archive.writestr(
            "blender_manifest.toml",
            manifest_for_platform(platform),
        )
        for archive_name, source_path in sources.items():
            archive.writestr(
                archive_name,
                bundled_source_bytes(source_path.name),
            )
    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)")