#!/usr/bin/env bash
# Install build dependencies for dora-zed on NVIDIA Jetson (ARM64 / L4T).
# Skips any component that is already installed.
# Uses aria2c for fast parallel downloads if available, otherwise falls back to wget.
#
# Supported JetPack versions for ZED SDK 5.x:
#   JetPack 6.0 GA  (L4T 36.3) — Jetson Orin
#   JetPack 6.1/6.2 (L4T 36.4)
#   JetPack 6.2.2   (L4T 36.5)
#   JetPack 7.0     (L4T 38.2) — Jetson Thor
#   JetPack 7.1     (L4T 38.4)
#
# ONNX Runtime: Microsoft does not publish official GPU binaries for ARM64.
#   The native engine (ZED SDK built-in TensorRT) works without ONNX Runtime.
#   To use the yolo_onnx engine, build ONNX Runtime from source:
#   https://onnxruntime.ai/docs/build/eps.html#tensorrt
set -euo pipefail

ZED_VERSION="5.2"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TMP_DIR="${SCRIPT_DIR}/../tmp"
mkdir -p "$TMP_DIR"

download() {
    local url="$1" dest="$2"
    if command -v aria2c &>/dev/null; then
        aria2c -x 16 -s 16 -d "$(dirname "$dest")" -o "$(basename "$dest")" "$url"
    else
        wget -O "$dest" "$url"
    fi
}

# --- Detect L4T version ---
if [ ! -f /etc/nv_tegra_release ]; then
    echo "Error: /etc/nv_tegra_release not found. Is this a Jetson device?" >&2
    exit 1
fi
L4T_VERSION=$(head -1 /etc/nv_tegra_release | grep -oP 'R\K[0-9]+\.[0-9]+' | head -1)
L4T_MAJOR=$(echo "$L4T_VERSION" | cut -d'.' -f1)
echo ">>> Detected L4T ${L4T_VERSION} (major: ${L4T_MAJOR})"

case "$L4T_MAJOR" in
    36) L4T_PATH="l4t36" ;;
    38) L4T_PATH="l4t38" ;;
    35) L4T_PATH="l4t35" ;;
    *)
        echo "Warning: L4T ${L4T_MAJOR} not tested. Attempting l4t${L4T_MAJOR}." >&2
        L4T_PATH="l4t${L4T_MAJOR}"
        ;;
esac

# --- ZED SDK ---
if [ -d /usr/local/zed ]; then
    echo ">>> ZED SDK already installed at /usr/local/zed — skipping."
else
    echo ">>> Downloading ZED SDK ${ZED_VERSION} for Jetson (${L4T_PATH})..."
    download \
        "https://download.stereolabs.com/zedsdk/${ZED_VERSION}/${L4T_PATH}/jetsons" \
        "${TMP_DIR}/ZED_SDK.run"
    chmod +x "${TMP_DIR}/ZED_SDK.run"
    echo ">>> Installing ZED SDK (silent)..."
    sudo "${TMP_DIR}/ZED_SDK.run" -- silent skip_python skip_tools
fi

# --- ONNX Runtime ---
echo ""
echo ">>> ONNX Runtime (ARM64 GPU): no official pre-built binary available."
echo "    The native ZED engine works without ONNX Runtime."
echo "    To use the yolo_onnx engine, build from source:"
echo "      https://onnxruntime.ai/docs/build/eps.html#tensorrt"
echo ""

echo ">>> Done. You can now run: mkdir -p build && cd build && cmake .. && make install"
echo ">>> Downloaded files are in ${TMP_DIR}"