#!/usr/bin/env python3
"""
hub_install.py — download REAL HuggingFace weights for a zelpi hub model
registry entry (cli/hub.mjs's MODELS[]). Distinct from `hub pull`, which only
registers a model in zelpi's own simulated fleet-training system (lib/engine.ts)
and never touches a real checkpoint — this script does the actual download,
via huggingface_hub.snapshot_download, into a local cache directory.

Runs with whatever ambient Python invokes it — only needs `huggingface_hub`
(stdlib-adjacent, no torch/heavy ML deps) to be present in that environment,
which the CLI ensures via a small dedicated venv (see cli/commands.mjs).

Usage:
  python hub_install.py --repo-id lerobot/smolvla_base --out ~/.zelpi/models/smolvla
"""
from __future__ import annotations

import argparse
import json
import pathlib
import shutil
import sys


def _check_disk_space(path: pathlib.Path, required_gb: float = 2.0) -> None:
    p = path
    while not p.exists() and p.parent != p:
        p = p.parent
    try:
        free_gb = shutil.disk_usage(p).free / (1024 ** 3)
        if free_gb < required_gb:
            print(f"WARNING: only {free_gb:.1f} GB free at {p} — download may need ~{required_gb:.0f} GB", file=sys.stderr)
    except Exception as e:
        print(f"WARNING: could not check disk space: {e}", file=sys.stderr)


def main():
    ap = argparse.ArgumentParser(description="Download real HF weights for a zelpi hub model")
    ap.add_argument("--repo-id", required=True, help="HuggingFace repo id, e.g. lerobot/smolvla_base")
    ap.add_argument("--out", required=True, help="local directory to download into")
    args = ap.parse_args()

    try:
        from huggingface_hub import snapshot_download
    except ImportError as e:
        print("HUB_INSTALL_RESULT " + json.dumps({"code": "missing_dependency", "message": str(e)}))
        sys.exit(1)

    out_dir = pathlib.Path(args.out).expanduser()
    out_dir.mkdir(parents=True, exist_ok=True)
    _check_disk_space(out_dir)

    print(f"Downloading {args.repo_id} -> {out_dir} ...", file=sys.stderr)
    try:
        local_path = snapshot_download(repo_id=args.repo_id, local_dir=str(out_dir))
    except Exception as e:
        print("HUB_INSTALL_RESULT " + json.dumps({"code": "download_failed", "message": str(e)}))
        sys.exit(1)

    total_bytes = sum(f.stat().st_size for f in out_dir.rglob("*") if f.is_file())
    file_count = sum(1 for f in out_dir.rglob("*") if f.is_file())
    result = {
        "code": "success",
        "repo_id": args.repo_id,
        "path": str(local_path),
        "size_gb": round(total_bytes / (1024 ** 3), 3),
        "files": file_count,
    }
    print(f"Done: {file_count} files, {result['size_gb']} GB", file=sys.stderr)
    print("HUB_INSTALL_RESULT " + json.dumps(result))


if __name__ == "__main__":
    main()