#!/bin/bash
# Print commands and their arguments as they are executed.
## Quick usage example:
# ./get-cuda-mirror.sh --nvidia-driver=575 --cuda-toolkit=12-9 --nvitop=true

# === Parse Parameters ===
NVIDIA_DRIVER=""
CUDA_TOOLKIT=""
NVITOP=false
VM=false
FORCE=false
UNINSTALL_ALL=false

for arg in "$@"; do
    case $arg in
        --nvidia-driver=*)
            NVIDIA_DRIVER="${arg#*=}"
            shift
            ;;
        --cuda-toolkit=*)
            CUDA_TOOLKIT="${arg#*=}"
            shift
            ;;
        --nvitop=true)
            NVITOP=true
            shift
            ;;
        --vm=true)
            VM=true
            shift
            ;;
        --force-install=true)
            FORCE=true
            shift
            ;;
        --uninstall-all=true)
            UNINSTALL_ALL=true
            shift
            ;;
        *)
            echo "Unknown option: $arg"
            ;;
    esac
done

# if all options are empty or false, print usage and exit
if [[ -z "$NVIDIA_DRIVER" && -z "$CUDA_TOOLKIT" && "$NVITOP" == false && "$VM" == false && "$FORCE" == false && "$UNINSTALL_ALL" == false ]]; then
    echo "Usage: ./get-cuda-mirror.sh [--nvidia-driver=VERSION] [--cuda-toolkit=VERSION] [--nvitop=true] [--vm=true] [--force-install=true] [--uninstall-all=true]"
    echo "Examples:"
    echo "  Install: ./get-cuda-mirror.sh --nvidia-driver=575 --cuda-toolkit=12-9 --nvitop=true"
    echo "  Install for VM: ./get-cuda-mirror.sh --nvidia-driver=575 --cuda-toolkit=12-9 --nvitop=true --vm=true"
    echo "  Force Install: ./get-cuda-mirror.sh --force-install=true"
    echo "  Uninstall: ./get-cuda-mirror.sh --uninstall-all=true"
    exit 0
fi

# Detect Ubuntu version like ubuntu2204
linux_version="$(. /etc/os-release && echo ${ID})"
version_id="$(. /etc/os-release && echo ${VERSION_ID//./})"

# Detect architecture (e.g. amd64, arm64)
distro="${linux_version}${version_id}"
arch="$(lscpu | grep "Architecture:" | awk '{print $2}')"

echo "Check distro: $distro"
echo "Check arch: $arch"

echo "Install the new cuda-keyring package:"
wget https://mirrors.siam.ai/compute/cuda/repos/$distro/$arch/cuda-keyring_1.1-1_all.deb
dpkg -i cuda-keyring_1.1-1_all.deb
apt update

echo "Enroll the new signing key"
wget https://mirrors.siam.ai/compute/cuda/repos/$distro/$arch/cuda-archive-keyring.gpg
sudo mv cuda-archive-keyring.gpg /usr/share/keyrings/cuda-archive-keyring.gpg

echo "Enable the network repository"
echo "deb [signed-by=/usr/share/keyrings/cuda-archive-keyring.gpg] https://mirrors.siam.ai/compute/cuda/repos/$distro/$arch/ /" \
    | sudo tee /etc/apt/sources.list.d/cuda-$distro-$arch.list

# === Optional Step 1. Uninstall all cuda related packages ===
if [[ "$UNINSTALL_ALL" == true ]]; then
    echo "==== Uninstalling all cuda related packages ===="
    sudo apt-get -y remove --purge '^nvidia-.*' '^cuda-.*' '^libnvidia-.*' 'nvidia-cuda-toolkit' 'nvidia-container-toolkit'
    sudo apt-get -y autoremove
    echo "==== All cuda related packages uninstalled ===="
    exit 0
fi

# === Optional Step 2. Force Install specific versions of NVIDIA Driver, CUDA Toolkit, and NVITOP ===
if [[ "$FORCE" == true ]]; then
    NVIDIA_DRIVER="575_575.57.08"
    CUDA_TOOLKIT="12-9"
    NVITOP=true
    echo "=== Downloading libnvidia-gl package ==="
    wget -O /tmp/libnvidia-gl-575_575.57.08-0ubuntu1_amd64.deb \
    "https://mirrors.siam.ai/compute/cuda/repos/ubuntu2204/x86_64/libnvidia-gl-575_575.57.08-0ubuntu1_amd64.deb"

    echo "=== Installing libnvidia-gl package ==="
    apt install -y /tmp/libnvidia-gl-575_575.57.08-0ubuntu1_amd64.deb

    echo "=== Downloading NVIDIA driver package ==="
    wget -O /tmp/nvidia-driver-575_575.57.08-0ubuntu1_amd64.deb \
    "https://mirrors.siam.ai/compute/cuda/repos/ubuntu2204/x86_64/nvidia-driver-575_575.57.08-0ubuntu1_amd64.deb"

    echo "=== Installing NVIDIA driver package ==="
    apt install -y /tmp/nvidia-driver-575_575.57.08-0ubuntu1_amd64.deb

    echo "=== Downloading NVIDIA Fabric Manager package ==="
    wget -O /tmp/nvidia-fabricmanager-575_575.57.08-1_amd64.deb \
    "https://mirrors.siam.ai/compute/cuda/repos/ubuntu2204/x86_64/nvidia-fabricmanager-575_575.57.08-1_amd64.deb"

    echo "=== Installing NVIDIA Fabric Manager package ==="
    apt install -y /tmp/nvidia-fabricmanager-575_575.57.08-1_amd64.deb

    echo "=== Enabling and starting NVIDIA Fabric Manager ==="
    systemctl enable nvidia-fabricmanager.service
    systemctl start nvidia-fabricmanager.service

    echo "=== Fixing dependencies and updating package cache ==="
    apt update -y
    apt install -f -y

    apt-get update
    echo "==== All Done ===="
    exit 0
fi

# === Optional Step 3. Install NVIDIA Driver ===
if [[ -n "$NVIDIA_DRIVER" ]]; then
    echo "==== Installing NVIDIA Driver version $NVIDIA_DRIVER ===="
    sudo apt-get -y install nvidia-driver-"$NVIDIA_DRIVER"
fi

# === Optional Step 4. Install CUDA Toolkit ===
if [[ -n "$CUDA_TOOLKIT" ]]; then
    echo "==== Installing CUDA Toolkit version $CUDA_TOOLKIT ===="
    sudo apt-get -y install cuda-toolkit-"$CUDA_TOOLKIT"
    sudo apt-get -y install cuda-"$CUDA_TOOLKIT"
    sudo apt-get -y install -y nvidia-container-toolkit # install toolkit
fi

# === Optional Step 5. Install nvitop ===
if [[ "$NVITOP" == true ]]; then
    echo "==== Installing nvitop ===="
    sudo apt-get -y install python3-pip
    pip3 install --upgrade pip
    pip3 install nvitop
fi

# === Optional Step 6. Disable nvlink for vm ===
if [[ "$VM" == true ]]; then
    echo "blacklist nvidia_uvm" | tee /etc/modprobe.d/nvlink-denylist.conf && echo "options nvidia NVreg_NvLinkDisable=1" | tee /etc/modprobe.d/disable-nvlink.conf
    update-initramfs -u
    echo "!!! For VM, please reboot the system to take effect !!!"
fi

apt update
echo "==== All Done ===="