import os, time, urllib.request
from pathlib import Path

repo = "90-008/Swift-1.5-Qwen3.8-Flash-Next-REAP-Pruned-GGUF"
ROOT = Path(__file__).resolve().parent.parent
out_dir = ROOT / "models" / "K256_reap_l2_sum"
out_dir.mkdir(parents=True, exist_ok=True)
tok_path = Path.home() / ".cache/huggingface/token"
tok = tok_path.read_text().strip() if tok_path.exists() else ""

shards = [
    "K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00001-of-00002.gguf",
    "K256_reap_l2_sum/Swift-Qwen3.8-Flash-Next-GSQ-RCO-IQ3_S-K256_reap_l2_sum-00002-of-00002.gguf",
]

for rel in shards:
    fname = rel.split("/")[-1]
    dest = out_dir / fname
    url = f"https://huggingface.co/{repo}/resolve/main/{rel}"
    existing = dest.stat().st_size if dest.exists() else 0
    headers = {"User-Agent": "Mozilla/5.0"}
    if tok:
        headers["Authorization"] = f"Bearer {tok}"
    # Check total size first
    req_head = urllib.request.Request(url, method="HEAD", headers=headers)
    with urllib.request.urlopen(req_head, timeout=30) as rh:
        total_sz = int(rh.headers.get("Content-Length", "0"))
    if existing > 0 and existing == total_sz:
        print(f"[SKIP] {fname} already complete ({total_sz / 1e9:.2f} GB)", flush=True)
        continue
    if existing > 0:
        headers["Range"] = f"bytes={existing}-"
    print(f"[START] Downloading {fname} (from {existing / 1e9:.2f} / {total_sz / 1e9:.2f} GB)...", flush=True)
    req = urllib.request.Request(url, headers=headers)
    t0 = time.monotonic()
    last_log = t0
    mode = "ab" if existing > 0 else "wb"
    cur = existing
    with urllib.request.urlopen(req, timeout=60) as r, open(dest, mode) as f:
        while True:
            chunk = r.read(16 * 1024 * 1024)
            if not chunk:
                break
            f.write(chunk)
            cur += len(chunk)
            now = time.monotonic()
            if now - last_log >= 15.0:
                mb_s = (cur - existing) / 1e6 / max(1e-3, now - t0)
                pct = 100.0 * cur / max(1, total_sz)
                print(f"  {fname}: {cur / 1e9:.2f} / {total_sz / 1e9:.2f} GB ({pct:.1f}%, {mb_s:.1f} MB/s)", flush=True)
                last_log = now
    dt = time.monotonic() - t0
    print(f"[DONE] {fname}: {cur / 1e9:.2f} GB in {dt:.1f}s", flush=True)
print("ALL K256 SHARDS DOWNLOADED!", flush=True)