#!/bin/bash
# Download MobileNet V2 SSD (COCO, INT8 quantized) for TFLite
#
# Model: ssd_mobilenet_v2_coco_quant_postprocess.tflite (~6MB)
# Source: TensorFlow Model Zoo
# Classes: 80 COCO object classes
#
# This script is idempotent - skips download if model exists.

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODEL_DIR="${SCRIPT_DIR}/models"
MODEL_FILE="${MODEL_DIR}/ssd_mobilenet_v2_coco_quant_postprocess.tflite"
ZIP_URL="https://storage.googleapis.com/download.tensorflow.org/models/tflite/coco_ssd_mobilenet_v2_quantized_300x300_uint8_20200430.zip"
ZIP_FILE="${MODEL_DIR}/model.zip"

mkdir -p "${MODEL_DIR}"

if [ -f "${MODEL_FILE}" ]; then
    echo "Model already exists: ${MODEL_FILE}"
    echo "Size: $(du -h "${MODEL_FILE}" | cut -f1)"
    exit 0
fi

echo "Downloading MobileNet V2 SSD (quantized)..."
curl -L -o "${ZIP_FILE}" "${ZIP_URL}"

echo "Extracting model..."
unzip -o "${ZIP_FILE}" -d "${MODEL_DIR}"

# Clean up zip and extra files
rm -f "${ZIP_FILE}"
rm -f "${MODEL_DIR}/labelmap.txt"  # We use our own coco_labels.txt

if [ -f "${MODEL_FILE}" ]; then
    echo "Model downloaded: ${MODEL_FILE}"
    echo "Size: $(du -h "${MODEL_FILE}" | cut -f1)"
else
    # The zip might extract with a slightly different path
    FOUND=$(find "${MODEL_DIR}" -name "*.tflite" -type f | head -1)
    if [ -n "${FOUND}" ] && [ "${FOUND}" != "${MODEL_FILE}" ]; then
        mv "${FOUND}" "${MODEL_FILE}"
        echo "Model downloaded: ${MODEL_FILE}"
        echo "Size: $(du -h "${MODEL_FILE}" | cut -f1)"
    else
        echo "ERROR: Model file not found after extraction"
        exit 1
    fi
fi

# Clean up any extracted subdirectories
find "${MODEL_DIR}" -mindepth 1 -type d -exec rm -rf {} + 2>/dev/null || true

echo "Done!"