import os
import subprocess
from huggingface_hub import snapshot_download, login

def download_models():
    token = os.environ.get("HF_TOKEN", "HF_TOKEN").strip()

    if not token:
        raise RuntimeError("HF_TOKEN is not set.")

    login(token=token)
    print("Authenticated with HuggingFace.")

    models = [
        "drrobot9/wav2vec2-nigerian-language-identifier-v8",
        "NCAIR1/Yoruba-ASR",
        "NCAIR1/Igbo-ASR",
        "NCAIR1/Hausa-ASR",
        "NCAIR1/NigerianAccentedEnglish",
        "NCAIR1/N-ATLaS",
    ]

    ignore = [
        "*.msgpack", "*.h5", "flax_model*", "tf_model*",
        "rust_model*", "optimizer.pt", "rng_state*",
        "training_args.bin", "trainer_state.json", "scheduler.pt",
    ]

    for model_id in models:
        print(f"Downloading {model_id}...")
        snapshot_download(
            repo_id=model_id,
            token=token,
            ignore_patterns=ignore,
        )
        print(f"  Done: {model_id}")

    print("All models downloaded.")


if __name__ == "__main__":
    download_models()

    # Start the FastAPI server
    subprocess.run(
        ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"],
        check=True,
    )