"""Self-heal model weights into /models. Idempotent — only missing files
are fetched. Called by the ml-worker entrypoint before Celery starts.
"""

import os
import sys
from pathlib import Path

MODEL_ROOT = Path(os.environ.get("ML_MODEL_DIR", "/models"))
CAMIE_REPO = os.environ.get("CAMIE_HF_REPO", "Camais03/camie-tagger-v2")
SIGLIP_REPO = os.environ.get(
    "SIGLIP_HF_REPO", "google/siglip-so400m-patch14-384"
)


def _snapshot(repo_id: str, dest: Path, allow_patterns: list[str] | None) -> None:
    from huggingface_hub import snapshot_download

    dest.mkdir(parents=True, exist_ok=True)
    snapshot_download(
        repo_id=repo_id,
        local_dir=str(dest),
        allow_patterns=allow_patterns,
    )


def ensure_camie() -> None:
    dest = MODEL_ROOT / "camie"
    if (dest / "model.onnx").is_file() and (dest / "selected_tags.csv").is_file():
        print(f"[download_models] Camie present at {dest}")
        return
    print(f"[download_models] Fetching {CAMIE_REPO} -> {dest}")
    _snapshot(CAMIE_REPO, dest, ["model.onnx", "selected_tags.csv", "*.json"])


def ensure_siglip() -> None:
    dest = MODEL_ROOT / "siglip"
    if (dest / "config.json").is_file() and any(dest.glob("*.safetensors")):
        print(f"[download_models] SigLIP present at {dest}")
        return
    print(f"[download_models] Fetching {SIGLIP_REPO} -> {dest}")
    _snapshot(SIGLIP_REPO, dest, None)


def main() -> int:
    ensure_camie()
    ensure_siglip()
    print("[download_models] Done.")
    return 0


if __name__ == "__main__":
    sys.exit(main())