import os
import requests

MODEL_DIR = "zipvoice_finetune"
os.makedirs(MODEL_DIR, exist_ok=True)

files = {
    "iter-525000-avg-2.pt": "https://huggingface.co/datasets/kjanh/demo_zip/resolve/main/epoch-46-all-speak-600h-en-norm.pt",
    "model.json": "https://huggingface.co/datasets/kjanh/demo_zip/resolve/main/model.json",
    "tokens.txt": "https://huggingface.co/datasets/kjanh/demo_zip/resolve/main/tokens.txt",
}

HF_TOKEN = os.getenv("HF_TOKEN")

def download_with_token(url, dest_path):
    if os.path.exists(dest_path):
        print(f"✔ File tồn tại: {dest_path}")
        return

    if HF_TOKEN is None:
        raise RuntimeError("❌ Missing HF_TOKEN in Secrets!")

    print(f"⬇ Downloading {dest_path} ...")

    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    r = requests.get(url, headers=headers, stream=True)
    r.raise_for_status()

    with open(dest_path, "wb") as f:
        for chunk in r.iter_content(1024 * 1024):
            f.write(chunk)

    print(f"✅ Downloaded {dest_path}")
# demo
def download_all_models():
    for filename, url in files.items():
        dest = os.path.join(MODEL_DIR, filename)
        download_with_token(url, dest)