#!/bin/bash
# Copyright 2019-2020 Huawei Technologies Co., Ltd
#
# 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
#
# http://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.
# ============================================================================

set -e
BASEPATH=$(cd "$(dirname $0)"; pwd)
OUTPUT_PATH="${BASEPATH}/output"
export BUILD_PATH="${BASEPATH}/build/"

# print usage message
usage()
{
  echo "Usage:"
  echo "sh build.sh [-h] [-f] [-a] [-c] [-p] [-v] [-g] [-j[n]]"
  echo ""
  echo "Options:"
  echo "    -h Print usage"
  echo "    -f Build fe ut/st"
  echo "    -a Build aicpu ut/st"
  echo "    -c Build ut/st with coverage tag"
  echo "    -u Build all ut/st"
  echo "    -p Build inference or train"
  echo "    -v Display build command"
  echo "    -d Enable download cmake compile dependency from gitee , default off"
  echo "    -j[n] Set the number of threads used for building AIR, default is 8"
  echo "to be continued ..."
}

# check value of input is 'on' or 'off'
# usage: check_on_off arg_value arg_name
check_on_off()
{
  if [[ "X$1" != "Xon" && "X$1" != "Xoff" ]]; then
    echo "Invalid value $1 for option -$2"
    usage
    exit 1
  fi
}

# parse and set options
checkopts()
{
  VERBOSE=""
  THREAD_NUM=8
  ENABLE_FE_LLT="off"
  ENABLE_AICPU_LLT="off"
  ENABLE_COV="off"
  PLATFORM="all"
  ENABLE_GITEE="off"
  PRODUCT="normal"
  ENABLE_UT="off"
  # Process the options
  while getopts 'fachj:vp:sgu' opt
  do
    OPTARG=$(echo ${OPTARG} | tr '[A-Z]' '[a-z]')
    case "${opt}" in
      f)
        ENABLE_FE_LLT="on"
        ;;
      a)
        ENABLE_AICPU_LLT="on"
        ;;
      c)
        ENABLE_COV="on"
        ;;
      h)
        usage
        exit 0
        ;;
      j)
        THREAD_NUM=$OPTARG
        ;;
      v)
        VERBOSE="VERBOSE=1"
        ;;
      p)
        PLATFORM=$OPTARG
        ;;
      s)
        ENABLE_GITEE="on"
        ;;
      g)
        PRODUCT=$OPTARG
        ;;
      u)
        ENABLE_UT="on"
        ;;
      *)
        echo "Undefined option: ${opt}"
        usage
        exit 1
    esac
  done
}

mk_dir() {
    local create_dir="$1"  # the target to make

    mkdir -pv "${create_dir}"
    echo "created ${create_dir}"
}

build_air()
{
  echo "create build directory and build AIR";

  mk_dir "${BUILD_PATH}"
  cd "${BUILD_PATH}"
  CMAKE_ARGS="-DBUILD_PATH=$BUILD_PATH"

  if [[ "X$ENABLE_COV" = "Xon" ]]; then
    CMAKE_ARGS="${CMAKE_ARGS} -DENABLE_COV=ON"
  fi

  if [[ "X$ENABLE_FE_LLT" = "Xon" ]]; then
    CMAKE_ARGS="${CMAKE_ARGS} -DENABLE_FE_LLT=ON"
    FE_LLT_TARGET="fe_ut fe_st"
  fi

  if [[ "X$ENABLE_AICPU_LLT" = "Xon" ]]; then
    CMAKE_ARGS="${CMAKE_ARGS} -DENABLE_AICPU_LLT=ON"
    CPU_LLT_TARGET="run_cpu_test"
  fi

  if [[ "X$ENABLE_UT" = "Xon" ]]; then
    CMAKE_ARGS="${CMAKE_ARGS} -DENABLE_AICPU_LLT=ON -DENABLE_FE_LLT=ON"
    CPU_LLT_TARGET="run_cpu_test"
    FE_LLT_TARGET="fe_ut fe_st"
  fi

  if [[ "X$ENABLE_GITEE" = "Xon" ]]; then
    CMAKE_ARGS="${CMAKE_ARGS} -DENABLE_GITEE=ON"
  fi

  CMAKE_ARGS="${CMAKE_ARGS} -DBUILD_OPEN_PROJECT=True -DCMAKE_INSTALL_PREFIX=${OUTPUT_PATH} -DPLATFORM=${PLATFORM} -DPRODUCT=${PRODUCT}"
  echo "${CMAKE_ARGS}"
  cmake ${CMAKE_ARGS} ..
  if [ $? -ne 0 ]
  then
    echo "execute command: cmake ${CMAKE_ARGS} .. failed."
    return 1
  fi
  COMMON_TARGET="fe.ini fusion_config.json"
  TARGET=${COMMON_TARGET}
  if [ "x${PLATFORM}" = "xtrain" ]
  then
    CPU_TARGET="aicpu_engine_common aicpu_ascend_engine aicpu_ascend_builder aicpu_tf_engine aicpu_tf_builder"
    FE_TARGET="aicore_utils opskernel fe"
    TARGET="${TARGET} ${CPU_TARGET} ${FE_TARGET}"
  elif [ "x${PLATFORM}" = "xinference" ]
  then
    CPU_TARGET="aicpu_builder_common atc_aicpu_engine_common atc_aicpu_ascend_engine \
                atc_aicpu_ascend_builder atc_aicpu_tf_engine atc_aicpu_tf_builder"
    FE_TARGET="atc_aicore_utils atc_opskernel atc_fe"
    TARGET="${TARGET} ${CPU_TARGET} ${FE_TARGET}"
  elif [ "x${PLATFORM}" = "xconfig" ]
  then
    TARGET="fe.ini fusion_config.json"
  elif [ "x${PLATFORM}" = "xall" ]
  then
    # build all the target
    TARGET=""
  fi

  if [ "X$ENABLE_AICPU_LLT" = "Xon" ] && [ "X$ENABLE_FE_LLT" != "Xon" ];then
    TARGET="${CPU_LLT_TARGET}"
    make ${VERBOSE} ${TARGET} -j${THREAD_NUM}
  elif [ "X$ENABLE_FE_LLT" = "Xon" ] && [ "X$ENABLE_AICPU_LLT" != "Xon" ];then
    TARGET="${FE_LLT_TARGET}"
    make ${VERBOSE} ${TARGET} -j${THREAD_NUM}
  elif [ "X$ENABLE_UT" = "Xon" ];then
    TARGET="${CPU_LLT_TARGET} ${FE_LLT_TARGET}"
    make ${VERBOSE} ${TARGET} -j${THREAD_NUM}
  else
    make ${VERBOSE} ${TARGET} -j${THREAD_NUM} && make install
  fi

  if [ $? -ne 0 ]
  then
    echo "execute command: make ${VERBOSE} -j${THREAD_NUM} && make install failed."
    return 1
  fi
  echo "AIR build success!"
}

generate_package()
{
  cd "${BASEPATH}"

  AIR_LIB_PATH="lib"
  FWK_PATH="fwkacllib/lib64"
  ATC_PATH="atc/lib64"
  OPSKERNEL_PATH="plugin/opskernel"
  OPSKERNEL_FE_CONFIG_PATH="plugin/opskernel/fe_config"
  OPSKERNEL_CPU_CONFIG_PATH="plugin/opskernel/config"

  COMMON_LIB=("libaicore_utils.so" "libopskernel.so" "libaicpu_engine_common.so")
  OPSKERNEL_LIB=("libfe.so" "libaicpu_ascend_engine.so" "libaicpu_ascend_builder.so"
                 "libaicpu_tf_engine.so" "libaicpu_tf_builder.so")
  PLUGIN_OPSKERNEL_FE=("fusion_config.json" "fe.ini")
  PLUGIN_OPSKERNEL_CPU=("init.conf" "ir2tf_op_mapping_lib.json")

  rm -rf ${OUTPUT_PATH:?}/${FWK_PATH}/
  rm -rf ${OUTPUT_PATH:?}/${ATC_PATH}/

  mk_dir "${OUTPUT_PATH}/${FWK_PATH}/${OPSKERNEL_FE_CONFIG_PATH}"
  mk_dir "${OUTPUT_PATH}/${FWK_PATH}/${OPSKERNEL_CPU_CONFIG_PATH}"
  mk_dir "${OUTPUT_PATH}/${ATC_PATH}/${OPSKERNEL_FE_CONFIG_PATH}"
  mk_dir "${OUTPUT_PATH}/${ATC_PATH}/${OPSKERNEL_CPU_CONFIG_PATH}"

  cd "${OUTPUT_PATH}"

  find ./ -name air.tar -exec rm {} \;

  MAX_DEPTH=1

  for lib in "${PLUGIN_OPSKERNEL_CPU[@]}";
  do
    find ${OUTPUT_PATH}/${AIR_LIB_PATH} -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${FWK_PATH}/${OPSKERNEL_CPU_CONFIG_PATH} \;
    find ${OUTPUT_PATH}/${AIR_LIB_PATH} -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${ATC_PATH}/${OPSKERNEL_CPU_CONFIG_PATH} \;
  done

  for lib in "${COMMON_LIB[@]}";
  do
    find ${OUTPUT_PATH}/${AIR_LIB_PATH}/atclib -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${ATC_PATH} \;
    find ${OUTPUT_PATH}/${AIR_LIB_PATH} -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${FWK_PATH} \;
  done

  for lib in "${OPSKERNEL_LIB[@]}";
  do
    find ${OUTPUT_PATH}/${AIR_LIB_PATH}/atclib -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${ATC_PATH}/${OPSKERNEL_PATH} \;
  done

  for lib in "${OPSKERNEL_LIB[@]}";
  do
    find ${OUTPUT_PATH}/${AIR_LIB_PATH} -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${FWK_PATH}/${OPSKERNEL_PATH} \;
  done

  for lib in "${PLUGIN_OPSKERNEL_FE[@]}";
  do
    find ${OUTPUT_PATH}/${AIR_LIB_PATH} -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${FWK_PATH}/${OPSKERNEL_FE_CONFIG_PATH} \;
    find ${OUTPUT_PATH}/${AIR_LIB_PATH} -maxdepth ${MAX_DEPTH} -name "$lib" -exec cp -f {} ${OUTPUT_PATH}/${ATC_PATH}/${OPSKERNEL_FE_CONFIG_PATH} \;
  done

  if [ "x${PLATFORM}" = "xtrain" ]
  then
    tar -cf air.tar fwkacllib
  elif [ "x${PLATFORM}" = "xinference" ]
  then
    tar -cf air.tar atc
  elif [ "x${PLATFORM}" = "xall" ]
  then
    tar -cf air.tar fwkacllib atc
  fi
}

main() {
  # AIR build start
  echo "---------------- AIR build start ----------------"
  checkopts "$@"
  chmod 755 third_party/metadef/graph/stub/gen_stubapi.py
  g++ -v

  mk_dir ${OUTPUT_PATH}
  build_air || { echo "AIR build failed."; exit 1; }
  echo "---------------- AIR build finished ----------------"

  chmod -R 750 ${OUTPUT_PATH}
  find ${OUTPUT_PATH} -name "*.so*" -print0 | xargs -0 chmod 500

  echo "---------------- AIR output generated ----------------"


  if [[ "X$ENABLE_UT" = "Xoff" ]] && [[ "X$ENABLE_FE_LLT" = "Xoff" ]] && [[ "X$ENABLE_AICPU_LLT" = "Xoff" ]]; then
    generate_package
  fi
  echo "---------------- AIR package archive generated ----------------"
}

main "$@"