#!/bin/bash
#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
#  with the License. A copy of the License is located at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
#  OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
#  and limitations under the License.

set -x

SEMAPHORE_DIR="/root/bootstrap/semaphore"
INSTALL_POST_REBOOT_FINISHED_LOCK="${SEMAPHORE_DIR}/install_post_reboot_finished.lock"

PREBAKING_AMI="true"
MODULE_ID="vdi-app"

while getopts m:g:p:e:n:o:d:i:t: opt
do
    case "${opt}" in
      m) MODULE_ID=${OPTARG};;
      g) GPU_FAMILY=${OPTARG};;
      p) PREBAKING_AMI=${OPTARG};;
      e) ENVIRONMENT_NAME=${OPTARG};;
      n) PROJECT_NAME=${OPTARG};;
      o) SESSION_OWNER=${OPTARG};;
      d) SESSION_TYPE=${OPTARG};;
      i) SESSION_ID=${OPTARG};;
      t) CUSTOM_BROKER_URL=${OPTARG};;
      ?) echo "Invalid option for install_post_reboot.sh script: -${opt}."
         exit 1;;
    esac
done


if [[ ! $GPU_FAMILY =~ ^(NONE|NVIDIA|AMD)$ ]]; then
  echo "ERROR: -g argument must be either NONE, NVIDIA, or AMD (GPU Instance)."
  exit 1
fi

if [[ $INSTALL_SCRIPT_EXEC == 0 ]]; then
  echo "ERROR: install.sh must be executed first."
  exit 1
fi

if [[ ! $PREBAKING_AMI =~ ^(true|false)$ ]]; then
  echo "ERROR: -p argument must be either true or false."
  exit 1
fi

SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )

if [[ ! -f ${INSTALL_POST_REBOOT_FINISHED_LOCK} ]]; then

  timestamp=$(date +%s)
  exec > /root/bootstrap/logs/install_post_reboot.log.${timestamp} 2>&1

  echo -n "no" > ${BOOTSTRAP_DIR}/reboot_required.txt

  source "${SCRIPT_DIR}/../../common/linux/bootstrap_common.sh"
  BASE_OS=$(get_base_os)

  if [[ ! -f ${BOOTSTRAP_DIR}/res_installed_all_packages.log ]]; then
    if [[ $SESSION_ID != NONE ]]; then
      # Begin: Install Nice DCV Server
      /bin/bash "${SCRIPT_DIR}/../../dcv/linux/dcv_server.sh" -o $BASE_OS -g $GPU_FAMILY -s "${SCRIPT_DIR}/.."
      # End: Install Nice DCV Server

      # Begin: Install Nice DCV Session Manager Agent
      /bin/bash "${SCRIPT_DIR}/../../dcv/linux/dcv_session_manager_agent.sh" -o $BASE_OS -s "${SCRIPT_DIR}/.."
      # End: Install Nice DCV Session Manager Agent
    fi

    # Begin: Install AWS Systems Manager Agent
    # Install SSM agent at the end to avoid unexpected reboot for security patching
    /bin/bash "${SCRIPT_DIR}/../../common/linux/aws_ssm.sh" -o $BASE_OS -s "${SCRIPT_DIR}/.."
    # End: Install AWS Systems Manager Agent

    echo "$(date)" >> /root/bootstrap/res_installed_all_packages.log
    set_reboot_required "DCV and any associated GPU drivers have been installed, reboot required for changes to take effect..."
  else
     log_info "Found ${BOOTSTRAP_DIR}/res_installed_all_packages.log... skipping package installation..."
  fi

  mkdir -p ${SEMAPHORE_DIR}
  echo $(date +%s) > ${INSTALL_POST_REBOOT_FINISHED_LOCK}

  REBOOT_REQUIRED=$(cat /root/bootstrap/reboot_required.txt)
  if [[ "${REBOOT_REQUIRED}" == "yes" && "${PREBAKING_AMI}" == "false" ]]; then
    log_info "Triggering 2nd Reboot Required for RES Configuration for VDI"
    echo -n "no" > /root/bootstrap/reboot_required.txt

    reboot

    # Exit to avoid running the rest of script before reboot is completed.
    exit 0
  fi
fi

# Prebake running on EC2 Image Builder does not need to run the configure.sh script
if [[ "${PREBAKING_AMI}" == "false" ]]; then
  /bin/bash "${SCRIPT_DIR}/../../common/linux/install_app.sh" -s "${SCRIPT_DIR}/.." -c "virtual-desktop-app" -m "${MODULE_ID}" -e "${ENVIRONMENT_NAME}" -n "${PROJECT_NAME}" -o "${SESSION_OWNER}" -d "${SESSION_TYPE}" -i "${SESSION_ID}" -t "${CUSTOM_BROKER_URL}"
fi
