#!/usr/bin/env python
"""Disk-guarded snapshot download of an open Lean-prover teacher into HF_HOME (/apps only).

    source /apps/development/rlhf/env/apps_only.sh
    PY=/apps/development/rlhf/mistral-prm-ppo/.venv/bin/python   # any venv with huggingface_hub
    align/.venv/bin/python scripts/hostguard.py run --cores 2 --mem-gb 8 --label dl --wait -- \
        $PY scripts/teachers/download_teacher.py REPO_ID [--revision SHA] [--min-free-gb 60]

Refuses (exit 2) when free space on /apps minus the repo's weight bytes would fall below --min-free-gb. Sizes come from the
hub API (`blobs=true`), so the check happens BEFORE any byte is fetched. 2 workers, no xet.
"""
from __future__ import annotations

import argparse
import os
import shutil
import sys

os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
if not os.environ.get("HF_HOME", "").startswith("/apps/"):
    sys.exit("HF_HOME must be under /apps (source env/apps_only.sh)")

from huggingface_hub import HfApi, snapshot_download  # noqa: E402


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("repo")
    ap.add_argument("--revision", default=None)
    ap.add_argument("--min-free-gb", type=float, default=60.0)
    a = ap.parse_args()
    info = HfApi().model_info(a.repo, revision=a.revision, files_metadata=True)
    total = sum(s.size or 0 for s in info.siblings)
    # a resumed download already holds part of the bytes (complete blobs and *.incomplete): count only what is still missing
    blobs = os.path.join(os.environ["HF_HOME"], "hub", "models--" + a.repo.replace("/", "--"), "blobs")
    have = sum(e.stat().st_size for e in os.scandir(blobs)) if os.path.isdir(blobs) else 0
    need = max(total - have, 0)
    free = shutil.disk_usage("/apps").free
    after = (free - need) / 1e9
    print(f"{a.repo}@{info.sha}: total {total / 1e9:.1f} GB, have {have / 1e9:.1f}, need {need / 1e9:.1f} GB, free {free / 1e9:.1f} GB, free after {after:.1f} GB (floor {a.min_free_gb})", flush=True)
    if after < a.min_free_gb:
        print("REFUSED: would cross the free-space floor", file=sys.stderr)
        return 2
    p = snapshot_download(a.repo, revision=info.sha, max_workers=2)
    print("snapshot:", p, flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())