"""
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 are skipped.

For models that require export (OpenCLIP, YOLOv8n), see export_models.py.
Those must be exported once on any machine with pip, then placed in
the models volume before the worker starts.
"""
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, description)
# Models with url=None must be pre-exported via export_models.py.
DOWNLOADS = [
    # YuNet — face detection (Apache 2.0, opencv_zoo, ~233 KB)
    (
        "face/yunet.onnx",
        "https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx",
        "YuNet face detector",
    ),
    # SFace — face recognition (Apache 2.0, opencv_zoo, ~37 MB)
    (
        "face/sface.onnx",
        "https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx",
        "SFace face recognizer",
    ),
]

# Models that need manual export via export_models.py
EXPORTS = [
    ("embed/visual.onnx", "OpenCLIP ViT-B/32 visual encoder"),
    ("embed/textual.onnx", "OpenCLIP ViT-B/32 textual encoder"),
    ("detect/yolov8n.onnx", "YOLOv8n object detector"),
]


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)

    # Download auto-downloadable models
    for rel_path, url, desc in DOWNLOADS:
        dest = base / rel_path
        dest.parent.mkdir(parents=True, exist_ok=True)

        if dest.exists():
            logger.debug("Already exists: %s (%s)", dest, desc)
            continue

        logger.info("Downloading %s → %s", desc, dest)
        try:
            urlretrieve(url, str(dest))
            size_kb = dest.stat().st_size / 1024
            logger.info("Downloaded %s (%.0f KB)", desc, size_kb)
        except Exception as e:
            logger.error("Failed to download %s: %s", desc, e)
            if dest.exists():
                dest.unlink()

    # Check for manually-exported models
    missing = []
    for rel_path, desc in EXPORTS:
        dest = base / rel_path
        if not dest.exists():
            missing.append((rel_path, desc))

    if missing:
        logger.warning(
            "Missing %d model(s) that require manual export via export_models.py:",
            len(missing),
        )
        for rel_path, desc in missing:
            logger.warning("  %s — %s", base / rel_path, desc)
        logger.warning(
            "Run: python -m app.services.vision.export_models --models-dir %s",
            base,
        )
    else:
        logger.info("All model files present in %s", base)


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