#!/usr/bin/env bash
# Download ONNX models required by scene_analyze and build_gallery.
# Run from the project root: bash scripts/download_models.sh

set -euo pipefail

MODELS_DIR="${1:-models}"
mkdir -p "$MODELS_DIR"

# ── YuNet face detection ──────────────────────────────────────────────────────
YUNET_URL="https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx"
YUNET_FILE="$MODELS_DIR/face_detection_yunet_2023mar.onnx"
if [ ! -f "$YUNET_FILE" ]; then
    echo "Downloading YuNet…"
    curl -L "$YUNET_URL" -o "$YUNET_FILE"
else
    echo "YuNet already present: $YUNET_FILE"
fi

# ── ArcFace face recognition (buffalo_l / w600k_r50) ─────────────────────────
# This model is part of InsightFace's buffalo_l pack.
# We download and unpack only the recognition model.
ARCFACE_FILE="$MODELS_DIR/arcface_w600k_r50.onnx"
if [ ! -f "$ARCFACE_FILE" ]; then
    echo "Downloading ArcFace (buffalo_l)…"
    TMP_ZIP=$(mktemp /tmp/buffalo_l.XXXXXX.zip)
    curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_l.zip" \
        -o "$TMP_ZIP"
    # The zip contains: 1k3d68.onnx  2d106det.onnx  det_10g.onnx  genderage.onnx  w600k_r50.onnx
    unzip -jo "$TMP_ZIP" "w600k_r50.onnx" -d "$MODELS_DIR"
    mv "$MODELS_DIR/w600k_r50.onnx" "$ARCFACE_FILE"
    rm "$TMP_ZIP"
else
    echo "ArcFace already present: $ARCFACE_FILE"
fi

# ── ArcFace face recognition (buffalo_s / w600k_mbf — MobileFaceNet) ─────────
# Lighter backbone (13 MB vs 174 MB for R50) — same 512-dim output, faster inference.
ARCFACE_MBF_FILE="$MODELS_DIR/arcface_w600k_mbf.onnx"
if [ ! -f "$ARCFACE_MBF_FILE" ]; then
    echo "Downloading ArcFace MobileFaceNet (buffalo_s)…"
    TMP_ZIP=$(mktemp /tmp/buffalo_s.XXXXXX.zip)
    curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_s.zip" \
        -o "$TMP_ZIP"
    unzip -jo "$TMP_ZIP" "w600k_mbf.onnx" -d "$MODELS_DIR"
    mv "$MODELS_DIR/w600k_mbf.onnx" "$ARCFACE_MBF_FILE"
    rm "$TMP_ZIP"
else
    echo "ArcFace MBF already present: $ARCFACE_MBF_FILE"
fi

# ── SCRFD-500MF face detection (InsightFace buffalo_sc) ───────────────────────
# buffalo_sc.zip contains det_500m.onnx (SCRFD-500MF with 5 keypoints).
# If the unzip fails (file not found in archive), download manually from:
#   https://huggingface.co/deepinsight/insightface/resolve/main/models/buffalo_sc/det_500m.onnx
SCRFD_FILE="$MODELS_DIR/scrfd_500m_bnkps.onnx"
if [ ! -f "$SCRFD_FILE" ]; then
    echo "Downloading SCRFD-500MF (buffalo_sc)…"
    TMP_ZIP=$(mktemp /tmp/buffalo_sc.XXXXXX.zip)
    curl -L "https://github.com/deepinsight/insightface/releases/download/v0.7/buffalo_sc.zip" \
        -o "$TMP_ZIP"
    unzip -jo "$TMP_ZIP" "det_500m.onnx" -d "$MODELS_DIR"
    mv "$MODELS_DIR/det_500m.onnx" "$SCRFD_FILE"
    rm "$TMP_ZIP"
else
    echo "SCRFD-500MF already present: $SCRFD_FILE"
fi

# ── TransNetV2 shot-boundary detection (scene detector, opt-in) ───────────────
# ONNX export (elya5/transnetv2, MIT). Fixed input 1x100x27x48x3 (RGB 0-255),
# primary output "534" = per-frame boundary logits. Used only with --scene-detect.
SCENE_FILE="$MODELS_DIR/transnetv2.onnx"
SCENE_SHA="c4d54a682bace32f25136ef83ca2c9d403e8f8193775efeb995172a0d95a8e0c"
if [ ! -f "$SCENE_FILE" ]; then
    echo "Downloading TransNetV2…"
    curl -L "https://huggingface.co/elya5/transnetv2/resolve/main/transnetv2.onnx" \
        -o "$SCENE_FILE"
    if command -v sha256sum >/dev/null; then
        echo "$SCENE_SHA  $SCENE_FILE" | sha256sum -c - \
            || echo "WARNING: TransNetV2 sha256 mismatch (upstream may have changed)"
    fi
else
    echo "TransNetV2 already present: $SCENE_FILE"
fi

echo ""
echo "Models ready in $MODELS_DIR/:"
ls -lh "$MODELS_DIR"