#!/usr/bin/env python3
"""从 ModelScope 镜像下载模型权重（HF 直连被墙时的替代路径）。

ModelScope 的下载地址是单文件 API，不支持 huggingface_hub，所以自己拉清单 + 断点续传。
已完整的文件按大小跳过，中断后重跑即可续传。
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import urllib.parse
import urllib.request

API = "https://modelscope.cn/api/v1/models/{repo}/repo/files?Revision={rev}"
FILE_API = "https://modelscope.cn/api/v1/models/{repo}/repo?Revision={rev}&FilePath={path}"


def list_files(repo: str, rev: str) -> list[dict]:
    url = API.format(repo=repo, rev=rev)
    with urllib.request.urlopen(url, timeout=60) as r:
        data = json.load(r)["Data"]["Files"]
    return [f for f in data if f["Type"] == "blob" and f["Size"] > 0]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="Cloudflare/clef-flash")
    ap.add_argument("--rev", default="master")
    ap.add_argument("--out", default="models/clef-flash")
    ap.add_argument("--verify", action="store_true", help="只校验完整性，不下载")
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    files = list_files(args.repo, args.rev)
    total = sum(f["Size"] for f in files)
    print(f"[list] {len(files)} 个文件，合计 {total/1024**3:.1f} GB")

    if args.verify:
        missing = [f["Name"] for f in files
                   if not (os.path.exists(os.path.join(args.out, f["Name"]))
                           and os.path.getsize(os.path.join(args.out, f["Name"])) == f["Size"])]
        if missing:
            print(f"[verify] 未完成，缺 {len(missing)} 个文件：{missing}", file=sys.stderr)
            sys.exit(1)
        print(f"[verify] 全部完整 {total/1024**3:.1f} GB")
        return

    done_bytes = 0
    for i, f in enumerate(files, 1):
        name, size = f["Name"], f["Size"]
        dest = os.path.join(args.out, name)
        if os.path.exists(dest) and os.path.getsize(dest) == size:
            print(f"[skip] {name} 已完整")
            done_bytes += size
            continue
        url = FILE_API.format(repo=args.repo, rev=args.rev,
                              path=urllib.parse.quote(name, safe=""))
        print(f"[get ] ({i}/{len(files)}) {name} {size/1024**3:.2f} GB", flush=True)
        # -C - 续传；--retry 应对镜像抖动
        rc = subprocess.call([
            "curl", "-L", "-C", "-", "--retry", "8", "--retry-delay", "5",
            "--connect-timeout", "30", "-sS", "-o", dest, url,
        ])
        got = os.path.getsize(dest) if os.path.exists(dest) else 0
        if rc != 0 or got != size:
            print(f"[fail] {name}: rc={rc} got={got} want={size}", file=sys.stderr)
            sys.exit(1)
        done_bytes += size
        print(f"[ ok ] {name}  ({done_bytes/1024**3:.1f}/{total/1024**3:.1f} GB)", flush=True)

    print(f"DONE {args.out} {total/1024**3:.1f} GB")


if __name__ == "__main__":
    main()