#!/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

readonly PROG_NAME=$(basename "$0")
readonly COS_GPU_INSTALLER="${COS_GPU_INSTALLER:-\
gcr.io/cos-cloud/cos-gpu-installer:v2.0.11}"
readonly OS_RELEASE="/etc/os-release"
readonly EXTENSIONS_CACHE="/var/lib/cos-extensions"

usage() {
  cat <<EOF

${PROG_NAME}: Utility to manage COS extensions.

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

Options:
    -h, --help    print help message

Commands:
    list                                  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 --use-build-tools-cache
    The gpu extension can be invoked with a --use-build-tools-cache
    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. It may also help to save time on downloading the
    toolchain during the installation, after the first run.
}
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)
      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() {
  printf "Available extensions for COS version %s-%s:\n\n" \
    "${VERSION_ID}" "${BUILD_ID}"
  echo "[gpu]"
  run_gpu_installer list 2>/dev/null
}

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 installer_args=()
  for i in "$@"; do
    if [[ $i == '--use-build-tools-cache' ]]; then
      use_build_cache=true
    else
      installer_args+=($i)
    fi
  done

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

  if [[ "$use_build_cache" = true ]]; then
    docker_args+=(--volume ${EXTENSIONS_CACHE}/:/build/)
  fi

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

main() {
  source "${OS_RELEASE}"
  parse_args "$@"
}

main "$@"