import os
import shutil
from huggingface_hub import hf_hub_download

BLUE_REPO = "notmax123/blue-onnx"
RENIKUD_REPO = "thewh1teagle/renikud"

BLUE_FILES = [
    "text_encoder.onnx",
    "reference_encoder.onnx",
    "vector_estimator.onnx",
    "vocoder.onnx",
    "duration_predictor.onnx",
    "length_pred_style.onnx",
    "stats.npz",
    "uncond.npz",
]

def _is_valid(path, min_bytes=100_000):
    return os.path.exists(path) and os.path.getsize(path) >= min_bytes

def download_blue_models(dest_dir="onnx_models"):
    os.makedirs(dest_dir, exist_ok=True)
    token = os.environ.get("HF_TOKEN") or None

    for filename in BLUE_FILES:
        dest = os.path.join(dest_dir, filename)
        if _is_valid(dest):
            print(f"Already present: {dest} ({os.path.getsize(dest):,} bytes)")
            continue
        print(f"Downloading {BLUE_REPO}/{filename} …")
        try:
            cached = hf_hub_download(
                repo_id=BLUE_REPO,
                filename=filename,
                repo_type="model",
                token=token,
            )
            shutil.copy2(cached, dest)
            print(f"  → {dest} ({os.path.getsize(dest):,} bytes)")
        except Exception as e:
            print(f"  FAILED {filename}: {e}")
            raise

def download_renikud(dest="renikud.onnx"):
    if _is_valid(dest, min_bytes=1_000_000):
        print(f"Already present: {dest} ({os.path.getsize(dest):,} bytes)")
        return
    token = os.environ.get("HF_TOKEN") or None
    print(f"Downloading renikud model …")
    cached = hf_hub_download(
        repo_id=RENIKUD_REPO,
        filename="model.onnx",
        repo_type="model",
        token=token,
    )
    shutil.copy2(cached, dest)
    print(f"  → {dest} ({os.path.getsize(dest):,} bytes)")

if __name__ == "__main__":
    download_blue_models()
    download_renikud()