#!/usr/bin/env bash
# SPDX-License-Identifier: MIT
# Copyright (C) Intel Corporation
#
# Download the YOLOv8 license plate detector and PaddleOCR PP-OCRv4
# recognizer (both as OpenVINO IR) for use with Intel DLStreamer.
# Usage: ./export_and_quantize.sh [MODELS_DIR]
# Example: ./export_and_quantize.sh ./models

set -euo pipefail

MODELS_DIR="${1:-./models}"
LP_DETECTOR_NAME="yolov8_license_plate_detector"
LP_DETECTOR_URL="https://github.com/open-edge-platform/edge-ai-resources/raw/main/models/license-plate-reader.zip"

mkdir -p "${MODELS_DIR}"

echo "--- Downloading sample test video ---"
SAMPLE_VIDEO_URL="https://github.com/open-edge-platform/edge-ai-resources/raw/main/videos/ParkingVideo.mp4"
if [[ ! -f ParkingVideo.mp4 ]]; then
    wget -q -O ParkingVideo.mp4 "${SAMPLE_VIDEO_URL}"
    echo "Downloaded: ParkingVideo.mp4"
else
    echo "Already present: ParkingVideo.mp4"
fi

echo "--- Downloading ${LP_DETECTOR_NAME} (OpenVINO IR) ---"
LP_DIR="${MODELS_DIR}/${LP_DETECTOR_NAME}"
mkdir -p "${LP_DIR}"
TMP_ZIP="$(mktemp --suffix=.zip)"
trap 'rm -f "${TMP_ZIP}"' EXIT
wget -q -O "${TMP_ZIP}" "${LP_DETECTOR_URL}"
unzip -oq "${TMP_ZIP}" -d "${LP_DIR}"

LP_XML="$(find "${LP_DIR}" -name "*retrained*.xml" -o -name "*license*plate*.xml" | head -n1)"
if [[ -z "${LP_XML}" ]]; then
    LP_XML="$(find "${LP_DIR}" -path "*/yolov8n/*.xml" | head -n1)"
fi
if [[ -z "${LP_XML}" ]]; then
    echo "Error: license plate detector .xml not found under ${LP_DIR}" >&2
    exit 1
fi
echo "Found detector model: ${LP_XML}"

OCR_XML="$(find "${LP_DIR}" -path "*ch_PP-OCRv4_rec_infer*.xml" | head -n1)"
if [[ -z "${OCR_XML}" ]]; then
    echo "Error: PaddleOCR PP-OCRv4 .xml not found under ${LP_DIR}" >&2
    exit 1
fi
echo "Found OCR model: ${OCR_XML}"

echo "--- Done ---"
echo "Detector : ${LP_XML}"
echo "OCR      : ${OCR_XML}"
echo "Sample   : $(pwd)/ParkingVideo.mp4"