"""Скачать колёса PyTorch с CUDA для сборки образа без сети (в docker/wheels).

Нужно, когда у машины, где собирается образ, нет выхода в интернет (или он
сломан, как у WSL за VPN), а у соседней машины есть. Скачивать можно хоть на
Windows: колёса берутся под Linux x86_64 и Python 3.10 — ровно под образ.

Тонкость: `pip download --platform` не пересчитывает маркеры окружения и на
Windows молча пропускает зависимости «только для Linux» — все библиотеки CUDA
и triton. Поэтому их список берётся из метаданных самого колеса torch и
докачивается явно, с точными версиями.

    python docker/fetch_wheels.py
"""
from __future__ import annotations

import re
import subprocess
import sys
import zipfile
from pathlib import Path

TORCH = "torch==2.9.1"
INDEX = "https://download.pytorch.org/whl/cu128"
OUT = Path(__file__).resolve().parent / "wheels"
TARGET = ["--platform", "manylinux_2_28_x86_64", "--platform", "manylinux_2_27_x86_64",
          "--platform", "manylinux_2_17_x86_64", "--platform", "manylinux2014_x86_64",
          "--python-version", "3.10", "--implementation", "cp", "--abi", "cp310",
          "--only-binary=:all:"]


def pip_download(*args: str) -> None:
    cmd = [sys.executable, "-m", "pip", "download", "--index-url", INDEX,
           "--extra-index-url", "https://pypi.org/simple", "-d", str(OUT), *args]
    subprocess.run(cmd, check=True)


def main() -> None:
    OUT.mkdir(exist_ok=True)
    pip_download(TORCH, *TARGET)
    wheel = next(OUT.glob("torch-" + TORCH.split("==")[1] + "*.whl"))
    with zipfile.ZipFile(wheel) as z:
        meta = next(n for n in z.namelist() if n.endswith(".dist-info/METADATA"))
        lines = z.read(meta).decode("utf-8").splitlines()
    linux = []
    for line in lines:
        m = re.match(r"Requires-Dist: ([A-Za-z0-9_.\-]+==[^;\s]+); platform_system == \"Linux\"$",
                     line)
        if m:
            linux.append(m.group(1))
    print(f"зависимости только для Linux: {len(linux)}")
    pip_download("--no-deps", *TARGET, *linux)
    subprocess.run([sys.executable, "-m", "pip", "download", "pip", "--only-binary=:all:",
                    "--no-deps", "-d", str(OUT)], check=True)
    size = sum(p.stat().st_size for p in OUT.glob("*.whl")) / 2**30
    print(f"готово: {len(list(OUT.glob('*.whl')))} колёс, {size:.1f} ГБ в {OUT}")


if __name__ == "__main__":
    main()