"""Download every model the service uses into the image, at build time.

Run by the Dockerfile. Without it, each model downloaded the first time it was used after the Space
started, and with a 15-minute sleep the Space starts often: the first film after a wake waited
minutes for gigabytes of weights. HOME and HF_HOME point into the image (see the Dockerfile), so
the caches written here are the ones the service reads at run time.

Downloads only, wherever a download can be had without building the model: the build machine has
far less memory than the GPU it runs on, and loading emotion2vec here got the build OOM-killed.
Each model is fetched in its own process, so whatever one leaves in memory is gone before the next.
A failure is reported, not fatal: a model missing from the image still downloads on first use.
"""
import os
import subprocess
import sys

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # the build has no GPU


def whisper(size):
    from faster_whisper.utils import download_model
    download_model(size)


def hub(repo, **kw):
    from huggingface_hub import snapshot_download
    snapshot_download(repo, **kw)


def hsemotion():
    from hsemotion_onnx.facial_emotions import get_model_path
    get_model_path("enet_b0_8_best_afew")


def easyocr_models():
    # Builds the reader (small: the text detector and the English recogniser), which is how
    # EasyOCR downloads its weights.
    import easyocr
    easyocr.Reader(["en"], gpu=False, verbose=False)


def memorability_models():
    from huggingface_hub import hf_hub_download
    hf_hub_download("PerceptCLIP/PerceptCLIP_Memorability", "modeling.py")
    hf_hub_download("PerceptCLIP/PerceptCLIP_Memorability", "perceptCLIP_Memorability.pth")
    hub("openai/clip-vit-large-patch14", allow_patterns=["*.json", "*.txt", "*.safetensors"])
    hub("sentence-transformers/all-MiniLM-L6-v2")


def deepface_emotion():
    # DeepFace's emotion model is small; building it is how DeepFace downloads it.
    import numpy as np
    from deepface import DeepFace
    DeepFace.analyze(np.zeros((64, 64, 3), dtype=np.uint8), actions=["emotion"], enforce_detection=False, silent=True)


JOBS = {
    "whisper medium": lambda: whisper("medium"),
    "whisper small": lambda: whisper("small"),
    "whisper medium.en (Stage Read)": lambda: whisper("medium.en"),
    "whisper small.en (Stage Read)": lambda: whisper("small.en"),
    # funasr with hub="hf" reads this from the Hugging Face cache
    "emotion2vec": lambda: hub("emotion2vec/emotion2vec_plus_large"),
    "HSEmotion": hsemotion,
    "EasyOCR": easyocr_models,
    "GoEmotions": lambda: hub("SamLowe/roberta-base-go_emotions"),
    "memorability (CLIP, PerceptCLIP, MiniLM)": memorability_models,
    "DeepFace emotion": deepface_emotion,
    "wav2vec2 (Stage Read)": lambda: hub("audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim"),
}

if __name__ == "__main__":
    if len(sys.argv) > 1:
        JOBS[sys.argv[1]]()
        sys.exit(0)
    failed = []
    for name in JOBS:
        print(f"[prefetch] {name}...", flush=True)
        if subprocess.run([sys.executable, __file__, name]).returncode == 0:
            print(f"[prefetch] {name} done", flush=True)
        else:
            print(f"[prefetch] {name} FAILED (it will download on first use instead)", flush=True)
            failed.append(name)
    print(f"[prefetch] finished; {len(failed)} failed: {failed}", flush=True)