#!/bin/bash
#
# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

GPU_INSTALLER_DOWNLOAD_URL="${GPU_INSTALLER_DOWNLOAD_URL:-}"

get_major_version() {
  echo "$1" | cut -d "." -f 1
}

get_minor_version() {
  echo "$1" | cut -d "." -f 2
}

get_download_location() {
  # projects/000000000000/zones/us-west1-a -> us
  local -r instance_location="$(curl --http1.1 -sfS "http://metadata.google.internal/computeMetadata/v1/instance/zone" -H "Metadata-Flavor: Google" | cut -d '/' -f4 | cut -d '-' -f1)"
  declare -A location_mapping
  location_mapping=( ["us"]="us" ["asia"]="asia" ["europe"]="eu" )
  # Use us as default download location.
  echo "${location_mapping[${instance_location}]:-us}"
}

precompiled_installer_download_url() {
  local -r driver_version="$1"
  local -r milestone="$2"
  local -r build_id="${3//\./-}"  # 11895.86.0 -> 11895-86-0
  local -r major_version="$(get_major_version "${driver_version}")"
  local -r minor_version="$(get_minor_version "${driver_version}")"
  local -r download_location="$(get_download_location)"

  echo "https://storage.googleapis.com/nvidia-drivers-${download_location}-public/nvidia-cos-project/${milestone}/tesla/${major_version}_00/${driver_version}/NVIDIA-Linux-x86_64-${driver_version}_${milestone}-${build_id}.cos"
}

default_installer_download_url() {
  local -r driver_version="$1"
  local -r major_version="$(get_major_version "${driver_version}")"
  local -r minor_version="$(get_minor_version "${driver_version}")"
  local -r download_location="$(get_download_location)"

  if (( "${major_version}" < 390 )); then
    # Versions prior to 390 are downloaded from the upstream location.
    echo "https://us.download.nvidia.com/tesla/${driver_version}/NVIDIA-Linux-x86_64-${driver_version}.run"
  elif (( "${major_version}" == 390 )) && (( "${minor_version}" == 46 )); then
    # 390.46 is the only version residing in the TESLSA/ dir
    echo "https://storage.googleapis.com/nvidia-drivers-${download_location}-public/TESLA/NVIDIA-Linux-x86_64-${driver_version}.run"
  elif (( "${major_version}" == 396 )) && (( "${minor_version}" == 26 )); then
    # Different naming format for 396.26 including the -dignostic keyword.
    echo "https://storage.googleapis.com/nvidia-drivers-${download_location}-public/tesla/${driver_version}/NVIDIA-Linux-x86_64-${driver_version}-diagnostic.run"
  else
    # All other versions available in the gs conform to this naming convention.
    echo "https://storage.googleapis.com/nvidia-drivers-${download_location}-public/tesla/${driver_version}/NVIDIA-Linux-x86_64-${driver_version}.run"
  fi
}

get_gpu_installer_url() {
  if [[ -z "${GPU_INSTALLER_DOWNLOAD_URL}" ]]; then
    # First try to find the precompiled gpu installer.
    local -r url="$(precompiled_installer_download_url "$@")"
    if curl --http1.1 -s -I "${url}"  2>&1 | grep -q 'HTTP/2 200'; then
      GPU_INSTALLER_DOWNLOAD_URL="${url}"
    else
      # Fallback to default gpu installer.
      GPU_INSTALLER_DOWNLOAD_URL="$(default_installer_download_url "$@")"
    fi
  fi
  echo "${GPU_INSTALLER_DOWNLOAD_URL}"
}