#!/usr/bin/env python3
"""Download Leanstral 1.5 using only urllib (no huggingface_hub dependency)."""
import os, sys, urllib.request, json, time

MODEL = "Frosty40/Leanstral-1.5-119B-A6B-GGUF-NVFP4"
FILE = "Leanstral-1.5-119B-A6B-NVFP4.gguf"
HF_TOKEN = os.environ.get("HF_TOKEN", "")
DEST = os.path.expanduser("~/models/leanstral")
PATH = os.path.join(DEST, FILE)

os.makedirs(DEST, exist_ok=True)

# Get file metadata from HF API
api_url = f"https://huggingface.co/api/models/{MODEL}"
print(f"Querying {api_url}...")
req = urllib.request.Request(api_url)
if HF_TOKEN:
    req.add_header("Authorization", f"Bearer {HF_TOKEN}")
resp = urllib.request.urlopen(req)
data = json.loads(resp.read())

# Find the file in siblings
siblings = data.get("siblings", [])
file_info = None
for s in siblings:
    if s.get("rfilename") == FILE:
        file_info = s
        break

if file_info:
    size_gb = file_info.get("size", 0) / (1024**3)
    print(f"File size: {size_gb:.1f} GB")
else:
    print("File info not found via API, will attempt download anyway")

# Direct download URL
dl_url = f"https://huggingface.co/{MODEL}/resolve/main/{FILE}"
print(f"Downloading from {dl_url}")
print(f"To: {PATH}")

# Resume partial download
existing = 0
mode = "wb"
if os.path.exists(PATH):
    existing = os.path.getsize(PATH)
    if existing > 0:
        mode = "ab"
        print(f"Resuming at {existing / (1024**3):.1f} GB")

headers = {"User-Agent": "Mozilla/5.0"}
if existing > 0:
    headers["Range"] = f"bytes={existing}-"

req = urllib.request.Request(dl_url, headers=headers)
if HF_TOKEN:
    req.add_header("Authorization", f"Bearer {HF_TOKEN}")

resp = urllib.request.urlopen(req)
total = existing + int(resp.headers.get("Content-Length", 0))
if total == existing:
    print("Already complete!")
    sys.exit(0)

total_gb = total / (1024**3)
downloaded = existing
chunk = 32 * 1024 * 1024  # 32 MB chunks
t0 = time.time()

with open(PATH, mode) as f:
    while True:
        data = resp.read(chunk)
        if not data:
            break
        f.write(data)
        downloaded += len(data)
        elapsed = time.time() - t0
        rate = downloaded / (1024**3) / max(elapsed, 0.1)
        pct = downloaded / total * 100 if total > 0 else 0
        print(f"\r  {downloaded/(1024**3):.1f}/{total_gb:.1f} GB ({pct:.1f}%) at {rate:.1f} GB/s", end="")
        sys.stdout.flush()

print(f"\nDone! {PATH}")
print(f"Size: {os.path.getsize(PATH) / (1024**3):.1f} GB")