#!/usr/bin/env python3
"""Download everything onto the node's persistent disk. Safe to re-run (skips finished files).
Needs HF_TOKEN in the environment for the gated N-ATLaS repo."""
import os
import sys
import time

from huggingface_hub import hf_hub_download, snapshot_download

HOME = os.path.expanduser("~")
DATA_DIR = f"{HOME}/data/llava_pretrain"
LLM_NAME = os.environ.get("LLM_NAME", "NCAIR1/N-ATLaS")
VISION_NAME = os.environ.get("VISION_NAME", "google/siglip2-base-patch16-224")


def log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


if not os.environ.get("HF_TOKEN"):
    sys.exit("HF_TOKEN is not set - needed for the gated N-ATLaS repo")

# 1. Fail fast if the token can't see the gated model
try:
    hf_hub_download(LLM_NAME, "config.json")
    log(f"gated access OK for {LLM_NAME}")
except Exception as e:
    sys.exit(f"Cannot access {LLM_NAME}: {type(e).__name__}: {e}\n"
             "Check that access was granted on the model page and the token can read gated repos.")

t = time.time()
snapshot_download(LLM_NAME, allow_patterns=["*.json", "*.safetensors", "tokenizer*", "*.txt", "*.model", "*.jinja"])
log(f"LLM downloaded ({time.time() - t:.0f}s)")

t = time.time()
snapshot_download(VISION_NAME)
log(f"vision encoder downloaded ({time.time() - t:.0f}s)")

t = time.time()
os.makedirs(DATA_DIR, exist_ok=True)
hf_hub_download("liuhaotian/LLaVA-Pretrain", "blip_laion_cc_sbu_558k.json", repo_type="dataset", local_dir=DATA_DIR)
hf_hub_download("liuhaotian/LLaVA-Pretrain", "images.zip", repo_type="dataset", local_dir=DATA_DIR)
size_gb = os.path.getsize(f"{DATA_DIR}/images.zip") / 1e9
log(f"dataset downloaded: images.zip {size_gb:.1f} GB ({time.time() - t:.0f}s)")