#!/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

if [[ "$(uname -s)" == "Linux" ]]; then
  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}"
fi

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

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