"""
Download vision model weights on first worker boot.

Run as:  python -m app.services.vision.bootstrap_models

Or called from the vision worker entrypoint before Celery starts.
Downloads are idempotent — existing files with matching sizes are skipped.
"""
import logging
import os
from pathlib import Path
from urllib.request import urlretrieve

from app.config import settings

logger = logging.getLogger(__name__)

# (relative_path, url, expected_size_bytes_approx)
# Sizes are approximate — used only for skip-if-exists checks, not integrity.
MODELS = [
    # OpenCLIP ViT-B/32 — visual and textual encoders (ONNX)
    # These must be exported manually via export_openclip.py (see below).
    # Placeholder entries — bootstrap will warn if missing.
    ("embed/visual.onnx", None, None),
    ("embed/textual.onnx", None, None),

    # YOLOv8n — object detection
    # Export: `yolo export model=yolov8n.pt format=onnx imgsz=640`
    # Placeholder — must be exported from ultralytics offline.
    ("detect/yolov8n.onnx", None, None),

    # YuNet — face detection (Apache 2.0, opencv_zoo)
    (
        "face/yunet.onnx",
        "https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
        233_000,
    ),

    # SFace — face recognition (Apache 2.0, opencv_zoo)
    (
        "face/sface.onnx",
        "https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx",
        37_000_000,
    ),
]


def bootstrap(models_dir: str | None = None):
    """Ensure all model files are present. Download what we can, warn about
    files that need manual export."""
    base = Path(models_dir or settings.vision.models_dir)
    base.mkdir(parents=True, exist_ok=True)

    for rel_path, url, expected_size in MODELS:
        dest = base / rel_path
        dest.parent.mkdir(parents=True, exist_ok=True)

        if dest.exists():
            logger.debug("Model already exists: %s", dest)
            continue

        if url is None:
            logger.warning(
                "Model file %s not found and has no auto-download URL. "
                "See bootstrap_models.py for export instructions.",
                dest,
            )
            continue

        logger.info("Downloading %s → %s", url, dest)
        try:
            urlretrieve(url, str(dest))
            actual = dest.stat().st_size
            logger.info("Downloaded %s (%d bytes)", rel_path, actual)
        except Exception as e:
            logger.error("Failed to download %s: %s", rel_path, e)
            if dest.exists():
                dest.unlink()


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    bootstrap()