#!/bin/bash

# Retrieve and apply the Nvidia driver 

set -euo pipefail

PROG=install-nvidia-driver
PUBLIC_ADDR=nvcr.io/nvidia
VERSION=v0.17.0
MANIFEST="https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/${VERSION}/deployments/static/nvidia-device-plugin.yml"

registry_namespace=''
registry_host=''

die() {
    echo "$PROG error: $*" >&2
    exit 1
}

usage() {
    echo "$@" >&2
    echo >&2
    echo 'USAGE: install-nvidia-driver ARGUMENTS...

Required Arguments
  --registry-namespace NAMESPACE   Desired namespace in private registry. Often "wallaroo".
  --registry-host HOST             Hostname and port of private registry.
' >&2
    exit 1
}

command -v curl &> /dev/null || die "'curl' must be installed"
command -v kubectl &> /dev/null || die "'kubectl' must be installed"

while [ $# -gt 0 ]; do
    case "$1" in
	--registry-namespace)  registry_namespace=$2; shift 2;;
	--registry-host)       registry_host=$2; shift 2;;
	*) usage "unknown argument $1";;
    esac
done

[[ -n "$registry_namespace" ]]	|| usage "missing argument --registry-namespace"
[[ -n "$registry_host" ]]	|| usage "missing argument --registry-host"

prepare-manifest() {
    repo=${registry_host}/${registry_namespace}
    curl -LSso /tmp/nvidia-manifest.yaml $MANIFEST
    echo downloaded manifest
    sed -i.orig \
	-e "s!${PUBLIC_ADDR}!${repo}!" \
	-e '/ spec:/a\
      imagePullSecrets: \
      - name: regcred	
       ' \
	/tmp/nvidia-manifest.yaml
    echo edited manifest
}

apply-manifest() {
    kubectl apply -f /tmp/nvidia-manifest.yaml
    echo applied manifest
}

prepare-manifest
apply-manifest