#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091,SC2317
## This script aims to setup the OS for a (cloud) VM
## starting from the "Plain Ubuntu 24.04" image
#! Update and upgrade package manager
sudo apt update
#! Installing the essentials
sudo apt install -y build-essential zlib1g-dev libedit-dev \
	libssl-dev liblzma-dev libffi-dev libbz2-dev \
	libreadline-dev libsqlite3-dev bc clang
sudo apt upgrade -y
#! This last step requires reboot to load the new kernel (if available)
echo "YOU SHOULD NOW REBOOT AND THEN RE-EXECUTE THE SCRIPT COMMENTING LINE 14"
exit 1

#! Install CUDA latest from the NVIDIA repository
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt update
sudo apt -y install cuda-toolkit cuda-drivers nvidia-fabricmanager libnvidia-nscq
sudo systemctl enable nvidia-fabricmanager
sudo systemctl start nvidia-fabricmanager
#! Set GPU persistence mode
sudo nvidia-smi -pm 1

#! Install the CUDA toolkit for 12.8, version which is used to compile our PyTorch/FA dependencies
#! Check the output of `nvcc -V`
NVCC_OUTPUT=$(nvcc -V)
if [[ $NVCC_OUTPUT == *"release 12.8"* ]]; then
	echo "CUDA 12.8 is detected."
else
	#! Get and install CUDA 12.8.1 and its drivers
	wget https://developer.download.nvidia.com/compute/cuda/12.8.1/local_installers/cuda_12.8.1_570.124.06_linux.run
	sudo sh cuda_12.8.1_570.124.06_linux.run --toolkit --silent
fi
if [[ $PATH == *"cuda-12.8"* ]]; then
	echo "PATH variable is already set."
else
	#! Set the PATH env variables for the new CUDA version
	echo '# Adding CUDA 12.8 to the PATH environmental variables' >>~/.bashrc
	# shellcheck disable=SC2016
	echo 'export PATH=/usr/local/cuda-12.8/bin${PATH:+:${PATH}}' >>~/.bashrc
	export PATH=/usr/local/cuda-12.8/bin${PATH:+:${PATH}}
fi
if [[ $LD_LIBRARY_PATH == *"cuda-12.8"* ]]; then
	echo "LD_LIBRARY_PATH variable is already set."
else
	#! Set the LD_LIBRARY_PATH env variables for the new CUDA version
	# shellcheck disable=SC2016
	echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >>~/.bashrc
	export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
fi

#! Install `UV`
UV_VER_OUTPUT=$(uv --version)
if [[ $UV_VER_OUTPUT == *"uv "* ]]; then
	echo "UV is already installed."
else
	#! Getting `UV`
	curl -LsSf https://astral.sh/uv/install.sh | sh
fi

#! Upgrade pip
sudo pip install --upgrade pip
#! Monitoring utilities and cmake
sudo apt -y install bpytop nvitop
sudo pip install cmake
