#!/usr/bin/env bash
set -euo pipefail

binary="$1"
command="$2"
shift 2

if [[ -z "${RUNFILES_DIR:-}" && -d "$0.runfiles" ]]; then
  RUNFILES_DIR="$0.runfiles"
fi
if [[ -z "${RUNFILES_DIR:-}" ]]; then
  echo "Bazel runfiles are unavailable." >&2
  exit 1
fi

find_library_dir() {
  local name="$1"
  local path
  path="$(find -L "$RUNFILES_DIR" -name "$name" -print -quit)"
  if [[ -z "$path" ]]; then
    echo "Unable to find $name in Bazel runfiles." >&2
    exit 1
  fi
  dirname "$path"
}

cublas_dir="$(find_library_dir libcublas.so.12)"
cudnn_dir="$(find_library_dir libcudnn.so.9)"
export LD_LIBRARY_PATH="$cublas_dir:$cudnn_dir${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"

if [[ "$binary" != /* ]]; then
  binary="$RUNFILES_DIR/_main/$binary"
fi

exec "$binary" "$command" "$@"