#!/bin/bash
set -eu

keyring_version=1.1-1

arch=$(dpkg --print-architecture)
case "$arch" in
    amd64) arch=x86_64 ;;
esac
ubuntu_version=$(source /etc/os-release && echo ${VERSION_ID/.})
keyring_url=https://developer.download.nvidia.com/compute/cuda/repos/ubuntu$ubuntu_version/$arch/cuda-keyring_${keyring_version}_all.deb

description=$(cat <<EOF
This script is a convenience tool designed to download and enable the NVIDIA
CUDA repository using a package from the following URL:

  $keyring_url

This package will add the NVIDIA CUDA package repository to your system.

Note that the package and NVIDIA CUDA repository are not owned or controlled by
Canonical Ltd or the Ubuntu community.
EOF
)

thisprog=$(basename "$0")
show_usage()
{
    echo "Usage: $thisprog [-h | -y]"
}
show_help()
{
    show_usage
    echo
    echo "$description"
    echo
    echo "Options:"
    echo "  -h    Show this help"
    echo "  -y    Do not prompt to acknowledge the warning about running this script"
}

answer=

while getopts :yh option; do
    case $option in
        y) answer=y ;;
        h)
            show_help
            exit 0 ;;
        *)
            echo "Invalid option: -$OPTARG" >&2
            show_usage >&2
            exit 1 ;;
    esac
done
shift $((OPTIND-1))
if (( $# )); then
    show_usage >&2
    exit 1
fi

if dpkg-query -W cuda-keyring &>/dev/null; then
    echo "NVIDIA CUDA repository was already added"
    exit 0
fi

if (( EUID != 0 )); then
    echo -e "WARNING: This program should be run only by root\n" >&2
fi

echo "$description"
cat <<EOF

Note that if the download fails because the URL does not exist, it is likely
that the keyring has been rotated. In this case, you will need to add the
repository to your system manually. For more information, please see
https://docs.nvidia.com/cuda/cuda-installation-guide-linux/#network-repo-installation-for-ubuntu.

Note that the script sets priority 400 to all packages in the NVIDIA
CUDA repository via the /etc/apt/preferences.d/cuda-repository-400
file, whose records take precedence over those in the
/etc/apt/preferences/cuda-repository-600 file installed by the cuda-keyring
package. The Ubuntu archive is, in effect, preferred for packages present in
both, regardless of version, to avoid upgrading packages tied to the NVIDIA
driver version from the NVIDIA CUDA repository. However, it is preferred
even for packages not tied to the driver version, like sdkmanager. Please
see apt_preferences(5) manual page for more information.

EOF

if [[ -z $answer ]]; then
    read -r -p "Do you want to continue? [y/N] " answer
fi

case $answer in
    [yY]) ;;
    *) exit 1 ;;
esac

cat >/etc/apt/preferences.d/cuda-repository-pin-400 <<EOF
# This file forces the NVIDIA CUDA repository's pin priority to 400. It has
# precedence over /etc/apt/preferences.d/cuda-repository-pin-600 shipped in the
# cuda-keyring package, which sets the pin priority to 600.

Explanation: Do not upgrade Ubuntu packages from the NVIDIA CUDA repository
Package: *
Pin: release l=NVIDIA CUDA
Pin-Priority: 400
EOF

set -x
wget "$keyring_url" -O /tmp/cuda-keyring.deb
dpkg -i /tmp/cuda-keyring.deb
rm /tmp/cuda-keyring.deb
apt-get update
