#!/usr/bin/env bash
# Put a GPU-capable ONNX Runtime where the desktop app looks for one
# (docs/dev/inference.md §3): `runtime/` beside the models in the user data
# directory, ahead of the system library.
#
#   ./tools/fetch-desktop-runtime.sh [DEST]
#
# The source is the `onnxruntime-gpu` wheel: the one build that carries the
# CUDA *and* TensorRT providers against the cuDNN and TensorRT majors current
# on this machine. Distribution packages tend to have neither — Arch's
# `onnxruntime-opt-cuda` is built without TensorRT and against cuDNN 8 — and
# the probe rejects them correctly and leaves the app on the CPU provider,
# which is what this script exists to fix. Nothing NVIDIA is bundled here:
# the providers load CUDA, cuDNN and TensorRT from the system, and if those
# are missing the probe says so and the app stays on the CPU.
#
# This is the NVIDIA script. On AMD there is nothing to fetch: the
# distribution's ROCm build of ONNX Runtime (Arch's `onnxruntime-rocm`)
# carries the MIGraphX provider, and the app finds it in the system library
# directory (docs/inference.md §1.3).
set -euo pipefail
DEST="${1:-${XDG_DATA_HOME:-${HOME}/.local/share}/darkroom/runtime}"
WORK="$(mktemp -d -p /var/tmp fetch-desktop-runtime.XXXXXX)"
trap 'rm -rf "${WORK}"' EXIT

echo "==> downloading the onnxruntime-gpu wheel"
uv venv --python 3.12 "${WORK}/venv" >/dev/null
VIRTUAL_ENV="${WORK}/venv" uv pip install --quiet onnxruntime-gpu
CAPI="$(find "${WORK}/venv" -type d -path '*/onnxruntime/capi' | head -1)"
[[ -n "${CAPI}" ]] || { echo "error: no capi directory in the wheel" >&2; exit 1; }

mkdir -p "${DEST}"
# The runtime and its provider libraries; not the Python binding.
cp "${CAPI}"/libonnxruntime.so* "${CAPI}"/libonnxruntime_providers_*.so "${DEST}/"
echo "==> runtime in ${DEST}:"
ls -1 "${DEST}" | sed 's/^/    /'
echo "    (the app finds it on its next launch; Settings › About › Inference says what it chose)"