"""Ensure WD14 and SigLIP model files exist under $ML_MODEL_DIR.

Run at container start by entrypoint-ml.sh. huggingface_hub performs content-based
integrity checks: matching files are skipped, changed/missing files are fetched.
If the remote is unreachable but a local copy is already present, log a warning
and continue so a transient HF outage does not take a healthy worker down.
"""
from __future__ import annotations

import logging
import os
import sys

from huggingface_hub import hf_hub_download, snapshot_download
from huggingface_hub.utils import HfHubHTTPError

logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(name)s: %(message)s')
log = logging.getLogger('ensure_models')

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')
WD14_REVISION = os.environ.get('WD14_REVISION') or None
SIGLIP_REVISION = os.environ.get('SIGLIP_REVISION') or None

WD14_FILES = ('model.onnx', 'selected_tags.csv')
SIGLIP_CORE_FILES = (
    'config.json',
    'preprocessor_config.json',
    'tokenizer_config.json',
    'tokenizer.json',
    'special_tokens_map.json',
    'spiece.model',
    'model.safetensors',
)


def _all_present(dir_path: str, files: tuple[str, ...]) -> bool:
    return all(os.path.exists(os.path.join(dir_path, f)) for f in files)


def ensure_wd14() -> None:
    target = os.path.join(MODEL_DIR, 'wd14')
    os.makedirs(target, exist_ok=True)
    try:
        for fname in WD14_FILES:
            hf_hub_download(
                repo_id=WD14_REPO,
                filename=fname,
                revision=WD14_REVISION,
                local_dir=target,
            )
        log.info('WD14 ready at %s', target)
    except (HfHubHTTPError, OSError) as exc:
        if _all_present(target, WD14_FILES):
            log.warning('WD14 refresh failed (%s); continuing with local copy', exc)
        else:
            log.error('WD14 missing and cannot download: %s', exc)
            raise


def ensure_siglip() -> None:
    target = os.path.join(MODEL_DIR, 'siglip')
    os.makedirs(target, exist_ok=True)
    try:
        snapshot_download(
            repo_id=SIGLIP_REPO,
            revision=SIGLIP_REVISION,
            local_dir=target,
            allow_patterns=list(SIGLIP_CORE_FILES),
        )
        log.info('SigLIP ready at %s', target)
    except (HfHubHTTPError, OSError) as exc:
        if _all_present(target, SIGLIP_CORE_FILES):
            log.warning('SigLIP refresh failed (%s); continuing with local copy', exc)
        else:
            log.error('SigLIP missing and cannot download: %s', exc)
            raise


if __name__ == '__main__':
    ensure_wd14()
    ensure_siglip()
    sys.exit(0)