#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091
# Default project path
PROJECT_PATH="$HOME/projects/repo"

# Parse command-line options
if ! OPTIONS=$(getopt -o p: --long project_path: -n 'parse-options' -- "$@"); then
	echo "install_env.sh: Error parsing options" >&2
	exit 1
fi

eval set -- "$OPTIONS"

echo "install_env.sh: options parsed OPTIONS=$OPTIONS"

while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	--)
		shift
		break
		;;
	*)
		break
		;;
	esac
done
printf "install_env.sh: arguments=%s, first argument=%s\n" "$@" "$1"
echo "install_env.sh: Install env in PROJECT_PATH=$PROJECT_PATH"
#! Moving to the project folder
cd "$PROJECT_PATH" || exit
#! Activate UV environment

# This handles locking and synchronization of the environment
# They are both supposed to be very fast
# GIT_LFS_SKIP_SMUDGE is necessary to avoid downloading large files
GIT_LFS_SKIP_SMUDGE=true uv sync

# Install the pre-commit git hooks
# Supposed to be very fast
uv run pre-commit install

# Adding CUDA paths to environment variables
export PATH=/usr/local/cuda-12.4/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

if [[ $1 == "no_cuda" ]]; then
	echo "install_env.sh: skipping any GPU-related setting."
else
	#! Check the output of `nvcc -V`
	NVCC_OUTPUT=$(nvcc -V)
	if [[ $NVCC_OUTPUT == *"release 12.4"* ]]; then
		echo "install_env.sh: CUDA 12.4 is detected."
	else
		echo "install_env.sh: CUDA 12.4 not detected. Please install CUDA 12.4. Exiting..."
		exit 1
	fi
	#! The following packages may be necessary but require manual installation, see:
	#! Flash-attn bug: https://github.com/python-poetry/poetry/issues/8427
	# The UV variant of the bug: https://github.com/astral-sh/uv/issues/6437
	#! Install `flash-attn`
	if ! uv pip list | grep "flash"; then
		echo "install_env.sh: Installing flash-attn..."
		uv pip install -q --no-build-isolation flash-attn==2.7.4.post1
	else
		echo "install_env.sh: flash-attn is already installed."
	fi
	uv run pre-commit install
	#! Final message
	echo "install_env.sh: Environment is ready."
fi

# Set TRITON_CACHE_DIR to be system dependent, as such must be under '/home/<username>'
TRITON_CACHE_DIR="/home/$(whoami)/.triton_cache"
export TRITON_CACHE_DIR

# Set the tmeporary directory for Ray
RAY_TEMP_DIR="/tmp/ray"
export RAY_TEMP_DIR
