"""
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.
# InsightFace (RetinaFace + ArcFace) auto-downloads via the insightface
# package on first use — no manual download entries needed.
DOWNLOADS = []

# 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 file(s); attempting automatic export:",
            len(missing),
        )
        for rel_path, desc in missing:
            logger.warning("  %s — %s", base / rel_path, desc)

        from app.services.vision import export_models

        missing_paths = {rel for rel, _ in missing}

        # Only run the export functions whose outputs are actually missing.
        # Each function is mapped to the file(s) it produces.
        export_map = [
            (export_models.export_openclip, ["embed/visual.onnx", "embed/textual.onnx"]),
            (export_models.export_siglip2, ["embed_siglip2/visual.onnx", "embed_siglip2/textual.onnx"]),
            (export_models.export_yolov8n, ["detect/yolov8n.onnx"]),
        ]
        for export_fn, outputs in export_map:
            if not any(o in missing_paths for o in outputs):
                continue
            try:
                export_fn(base)
            except Exception as e:
                logger.error(
                    "Export step %s failed: %s. "
                    "Run `python -m app.services.vision.export_models "
                    "--models-dir %s` manually to retry.",
                    export_fn.__name__,
                    e,
                    base,
                )

        # Re-check what's still missing after the export pass.
        still_missing = [
            (rel_path, desc)
            for rel_path, desc in EXPORTS
            if not (base / rel_path).exists()
        ]
        if still_missing:
            for rel_path, desc in still_missing:
                logger.error("  still missing: %s — %s", base / rel_path, desc)
        else:
            logger.info("All model files present in %s", base)
    else:
        logger.info("All model files present in %s", base)

    # Signal readiness via Redis so the scan pipeline knows the vision
    # worker can accept tasks.
    try:
        import redis as _redis
        r = _redis.from_url(settings.redis_url)
        r.set("mulita:vision:ready", "1")
        logger.info("Set mulita:vision:ready in Redis")
    except Exception as e:
        logger.warning("Could not set vision readiness flag in Redis: %s", e)


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