# prepare_assets.py
# Download ONLY: GGUF (llama.cpp) + prefetch guardrail XNLI
import os, shutil
from pathlib import Path
from huggingface_hub import hf_hub_download, snapshot_download

BASE = Path(__file__).resolve().parent
os.environ.setdefault("HF_HOME", str(BASE / ".hf-cache"))  # cache lokal biar cepat restart

GGUF_REPO_ID  = os.getenv("GGUF_REPO_ID",  "unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF")
GGUF_FILENAME = os.getenv("GGUF_FILENAME", "DeepSeek-R1-0528-Qwen3-8B-Q4_K_M.gguf")
XNLI_REPO_ID  = os.getenv("XNLI_REPO_ID",  "joeddav/xlm-roberta-large-xnli")

def ensure_dir(p: Path):
    p.parent.mkdir(parents=True, exist_ok=True)

def main():
    print("=== PREPARE_ASSETS start ===")

    # 1) Download GGUF -> models/
    try:
        target = BASE / "models" / GGUF_FILENAME
        ensure_dir(target)
        local = hf_hub_download(repo_id=GGUF_REPO_ID, filename=GGUF_FILENAME, repo_type="model")
        shutil.copy(local, target)
        print(f"[OK] GGUF -> {target}")
    except Exception as e:
        print(f"[WARN] GGUF download gagal: {e}")

    # 2) Prefetch guardrail model ke cache (biar pipeline cepat)
    try:
        snapshot_download(repo_id=XNLI_REPO_ID)  # hanya ke cache
        print(f"[OK] Prefetch guardrail: {XNLI_REPO_ID}")
    except Exception as e:
        print(f"[WARN] Prefetch XNLI gagal: {e}")

    print("=== PREPARE_ASSETS done ===")

if __name__ == "__main__":
    main()