"""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:
    """Fetch Camie v2 weights + metadata.

    v2 layout (HuggingFace Camais03/camie-tagger-v2): the ONNX file is
    named camie-tagger-v2.onnx (not model.onnx) and tags ship inside
    camie-tagger-v2-metadata.json (not selected_tags.csv). Both at root.
    The repo also contains app/, game/, training/, images/ subdirs full
    of setup/demo files we don't need — allow_patterns scopes the fetch
    to just the inference essentials (~790 MB instead of ~2 GB).
    """
    dest = MODEL_ROOT / "camie"
    model_file = dest / "camie-tagger-v2.onnx"
    meta_file = dest / "camie-tagger-v2-metadata.json"
    if model_file.is_file() and meta_file.is_file():
        print(f"[download_models] Camie present at {dest}")
        return
    print(f"[download_models] Fetching {CAMIE_REPO} -> {dest}")
    _snapshot(
        CAMIE_REPO, dest,
        [
            "camie-tagger-v2.onnx",
            "camie-tagger-v2-metadata.json",
            "config.json",
            "config.yaml",
        ],
    )


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())