#!/bin/bash
#
# Copyright 2020 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# Utility to manage COS extensions.

set -o errexit
set -o pipefail
set -o nounset

PROG_NAME=$(basename "$0")
readonly PROG_NAME
readonly OS_RELEASE="/etc/os-release"
readonly EXTENSIONS_CACHE="/var/lib/cos-extensions"
readonly ARTIFACT_REGISTRY_REGIONS=("us" "eu" "asia")
DEFAULT_GPU_INSTALLER="gcr.io/cos-cloud/cos-gpu-installer:v2.0.26"

COS_GPU_INSTALLER="${COS_GPU_INSTALLER:-}"

set_cos_gpu_installer() {
  if [[ -n "${COS_GPU_INSTALLER}" ]]; then
    return
  fi

  # Change gpu container image name if location of the vm is available
  # and it's from cos-cloud project.
  local instance_zone
  instance_zone="$(curl -s "http://metadata.google.internal/computeMetadata/v1/instance/zone?alt=text" -H "Metadata-Flavor: Google" || true)"
  if [[ -n "${instance_zone}" ]]; then
    # Syntax of instance_zone: projects/<project-id>/zones/<region>
    # Example: projects/475556798229/zones/us-central1-a
    local instance_region
    instance_region="$(echo "${instance_zone}"|cut -d/ -f4|cut -d- -f1)"
    for region in "${ARTIFACT_REGISTRY_REGIONS[@]}"; do
      if [[ "${region}" == "${instance_region}" ]]; then
        DEFAULT_GPU_INSTALLER="${instance_region}.${DEFAULT_GPU_INSTALLER}"
        break
      fi
    done
  fi
  COS_GPU_INSTALLER="${DEFAULT_GPU_INSTALLER}"
}

usage() {
  cat <<EOF

${PROG_NAME}: Utility to manage COS extensions.

Usage:
    ${PROG_NAME} [OPTIONS] COMMAND [ARGS]...

Options:
    -h, --help    print help message

Commands:
    list [-- --gpu-installer]             list all available COS extensions and
                                          their versions.

    install <extension> [-- -<version>]   install a COS extension. If no version
                                          is given, then the default version
                                          will be installed.

Additional Description:
    ${PROG_NAME} install gpu -- --prepare-build-tools

    The gpu extension can be invoked with a --prepare-build-tools optional
    argument that can be used to cache the toolchain for the installer.
    Caching the toolchain carries the overhead of ~1GB disk space on the
    stateful partition.
    Using this command only populates the cache and does NOT install the GPU
    drivers thus may saves time on downloading the toolchain during subsequent
    installations.

    ${PROG_NAME} install gpu -- --clean-build--tools
    Use this optional command to delete the cache for the toolchain present on
    the stateful partition.
}
EOF
  exit "${1}"
}

parse_args() {
  local args
  if ! args=$(getopt --options "h" --longoptions "help" -- "$@"); then
    usage 1
  fi

  eval set -- "${args}"
  while true; do
    case "$1" in
      -h|--help)
        usage 0
        ;;
      --)
        shift
        break
        ;;
      *)
        usage 1
        ;;
    esac
  done

  if [[ "$#" -eq 0 ]]; then
    usage 1
  fi

  case "$1" in
    list)
        shift
        list "$@"
      ;;
    install)
      if [[ "$#" -eq 2 ]]; then
        install "$2"
      elif [[ "$#" -ge 3 ]]; then
        extension="$2"
        shift 2
        install "${extension}" "$@"
      else
        usage 1
      fi
      ;;
    *)
      usage 1
      ;;
  esac
}

list() {
  if [[ "$#" -eq 0 ]]; then
    # shellcheck disable=SC2154
    printf "Available extensions for COS version %s-%s:\n\n" \
      "${VERSION_ID}" "${BUILD_ID}"
    echo "[gpu]"
    echo "gpu installer: ${DEFAULT_GPU_INSTALLER}"
    run_gpu_installer list 2>/dev/null
  elif [[ "$#" -eq 1 ]]; then
    case "$1" in
      --gpu-installer)
        echo "${DEFAULT_GPU_INSTALLER}"
        ;;
      *)
        echo "Unsupported argument $1"
        usage 1
        ;;
    esac
  else
    usage 1
  fi
}

install() {
  case "$1" in
    gpu)
      shift
      run_gpu_installer install "-host-dir=/var/lib/nvidia" "$@"
      ;;
    *)
      echo "Unsupported extension $1"
      exit 1
      ;;
    esac
}

check_arch() {
  arch=$(uname -m)
  if  [[ ${arch} != "x86_64" ]]; then
    echo "GPU installation is only supported on X86 for now.
    Current architecture detected: ${arch}"
    exit 1
  fi
}

run_gpu_installer() {
  check_arch
  local use_build_cache=false
  local clean_cache=false

  local installer_args=()
  for i in "$@"; do
    if [[ ${i} == '--clean-build-tools' ]]; then
      clean_cache=true
    else
      if [[ ${i} == '--prepare-build-tools' ]]; then
        use_build_cache=true
      fi
      installer_args+=("${i}")
    fi
  done

  local docker_args=(
    --rm
    --name="cos-gpu-installer"
    --privileged
    --net=host
    --pid=host
    --volume /dev:/dev
    --volume /:/root
  )

  if [[ "${clean_cache}" = true ]]; then
    echo "Cleaning cache present at: ${EXTENSIONS_CACHE}"
    rm -rf "${EXTENSIONS_CACHE}"
  else
    # use extensions cache(if it exists) by default
    if [[ "${use_build_cache}" = true || -d ${EXTENSIONS_CACHE} ]]; then
      docker_args+=(--volume "${EXTENSIONS_CACHE}/:/build/")
    fi

    /usr/bin/docker run "${docker_args[@]}" "${COS_GPU_INSTALLER}"\
      "${installer_args[@]}"
  fi
}

main() {
  # shellcheck source=/etc/os-release
  source "${OS_RELEASE}"
  set_cos_gpu_installer
  parse_args "$@"
}

main "$@"