#!/usr/bin/env bash
# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Download the pre-trained FP32 Worker Safety Gear Detection model
# from Intel Edge AI Resources.
# Usage: ./export_and_quantize.sh [OUTPUT_DIR]
# Example: ./export_and_quantize.sh ./models

set -euo pipefail

OUTPUT_DIR="${1:-./models}"
MODEL_DIR="${OUTPUT_DIR}/worker_safety_gear_detection"
ZIP_URL="https://github.com/open-edge-platform/edge-ai-resources/raw/main/models/FP32/worker-safety-gear-detection.zip"
ZIP_FILE="${OUTPUT_DIR}/worker-safety-gear-detection.zip"
EXPECTED_XML="${MODEL_DIR}/deployment/Detection/model/model.xml"

echo "--- Installing dependencies ---"
pip install -qU openvino

# Ask for approval before downloading models and sample files
echo ""
echo "This script will download:"
echo "  - Worker Safety Gear Detection model"
echo "  - Sample test video"
echo ""
read -p "Continue with downloads? (yes/no): " APPROVAL
if [[ "${APPROVAL}" != "yes" ]]; then
    echo "Download cancelled by user."
    exit 0
fi
echo ""

echo "--- Downloading sample test video ---"
if [[ ! -f test_video.avi ]]; then
    wget -q -O test_video.avi \
        https://github.com/open-edge-platform/edge-ai-resources/raw/refs/heads/main/videos/Safety_Full_Hat_and_Vest.avi
    echo "Downloaded: test_video.avi"
else
    echo "Already present: test_video.avi"
fi

echo "--- Downloading worker-safety-gear-detection FP32 model ---"
mkdir -p "${OUTPUT_DIR}"
if [[ ! -f "${ZIP_FILE}" ]]; then
    wget -q -O "${ZIP_FILE}" "${ZIP_URL}"
    echo "Downloaded: ${ZIP_FILE}"
else
    echo "Already downloaded: ${ZIP_FILE}"
fi

echo "--- Extracting model ---"
rm -rf "${MODEL_DIR}"
mkdir -p "${MODEL_DIR}"
unzip -qo "${ZIP_FILE}" -d "${MODEL_DIR}"

# The archive may contain an extra top-level worker_safety_gear_detection dir.
if [[ -d "${MODEL_DIR}/worker_safety_gear_detection/deployment" ]]; then
    mv "${MODEL_DIR}/worker_safety_gear_detection"/* "${MODEL_DIR}/"
    rmdir "${MODEL_DIR}/worker_safety_gear_detection" || true
fi

echo "Extracted to: ${MODEL_DIR}"

# Locate the model XML.
if [[ ! -f "${EXPECTED_XML}" ]]; then
    echo "ERROR: expected model.xml not found at ${EXPECTED_XML}" >&2
    exit 1
fi
MODEL_XML="${EXPECTED_XML}"
echo "Model IR found: ${MODEL_XML}"
echo "--- Done ---"