"""
Thunder Compute (H100) üzerinde HuggingFace'ten V5 verisi + ckpt + kod indir.

Bootstrap script'i — sıfırdan kurulum:
  1. Veriyi data/ icine indir
  2. Checkpoint'leri runs/tr-200m-v5/ icine indir
  3. Kodu mevcut dizine indir (varsa overwrite)

Kullanim:
    pip install huggingface_hub
    python hf_pull_v5.py --user musabc        # her seyi indir
    python hf_pull_v5.py --user musabc --no-ckpt  # sadece data + code
    python hf_pull_v5.py --resume-token <TOKEN>   # private repolar icin

NOT: ~30GB indirme. Thunder Compute disk yeterli mi kontrol et:
     df -h .
"""

import argparse
import os
import sys
from pathlib import Path

try:
    from huggingface_hub import snapshot_download, hf_hub_download
except ImportError:
    print("huggingface_hub yok, kuruluyor...")
    os.system("pip install -U huggingface_hub")
    from huggingface_hub import snapshot_download, hf_hub_download

REPO_BASE = "nanogpt-tr-v5"
DEFAULT_USER = "musabc"

ROOT = Path(__file__).parent if "__file__" in dir() else Path(".")
DATA_DIR = ROOT / "data"
RUN_DIR = ROOT / "runs" / "tr-200m-v5"


def fmt_size(b):
    for u in ["B", "KB", "MB", "GB"]:
        if b < 1024:
            return f"{b:.1f} {u}"
        b /= 1024
    return f"{b:.1f} TB"


def check_disk(path: Path, need_gb: float):
    import shutil
    free_gb = shutil.disk_usage(path).free / 1e9
    print(f"  Disk free: {free_gb:.1f} GB  (gerekli: ~{need_gb} GB)")
    if free_gb < need_gb * 1.1:
        print(f"  ! YETERSIZ DISK")
        return False
    return True


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--user", type=str, default=DEFAULT_USER)
    parser.add_argument("--no-data", action="store_true")
    parser.add_argument("--no-ckpt", action="store_true")
    parser.add_argument("--no-code", action="store_true")
    parser.add_argument("--token", type=str, default=None,
                        help="HF token (private repolar icin)")
    parser.add_argument("--only-best", action="store_true",
                        help="Sadece best_ckpt.pt indir (latest'i atla)")
    parser.add_argument("--workers", type=int, default=8,
                        help="Paralel indirme worker (default 8)")
    args = parser.parse_args()

    token = args.token or os.environ.get("HF_TOKEN")

    data_repo = f"{args.user}/{REPO_BASE}-data"
    ckpt_repo = f"{args.user}/{REPO_BASE}-ckpts"
    code_repo = f"{args.user}/{REPO_BASE}-code"

    DATA_DIR.mkdir(parents=True, exist_ok=True)
    RUN_DIR.mkdir(parents=True, exist_ok=True)

    # Disk kontrol
    print(f"\n{'='*60}\nDISK & ORTAM\n{'='*60}")
    needed = 0
    if not args.no_data: needed += 30
    if not args.no_ckpt: needed += 4 if args.only_best else 8
    if not args.no_code: needed += 0.01
    check_disk(ROOT, needed)

    # DATA
    if not args.no_data:
        print(f"\n{'='*60}\nDATA pull ← {data_repo}\n{'='*60}")
        print("  ~30GB indirme — internet hizina gore 15-40 dk")
        snapshot_download(
            repo_id=data_repo,
            repo_type="dataset",
            local_dir=str(DATA_DIR),
            token=token,
            max_workers=args.workers,
            allow_patterns=["*.bin", "*.json", "*.md"],
        )
        # Boyut dogrula
        for f in DATA_DIR.glob("v5_*.bin"):
            print(f"  ✓ {f.name}: {fmt_size(f.stat().st_size)}")

    # CKPT
    if not args.no_ckpt:
        print(f"\n{'='*60}\nCKPT pull ← {ckpt_repo}\n{'='*60}")
        patterns = ["best_ckpt.pt", "README.md"]
        if not args.only_best:
            patterns.extend(["latest_ckpt.pt", "train.log"])
        snapshot_download(
            repo_id=ckpt_repo,
            repo_type="model",
            local_dir=str(RUN_DIR),
            token=token,
            max_workers=args.workers,
            allow_patterns=patterns,
        )
        for f in RUN_DIR.glob("*.pt"):
            print(f"  ✓ {f.name}: {fmt_size(f.stat().st_size)}")

    # CODE
    if not args.no_code:
        print(f"\n{'='*60}\nCODE pull ← {code_repo}\n{'='*60}")
        snapshot_download(
            repo_id=code_repo,
            repo_type="model",
            local_dir=str(ROOT),
            token=token,
            max_workers=args.workers,
            allow_patterns=["*.py", "*.md"],
        )

    print(f"\n{'='*60}\n✓ TAMAMLANDI\n{'='*60}")
    print(f"\nSonraki adımlar (Thunder Compute):")
    print(f"  pip install -r requirements.txt  # veya manuel torch, tokenizers, liger-kernel")
    print(f"  python 05_train_v5.py --compile --resume")
    print(f"\nGPU kontrol:")
    print(f"  nvidia-smi")


if __name__ == "__main__":
    main()