"""Download Inflect-Nano-v1 model weights from HuggingFace.

Run this once after cloning the repo:
    python download_weights.py
"""
from __future__ import annotations

import sys
from pathlib import Path

# Use requests or urllib — try both
try:
    import requests
    _has_requests = True
except ImportError:
    import urllib.request
    import urllib.error
    _has_requests = False

WEIGHTS_DIR = Path(__file__).resolve().parent / "weights"
HF_BASE = "https://huggingface.co/owensong/Inflect-Nano-v1/resolve/main/weights"

FILES = [
    "inflect_nano_v1_acoustic.pt",
    "inflect_nano_v1_vocoder.pt",
]


def download_file(url: str, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"Downloading {dest.name}...", end=" ", flush=True)

    if _has_requests:
        resp = requests.get(url, stream=True)
        resp.raise_for_status()
        total = int(resp.headers.get("content-length", 0))
        with dest.open("wb") as f:
            downloaded = 0
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)
                downloaded += len(chunk)
                if total:
                    pct = downloaded * 100 // total
                    print(f"\rDownloading {dest.name}... {pct}%", end="", flush=True)
        print(" done.")
    else:
        urllib.request.urlretrieve(url, str(dest))
        print("done.")


def main() -> None:
    missing = [f for f in FILES if not (WEIGHTS_DIR / f).is_file()]
    if not missing:
        print("All weights already present.")
        return

    print(f"Missing {len(missing)} weight file(s). Downloading from HuggingFace...\n")
    for filename in missing:
        url = f"{HF_BASE}/{filename}"
        dest = WEIGHTS_DIR / filename
        try:
            download_file(url, dest)
        except Exception as e:
            print(f"\n  FAILED: {e}")
            print(f"  Manual download: {url}")
            sys.exit(1)

    print(f"\nAll weights downloaded to {WEIGHTS_DIR}")


if __name__ == "__main__":
    main()