#!/usr/bin/env python3
"""把「纯 Python wheel」重打包成「编译后 .so wheel」。

输入: 一个正常 wheel（含 <pkg>/**/*.py + 正确的依赖元数据）和 Nuitka 编译出的
整包 <pkg>.cpython-XXX-darwin.so。输出: 用单个 .so 替换掉整个 .py 包目录的
平台相关 wheel —— 解开 .whl 只能看到机器码,源码不再随包分发。

保留: dist-info（含 METADATA 的依赖声明，pip 据此拉第三方依赖）。
重写: WHEEL（标记非 purelib + 平台 tag）、RECORD（按最终文件树重算哈希）。
丢弃: <pkg>/ 下全部 .py 与 __pycache__（已被 .so 取代）。

用法:
  repack_wheel_so.py --wheel dist/lambdagent-1.3.0-py3-none-any.whl \
                     --so build/lambdagent.cpython-311-darwin.so \
                     --package lambdagent \
                     --tag cp311-cp311-macosx_15_0_arm64 \
                     --out dist-secure/
"""
from __future__ import annotations
import argparse
import base64
import csv
import hashlib
import os
import shutil
import tempfile
import zipfile
from pathlib import Path


def _hash(path: Path) -> tuple[str, int]:
    data = path.read_bytes()
    digest = base64.urlsafe_b64encode(hashlib.sha256(data).digest()).rstrip(b"=").decode()
    return f"sha256={digest}", len(data)


def repack(wheel: Path, so: Path, package: str, tag: str, out_dir: Path) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory() as td:
        root = Path(td)
        with zipfile.ZipFile(wheel) as zf:
            zf.extractall(root)

        # 1) 删除整个 .py 包目录,换成单个 .so
        pkg_dir = root / package
        if pkg_dir.is_dir():
            shutil.rmtree(pkg_dir)
        shutil.copy2(so, root / so.name)

        # 2) 定位 dist-info
        dist_info = next(p for p in root.iterdir() if p.name.endswith(".dist-info"))

        # 3) 重写 WHEEL: 非 purelib + 平台 tag
        wheel_meta = dist_info / "WHEEL"
        lines = []
        for ln in wheel_meta.read_text().splitlines():
            if ln.startswith("Root-Is-Purelib:"):
                ln = "Root-Is-Purelib: false"
            elif ln.startswith("Tag:"):
                ln = f"Tag: {tag}"
            lines.append(ln)
        if not any(l.startswith("Root-Is-Purelib:") for l in lines):
            lines.append("Root-Is-Purelib: false")
        if not any(l.startswith("Tag:") for l in lines):
            lines.append(f"Tag: {tag}")
        wheel_meta.write_text("\n".join(lines) + "\n")

        # 4) 重算 RECORD（遍历最终文件树）
        record_path = dist_info / "RECORD"
        rows = []
        for f in sorted(root.rglob("*")):
            if f.is_dir():
                continue
            rel = f.relative_to(root).as_posix()
            if rel == f"{dist_info.name}/RECORD":
                rows.append([rel, "", ""])
                continue
            h, n = _hash(f)
            rows.append([rel, h, str(n)])
        with record_path.open("w", newline="") as fp:
            csv.writer(fp).writerows(rows)

        # 5) 重新打包,文件名换平台 tag
        name = wheel.name
        # <dist>-<ver>-<pytag>-<abitag>-<plat>.whl  → 替换尾部三段为新 tag
        head = "-".join(name.split("-")[:2])  # dist-version
        out_wheel = out_dir / f"{head}-{tag}.whl"
        if out_wheel.exists():
            out_wheel.unlink()
        with zipfile.ZipFile(out_wheel, "w", zipfile.ZIP_DEFLATED) as zf:
            # RECORD 最后写
            files = [r[0] for r in rows if not r[0].endswith("/RECORD")]
            for rel in files:
                zf.write(root / rel, rel)
            zf.write(record_path, f"{dist_info.name}/RECORD")
        return out_wheel


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--wheel", required=True, type=Path)
    ap.add_argument("--so", required=True, type=Path)
    ap.add_argument("--package", required=True)
    ap.add_argument("--tag", required=True)
    ap.add_argument("--out", required=True, type=Path)
    a = ap.parse_args()
    result = repack(a.wheel, a.so, a.package, a.tag, a.out)
    print(f"OK -> {result}")


if __name__ == "__main__":
    main()