#!/bin/sh
# This helper deliberately installs only Viska's optional CUDA Python extra.
# NVIDIA drivers remain the administrator's responsibility.
set -eu

venv_python=/opt/viska/venv/bin/python
wheel=/opt/viska/viska-0.5.60-py3-none-any.whl

set_cuda_paths() {
  cuda_libs=
  for cuda_lib in /opt/viska/venv/lib/python*/site-packages/nvidia/cublas/lib \
                  /opt/viska/venv/lib/python*/site-packages/nvidia/cudnn/lib; do
    if [ -d "$cuda_lib" ]; then
      cuda_libs="${cuda_libs:+$cuda_libs:}$cuda_lib"
    fi
  done
  if [ -n "$cuda_libs" ]; then
    export LD_LIBRARY_PATH="$cuda_libs${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
  fi
}

usage() {
  echo "usage: viska-gpu-setup status|install" >&2
  exit 2
}

[ "$#" -eq 1 ] || usage
case "$1" in
  status)
    [ -x "$venv_python" ] || { echo "Viska is not configured yet." >&2; exit 1; }
    set_cuda_paths
    exec "$venv_python" -m viska.core.cuda_runtime --json
    ;;
  install)
    [ -x "$venv_python" ] || { echo "Viska is not configured yet." >&2; exit 1; }
    if [ "$(id -u)" -ne 0 ]; then
      echo "Run 'sudo viska-gpu-setup install' to modify /opt/viska/venv." >&2
      exit 1
    fi
    echo "Viska: installing its optional CUDA runtime."
    echo "Viska: NVIDIA drivers are not installed by this helper."
    "$venv_python" -m pip install --upgrade "$wheel[cuda]"
    set_cuda_paths
    exec "$venv_python" -m viska.core.cuda_runtime --json --require-cuda
    ;;
  *) usage ;;
esac
