"""Cache trusted installed distributions from the existing ARM64 base image.

Rebuild wheel RECORDs from installed package files, retaining upstream WHEEL
tags and entry points. This avoids downloading the same native libraries again.
Only named scientific dependencies are copied; CUDA/PyTorch are not included.
"""
import importlib.metadata as metadata
import json
import zipfile
from pathlib import Path

from wheel.wheelfile import WheelFile

workspace = Path(__file__).resolve().parents[3]
output = workspace / ".cache/seed-wheels"
output.mkdir(parents=True, exist_ok=True)
records = []
for name in (
    "pyarrow", "scipy", "pandas", "matplotlib", "scikit-learn", "wandb",
    "contourpy", "kiwisolver", "fonttools", "sentry-sdk", "threadpoolctl", "joblib",
):
    dist = metadata.distribution(name)
    normalized = dist.metadata["Name"].replace("-", "_")
    wheel_metadata = dist.read_text("WHEEL")
    tags = [line.removeprefix("Tag: ") for line in wheel_metadata.splitlines() if line.startswith("Tag: ")]
    tag = next(tag for tag in tags if tag.split("-")[0] != "py2")
    destination = output / f"{normalized}-{dist.version}-{tag}.whl"
    if not destination.exists():
        with WheelFile(str(destination), "w", compression=zipfile.ZIP_STORED) as wheel:
            for item in dist.files or []:
                path = Path(str(item))
                if ".." in path.parts or path.suffix == ".pyc" or path.name in {"RECORD", "INSTALLER", "direct_url.json"}:
                    continue
                source = Path(dist.locate_file(item))
                if source.is_file():
                    wheel.write(str(source), str(path))
    records.append({"name": name, "version": dist.version, "wheel": str(destination), "source": str(dist.locate_file(""))})
    print(f"Cached {name}=={dist.version}", flush=True)
(output / "seed-constraints.txt").write_text("".join(f"{r['name']}=={r['version']}\n" for r in records))
(output / "seed-manifest.json").write_text(json.dumps({"source_image": "auto2606.sqsh", "packages": records}, indent=2) + "\n")