#!/bin/bash
# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================

set -e
set -u

REPO_HOST="https://ascend-cann.obs.myhuaweicloud.com"
REPO_TYPE="CANN"
HTTP_OK="200"

# package file name
FILE_NAME_X86="ai_cann_x86.tar.gz"
FILE_NAME_ARM="ai_cann_arm.tar.gz"

# default download path and install path
WORKING_DIR=$(pwd)
DEFAULT_ROOT_DIR="/usr/local/Ascend"
DEFAULT_NORMAL_DIR="${HOME}/Ascend"

# package install options, require bash >= 4.0
# toolkit包的安装脚本的环境适应性不好，容易失败，因此手动解压+拷贝需要的头文件的方式完成安装
install_toolkit() {
  local run_name=$1
  local install_path=$2

  local tmp_toolkit_extract_path="${install_path}/tmp_toolkit"
  local option="--noexec --extract=${tmp_toolkit_extract_path}"
  echo "extract package ${key} option: ${option}"
  ./${run_name} ${option} && \
  chmod +w $(realpath ${install_path}/latest/include) && \
  cp -r ${tmp_toolkit_extract_path}/toolkit/include/experiment ${install_path}/latest/include/ && \
  rm -rf ${tmp_toolkit_extract_path}
}
declare -A PACKAGE_OPTIONS
PACKAGE_OPTIONS["runtime"]="--full -q --install-path=%s"
PACKAGE_OPTIONS["compiler"]="--full -q --pylocal --install-path=%s"
PACKAGE_OPTIONS["fwkplugin"]="--full -q --pylocal --install-path=%s"
PACKAGE_OPTIONS["opp"]="--full -q --install-path=%s"
PACKAGE_OPTIONS["toolkit"]="install_toolkit"
PACKAGE_OPTIONS["aoe"]="--full -q --install-path=%s"
PACKAGE_OPTIONS["ncs"]="--full -q --install-path=%s"
PACKAGE_OPTIONS["hccl"]="--full -q --pylocal --install-path=%s"

usage() {
  echo "Usage:"
  echo "    bash update_newest_packages.sh [-k] [-a] [-i <path>] [-d <path>] [-t <YearMonthDay>]"
  echo "Description:"
  echo "    -h, Display this information."
  echo "    -k, Allow insecure server connections when using SSL."
  echo "    -a, Select aarch64 run package. (Default: x86_64)"
  echo "    -i <path>, Install path. (Default: /usr/local/Ascend if run as root)"
  echo "    -d <path>, Download path for saving packages. (Default: install path if specified by -i else current path)"
  echo "    -t <YearMonthDay>, Specify the package version. (e.g., -t 20230821)"
}

checkopts() {
  CUSTOM_INSTALL_PATH=""
  NEWEST_DATE=""
  CA_OPTION=""
  CUSTOM_DOWNLOAD_PATH=""
  FILE_NAME=${FILE_NAME_X86}
  while getopts 'i:t:d:kah' opt; do
    case "${opt}" in
    i)
      CUSTOM_INSTALL_PATH=$(realpath ${OPTARG})
      ;;
    t)
      NEWEST_DATE=${OPTARG}
      ;;
    k)
      CA_OPTION="-k"
      ;;
    a)
      FILE_NAME=${FILE_NAME_ARM}
      ;;
    d)
      CUSTOM_DOWNLOAD_PATH=$(realpath ${OPTARG})
      ;;
    h)
      usage
      exit 1
      ;;
    *)
      log_error "Undefined option: ${opt}"
      usage
      exit 2
      ;;
    esac
  done
}

# ---------- logger functions ----------
log_info() {
  content="[INFO] $(date '+%Y-%m-%d %H:%M:%S') $@"
  echo -e "${content}"
}

log_error() {
  content="[ERROR] $(date '+%Y-%m-%d %H:%M:%S') $@"
  echo -e "\033[31m"${content} "\033[0m"
}

# ---------- package downloader ----------
build_local_dir() {
  echo "$1_newest"
}
build_download_url() {
  local full_time=$1
  local url="${REPO_HOST}/${REPO_TYPE}/$(build_local_dir $full_time)/${FILE_NAME}"
  echo ${url}
}

build_download_path() {
  local base_dir=$1
  local pkg_name="${2}_newest"
  local path="${base_dir}/${pkg_name}"
  echo ${path}
}

verify_url() {
  local url=$1
  local status=0
  set +e
  status=$(curl -k -I -m 5 -o /dev/null -s -w %{http_code} ${url})
  set -e
  echo ${status}
}

probe_newest_package() {
  local url=""
  local status=0
  local start=$(date +%Y%m%d)
  local end=$(date -d "-28 day ${start}" +%Y%m%d)
  log_info "Start to probe the newest package from ${start} to ${end}"
  while [ ${start} -ge ${end} ]; do
    url=$(build_download_url ${start})
    if [[ $(verify_url ${url}) -eq ${HTTP_OK} ]]; then
      # find the latest package and update the global variable
      NEWEST_DATE=${start}
      log_info "The latest package version found: [${NEWEST_DATE}_newest]"
      break
    fi
    start=$(date -d "-1 day ${start}" +%Y%m%d)
    url=""
    # set request interval to avoid potential Access Denied
    sleep 0.2
  done

  if [[ -z ${url} ]]; then
    log_error "Failed to find the newest package in the last 4 weeks, please check."
    exit 1
  fi
}

# ---------- package installer ----------
install_package() {
  local run_name=$1 # CANN-runtime-yyy.run
  if [ ! -f "${run_name}" ]; then
      log_error "Package ${run_name} does not exist."
      exit 1
  fi

  local package_name=$2 # runtime
  local option_fmt=$3
  if [ "install_${package_name}" = "${option_fmt}" ]; then
    log_info "Installing package ${run_name}...."
    ${option_fmt} $run_name $4
  else
    local option_fmt=$(echo "$3" | sed 's/--/\\055-/g')
    local cmd_option=$(printf "${option_fmt}" "$4")
    log_info "Installing package ${run_name}, install option: ${cmd_option}"
    ./${run_name} ${cmd_option}
  fi


}
extract_and_install_package() {
  local newest_pkg_path=$1
  local tar_name=$2
  local install_path="$3/$4"
  log_info "Begin to extract and install packages to ${install_path}."

  # extract packages
  cd ${newest_pkg_path}
  local dir_name=${tar_name::-7}
  tar -xf ${tar_name}
  cd ${dir_name}
  chmod u+x CANN-*.run

  # install packages
  for key in "${!PACKAGE_OPTIONS[@]}"; do
    local pattern="CANN-${key}-[0-9a-z.]{10,}-linux.(x86_64|aarch64).run"
    local package_name=$(ls ./ | grep -Eo "${pattern}")
    local options="${PACKAGE_OPTIONS[${key}]}"
    install_package "${package_name}" "${key}" "${options}" "${install_path}"
  done
  rm -rf $3/latest
  ln -s ${install_path}/latest $3/latest
  log_info "Complete installing packages: [${!PACKAGE_OPTIONS[*]}]"
}

# ---------- main procedure ----------
# check options
checkopts "$@"

# setup the package version
if [[ ${NEWEST_DATE} =~ ^20 ]]; then
  # user-specified package version
  log_info "Download package with user-specified version ${NEWEST_DATE}"
  DOWNLOAD_URL=$(build_download_url ${NEWEST_DATE})
  if [[ $(verify_url ${DOWNLOAD_URL}) -ne ${HTTP_OK} ]]; then
    log_error "Invalid package version: ${NEWEST_DATE}_newest. Please check the version and network accessibility."
    exit 1
  fi
else
  # automatically probe package version
  probe_newest_package
  DOWNLOAD_URL=$(build_download_url ${NEWEST_DATE})
fi

# check and build install path if CUSTOM_INSTALL_PATH is not specified
if [[ -z ${CUSTOM_INSTALL_PATH} ]]; then
  if [[ $(id -u) -eq 0 ]]; then
    # root user
    CUSTOM_INSTALL_PATH=${DEFAULT_ROOT_DIR}
  else
    # normal user
    CUSTOM_INSTALL_PATH=${DEFAULT_NORMAL_DIR}
  fi
  if [[ -z "${CUSTOM_DOWNLOAD_PATH}" ]]; then
    CUSTOM_DOWNLOAD_PATH="${WORKING_DIR}"
  fi
else
  if [[ -z "${CUSTOM_DOWNLOAD_PATH}" ]]; then
    CUSTOM_DOWNLOAD_PATH="${CUSTOM_INSTALL_PATH}/$(build_local_dir ${NEWEST_DATE})"
  fi
fi

# check and build download path
if [[ -d ${CUSTOM_DOWNLOAD_PATH} ]]; then
  log_error "The newest package [${CUSTOM_DOWNLOAD_PATH}]  already exists, please remove and retry."
  exit 1
fi

# display path info
log_info "Download packages from: ${DOWNLOAD_URL}"
log_info "Save packages to: ${CUSTOM_DOWNLOAD_PATH}/${FILE_NAME}"
log_info "Package's install path: ${CUSTOM_INSTALL_PATH}"

# make install directory and download package
mkdir -p ${CUSTOM_DOWNLOAD_PATH}
curl ${CA_OPTION} -o ${CUSTOM_DOWNLOAD_PATH}/${FILE_NAME} ${DOWNLOAD_URL}

# extract and install packages: [runtime, compiler, fwkplugin, opp, toolkit, aoe, hccl, ncs]
extract_and_install_package "${CUSTOM_DOWNLOAD_PATH}" "${FILE_NAME}" "${CUSTOM_INSTALL_PATH}" "$(build_local_dir ${NEWEST_DATE})"

log_info "Updated CANN successfully!"
log_info "When building metadef: set cmake option -DASCEND_INSTALL_PATH=${CUSTOM_INSTALL_PATH}/latest"
log_info "When build air: set environment variable ASCEND_INSTALL_PATH=${CUSTOM_INSTALL_PATH}/latest"