#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""离线依赖取件 (在**有网**的机器上跑一次, 产物随包/U 盘带到内网机器).

三类:
  wheels  各平台 Python 依赖轮子 (win_amd64 / linux_x86_64 / linux_aarch64 / macos_arm64)
  python  便携 CPython 3.12 运行时 (麒麟/统信自带多为 3.7/3.9, 跑不了本项目依赖)
  ollama  Ollama 离线安装包 (Windows .exe ~1.5G / Linux .tar.zst ~1.4G; **体积大, 不进主程序包**, 单独带)
用法: python scripts/fetch_offline_deps.py [wheels|python|ollama|all] [--arch x86_64,aarch64]
产物带 SHA256 记到 vendor/MANIFEST.json; 断点续传 (再跑一次自动接着下)。模型不在此列 —— 见 §模型."""
from __future__ import annotations
import argparse, hashlib, json, subprocess, sys, time, urllib.request
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
PY_TAG = "20260901"; PY_VER = "3.12.14"
PY_BASE = f"https://github.com/astral-sh/python-build-standalone/releases/download/{PY_TAG}"
PY_FILES = {"linux_x86_64": f"cpython-{PY_VER}+{PY_TAG}-x86_64-unknown-linux-gnu-install_only_stripped.tar.gz",
            "linux_aarch64": f"cpython-{PY_VER}+{PY_TAG}-aarch64-unknown-linux-gnu-install_only_stripped.tar.gz",
            "win_amd64": f"cpython-{PY_VER}+{PY_TAG}-x86_64-pc-windows-msvc-install_only_stripped.tar.gz"}
OL_VER = "v0.33.3"; OL_BASE = f"https://github.com/ollama/ollama/releases/download/{OL_VER}"
OL_FILES = {"win_amd64": "OllamaSetup.exe", "linux_x86_64": "ollama-linux-amd64.tar.zst", "linux_aarch64": "ollama-linux-arm64.tar.zst"}
WHEEL_PLAT = {"win_amd64": ["win_amd64"],
              "linux_x86_64": ["manylinux_2_28_x86_64", "manylinux_2_27_x86_64", "manylinux_2_26_x86_64", "manylinux_2_24_x86_64", "manylinux_2_17_x86_64", "manylinux2014_x86_64"],
              "linux_aarch64": ["manylinux_2_28_aarch64", "manylinux_2_27_aarch64", "manylinux_2_26_aarch64", "manylinux_2_24_aarch64", "manylinux_2_17_aarch64", "manylinux2014_aarch64"]}


def sha256(p: Path) -> str:
    h = hashlib.sha256()
    with p.open("rb") as f:
        for b in iter(lambda: f.read(1 << 20), b""): h.update(b)
    return h.hexdigest()


def get(url: str, dst: Path):
    """断点续传下载 (curl); 已完整则跳过."""
    dst.parent.mkdir(parents=True, exist_ok=True)
    r = subprocess.run(["curl", "-sL", "-C", "-", "--retry", "5", "--retry-delay", "3", "-o", str(dst), url])
    if r.returncode: raise SystemExit(f"下载失败 {url} (rc={r.returncode}); 再跑一次会续传")
    return dst


def wheels(targets):
    for t in targets:
        d = ROOT / "wheels" / t; d.mkdir(parents=True, exist_ok=True)
        cmd = [sys.executable, "-m", "pip", "download", "-r", str(ROOT / "requirements.txt"), "--only-binary=:all:",
               "--python-version", "3.12", "--implementation", "cp", "-d", str(d), "-q"]
        for p in WHEEL_PLAT[t]: cmd += ["--platform", p]
        cmd += ["--abi", "cp312", "--abi", "abi3", "--abi", "none"]
        print(f"  轮子 {t} …", flush=True); subprocess.run(cmd, check=True)
        print(f"    {len(list(d.glob('*.whl')))} 个")


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("what", nargs="?", default="all", choices=["wheels", "python", "ollama", "all"])
    ap.add_argument("--targets", default="win_amd64,linux_x86_64,linux_aarch64"); a = ap.parse_args()
    ts = [t.strip() for t in a.targets.split(",") if t.strip()]
    man_p = ROOT / "vendor/MANIFEST.json"; man = json.loads(man_p.read_text(encoding="utf-8")) if man_p.exists() else {}
    if a.what in ("wheels", "all"): wheels(ts)
    if a.what in ("python", "all"):
        for t in ts:
            f = PY_FILES.get(t)
            if not f: continue
            p = get(f"{PY_BASE}/{f}", ROOT / "vendor/python" / f); man[f"python/{f}"] = dict(sha256=sha256(p), bytes=p.stat().st_size, url=f"{PY_BASE}/{f}")
            print(f"  便携 Python {t}: {p.stat().st_size/1e6:.0f} MB")
    if a.what in ("ollama", "all"):
        for t in ts:
            f = OL_FILES.get(t)
            if not f: continue
            print(f"  Ollama {t} ({f}) 下载中 (1.4–1.6 GB, 可中断续传) …", flush=True)
            p = get(f"{OL_BASE}/{f}", ROOT / "vendor/ollama" / f); man[f"ollama/{f}"] = dict(sha256=sha256(p), bytes=p.stat().st_size, url=f"{OL_BASE}/{f}", version=OL_VER)
            print(f"    {p.stat().st_size/1e6:.0f} MB")
    man_p.parent.mkdir(parents=True, exist_ok=True); man["_built"] = time.strftime("%Y-%m-%d %H:%M")
    man_p.write_text(json.dumps(man, ensure_ascii=False, indent=1), encoding="utf-8"); print(f"清单 → {man_p}")
    print("模型不在此列: 在有网机器 `ollama pull qwen3:8b deepseek-r1:14b bge-m3`, 再把 ~/.ollama/models (Windows %USERPROFILE%\\.ollama\\models) 整个目录拷到内网机同位置。")


if __name__ == "__main__":
    main()