"""Pre-fetch WD14 and SigLIP model files into $ML_MODEL_DIR at image build time."""
from __future__ import annotations
import os
import sys

from huggingface_hub import hf_hub_download, snapshot_download

MODEL_DIR = os.environ.get('ML_MODEL_DIR', '/models')
WD14_REPO = os.environ.get('WD14_REPO', 'SmilingWolf/wd-eva02-large-tagger-v3')
SIGLIP_REPO = os.environ.get('SIGLIP_REPO', 'google/siglip-so400m-patch14-384')


def download_wd14() -> None:
    target = os.path.join(MODEL_DIR, 'wd14')
    os.makedirs(target, exist_ok=True)
    hf_hub_download(repo_id=WD14_REPO, filename='model.onnx', local_dir=target, local_dir_use_symlinks=False)
    hf_hub_download(repo_id=WD14_REPO, filename='selected_tags.csv', local_dir=target, local_dir_use_symlinks=False)
    print(f"WD14 downloaded to {target}")


def download_siglip() -> None:
    target = os.path.join(MODEL_DIR, 'siglip')
    os.makedirs(target, exist_ok=True)
    snapshot_download(
        repo_id=SIGLIP_REPO,
        local_dir=target,
        local_dir_use_symlinks=False,
        allow_patterns=[
            'config.json',
            'preprocessor_config.json',
            'tokenizer_config.json',
            'tokenizer.json',
            'special_tokens_map.json',
            'spiece.model',
            'model.safetensors',
            'pytorch_model.bin',
        ],
    )
    print(f"SigLIP downloaded to {target}")


if __name__ == '__main__':
    download_wd14()
    download_siglip()
    sys.exit(0)