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

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

eval set -- "$OPTIONS"

echo "bitsandbytes_installer.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 "bitsandbytes_installer.sh: arguments=%s, first argument=%s\n" "$@" "$1"
echo "bitsandbytes_installer.sh: Install env in PROJECT_PATH=$PROJECT_PATH"
# 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}}
export BNB_CUDA_VERSION="124"
#! Moving to the project folder
cd "$PROJECTS_FOLDER_PATH" || exit
#! Clone the repository of Bits & Bytes
rm -rf bitsandbytes
git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git && cd bitsandbytes/ || exit
#! Install the dependencies
uv pip install -r requirements-dev.txt
#! Function to get the compute capability of the GPU
get_compute_capability() {
	# Use nvidia-smi to get the GPU name
	GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -n 1)

	# Use nvcc to get the compute capability
	COMPUTE_CAPABILITY=$(nvcc --list-gpu-code | grep "$GPU_NAME" | awk -F'compute_' '{print $2}' | head -n 1)

	echo "$COMPUTE_CAPABILITY"
}

#! Get the compute capability of the current GPU
COMPUTE_CAPABILITY=$(get_compute_capability)

#! Set the NO_CUBLASLT variable to OFF if the compute capabilities of the current GPU are greater than 7.5
if (($(echo "$COMPUTE_CAPABILITY >= 7.5" | bc -l))); then
	export NO_CUBLASLT=OFF
else
	export NO_CUBLASLT=ON
fi

echo "NO_CUBLASLT is set to $NO_CUBLASLT"
#! Compile and install the package
cmake -DCOMPUTE_BACKEND=cuda -DNO_CUBLASLT=$NO_CUBLASLT -S .
make
uv pip install .
#! Move the compiled .so files to $HOME/lib
mkdir -p "$HOME/lib"
mv bitsandbytes/*.so "$HOME/lib"
#! Create a symbolic link of the compiled .so files to the environment folder
ln -s $HOME/lib/libbitsandbytes*.so $VIRTUAL_ENV/lib/python3.11/site-packages/bitsandbytes/
#! Export the compiled libs to the environment variables and the .bashrc
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$HOME/lib"
echo '# Adding local libs (with bitsandbytes) to the LD_LIBRARY_PATH' >>$HOME/.bashrc
echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/lib" >>"$HOME/.bashrc"
