#!/bin/bash

# Retrieve and apply the Nvidia driver 

set -euo pipefail

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

region=''
export AWS_PAGER=""

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

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

Required Arguments
  --region REGION                AWS region. Examples: us-east-1, us-west-2, etc.
' >&2
    exit 1
}

command -v aws &> /dev/null || die "'aws' CLI must be installed and configured via env or ~/.aws/credentials"
command -v curl &> /dev/null || die "'curl' must be installed"
command -v jq &> /dev/null || die "'jq' must be installed"
command -v kubectl &> /dev/null || die "'kubectl' must be installed"

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

[[ -n "$region" ]] || usage "missing argument --region"

check-image() {
    aws ecr describe-images --region "$region" --repository-name "$ECR_REPO" --image-ids "imageTag=${VERSION}" >& /dev/null ||
	die "first use load-images-ecr --load-image $PUBLIC_REPO"
    echo image "$PUBLIC_REPO" exists in ecr
}

get-image-uri() {
    aws ecr describe-repositories --region "${region}" --repository-name "$ECR_REPO" |
	jq -r '.repositories | .[] | .repositoryUri'
}

prepare-manifest() {
    ecr_repo=$(get-image-uri):${VERSION}
    curl -LSso /tmp/nvidia-manifest.yaml $MANIFEST
    echo downloaded manifest
    sed -i.orig "s!${PUBLIC_IMAGE}!${ecr_repo}!" /tmp/nvidia-manifest.yaml
    echo edited manifest
}

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

check-image
prepare-manifest
apply-manifest