#!/bin/bash
afot_path=$(cd "$(dirname "$0")" && pwd)
config_file=${afot_path}/a-fot.ini

# profile名字
profile_name="profile.data"
# gcov名字
gcov_name="profile.gcov"
# 内核编译选项
kernel_configs=()

# 解析配置文件配置项
function parse_config() {
  if [[ ! -f ${config_file} ]]; then
    echo "[ERROR] Could not load config file at: ${config_file}, please check!"
    exit 1
  fi
  while read line || [[ -n ${line} ]]; do
    if [[ ! ${line} =~ "#" ]] && [[ ${line} != "" ]]; then
      key=$(echo ${line} | awk -F "=" '{print $1}')
      value=$(echo ${line} | awk -F "=" '{print $2}')
      if [[ ${key} =~ "CONFIG_" ]]; then
        kernel_configs+=("${key}=${value}")
      else
        eval "${key}=${value}"
      fi
    fi
  done <${config_file}
}

# 解析输入的参数
function parse_input_params() {
  while [ $# -ge 2 ]; do
    case $1 in
    --opt_mode)
      opt_mode=$2
      shift 2
      ;;
    --perf_time)
      perf_time=$2
      shift 2
      ;;
    --app_name)
      application_name=$2
      shift 2
      ;;
    --bin_file)
      bin_file=$2
      shift 2
      ;;
    --build_script)
      build_script=$2
      shift 2
      ;;
    --work_path)
      work_path=$2
      shift 2
      ;;
    --run_script)
      run_script=$2
      shift 2
      ;;
    --max_waiting_time)
      max_waiting_time=$2
      shift 2
      ;;
    --compiler)
      compiler=$2
      shift 2
      ;;
    --compiler_path)
      compiler_path=$2
      shift 2
      ;;  
    --config_file)
      config_file=$2
      shift 2
      ;;
    --check_success)
      check_success=$2
      shift 2
      ;;
    --build_mode)
      build_mode=$2
      shift 2
      ;;
    --pgo_mode)
      pgo_mode=$2
      shift 2
      ;;
    --pgo_phase)
      pgo_phase=$2
      shift 2
      ;;
    --kernel_src)
      kernel_src=$2
      shift 2
      ;;
    --kernel_name)
      kernel_name=$2
      shift 2
      ;;
    --run_time)
      run_time=$2
      shift 2
      ;;
    --CONFIG_*)
      kernel_configs+=("${1:2}=$2")
      shift 2
      ;;
    --last_time)
      last_time=$2
      shift 2
      ;;
    -s)
      silent=1
      shift
      ;;
    -n)
      disable_compilation=1
      shift
      ;;
    --makefile)
      makefile=$2
      shift 2
      ;;
    --kernel_config)
      kernel_config=$2
      shift 2
      ;;
    --data_dir)
      data_dir=$2
      shift 2
      ;;
    *)
      suggest_info
      exit 1
      ;;
    esac
  done

  if [[ $# == 1 ]]; then
    if [[ $1 == "-s" ]]; then
      silent=1
    elif [[ $1 == "-n" ]]; then
      disable_compilation=1
    else
      suggest_info
      exit 1
    fi
  fi
}

# 检查某个配置项是否存在
function check_item() {
  if [[ -z $1 ]]; then
    echo "[ERROR] The configuration item '$2' is missing, please check!"
    exit 1
  fi
}

function check_config_items() {
  check_item ${application_name} application_name
  check_item ${bin_file} bin_file
  check_item ${work_path} work_path
  check_item ${build_script} build_script
  check_item ${run_script} run_script
  check_item ${max_waiting_time} max_waiting_time
  check_item ${opt_mode} opt_mode
  check_item ${perf_time} perf_time
  check_item ${compiler} compiler
  check_item ${compiler_path} compiler_path
  check_item ${check_success} check_success
  check_item ${build_mode} build_mode
}

function suggest_info() {
  echo """
Usage: a-fot [OPTION1 ARG1] [OPTION2 ARG2] [...]

For perf mode:
--config_file       Path of configuration file
--opt_mode          Optimization modes (AutoFDO/AutoPrefetch/AutoBOLT/AutoPGO)
--perf_time         Perf sampling duration (unit: seconds)
--compiler          Compiler (gcc or llvm for now)
--compiler_path     Compiler path
--app_name          Application process name
--bin_file          Executable binary file path
--build_script      Application build script path
--work_path         Script working directory (used to compile the application and store the profile)
--run_script        Script path for running application
--max_waiting_time  Maximum binary startup time (unit: seconds)
--check_success     Check optimization result
--build_mode        Execute build script mode (Wrapper/Bear)

For kernel PGO mode:
--config_file       Path of configuration file
--opt_mode          Optimization mode (Auto_kernel_PGO)
--pgo_mode          PGO mode (arc/all)
--pgo_phase         Phase of kernel PGO (1/2)
--kernel_src        Kernel source directory
--kernel_name       Kernel local version name (will be appended with "-pgoing" or "-pgoed")
--work_path         Script working directory (used to store the profile and the log)
--run_script        Script path for running application
--compiler_path     Compiler gcc path
--CONFIG_...        Kernel building configuration
--last_time         Last time directory before rebooting (used to put log infos together)
-s                  Silent mode (reboot automatically after kernel installation)
-n                  Do not compile kernel automatically
--makefile          Makefile path of kernel
--kernel_config     Config file path of kernel
--data_dir          Profile path generated by kernel
"""
}

# 根据模式加载不同的优化脚本
function load_script() {
  case ${opt_mode} in
  "AutoFDO")
    source ${afot_path}/auto_fdo.sh
    ;;
  "AutoPrefetch")
    source ${afot_path}/auto_prefetch.sh
    ;;
  "AutoBOLT")
    source ${afot_path}/auto_bolt.sh
    ;;
  "AutoPGO")
    source ${afot_path}/auto_pgo.sh
    ;;
  *)
    echo "[ERROR] Optimization mode ${opt_mode} is not supported, check the configuration item: opt_mode"
    exit 1
    ;;
  esac
}

# 公共依赖检查项
function check_common_dependency() {
  get_arch=$(arch)
  if [[ ${get_arch} =~ "x86_64" || ${get_arch} =~ "aarch64" ]]; then
    echo "[INFO] Current arch: ${get_arch}"
  else
    echo "[ERROR] Unsupport arch: ${get_arch}"
    exit 1
  fi
  if ! type perf &>/dev/null; then
    echo "[ERROR] Optimization mode ${opt_mode} but perf is missing, try 'yum install perf'."
    exit 1
  fi
  is_file_exist ${build_script}
  is_file_exist ${run_script}
  case ${compiler} in 
  "gcc") 
    is_file_exist "${compiler_path}/bin/gcc"
    ;;
  "llvm") 
    is_file_exist "${compiler_path}/bin/clang" 
    ;;
  *)
    echo "[ERROR] compiler ${compiler} is not supported, check the configuration item: compiler"
    exit 1
    ;;
  esac
}

# 拆分编译数据库
function split_option() {
  if [ "$bear_prefix" ]; then
    python3 $afot_path/split_json.py -i $PWD/compile_commands.json
    mv $PWD/compile_commands.json $PWD/compile_commands_$1.json
    mv $PWD/compile_commands.fail.json $PWD/compile_commands.fail_$1.json
  fi
}

# 使用原始编译选项进行编译
function first_compilation() {
  echo "[INFO] Start raw compilation."
  is_file_exist ${build_script} "build_script"
  if [[ $build_mode =~ "Bear" ]]; then
    bear_prefix="bear -- "
    echo "[INFO] Build in Bear mode."
  else
    echo "[INFO] Build in Wrapper mode."
  fi
  $bear_prefix /bin/bash ${build_script} >>${log_file} 2>&1
  split_option first
  is_file_exist ${bin_file}
  is_success $?
}

# 创建wrapper之后的操作
function post_create_wrapper() {
  chmod 755 ${compiler_wrapper}/${c_compiler}
  chmod 755 ${compiler_wrapper}/${cxx_compiler}

  export CC=${compiler_wrapper}/${c_compiler}
  export CXX=${compiler_wrapper}/${cxx_compiler}
  if [ "$compiler" = "gcc" ]; then
    export LD_LIBRARY_PATH=${compiler_path}/lib64:${LD_LIBRARY_PATH}
  else
    export LD_LIBRARY_PATH=${compiler_path}/lib:${LD_LIBRARY_PATH}
  fi

  export PATH=${compiler_wrapper}:${PATH}
}

# 执行应用程序执行脚本
function execute_run_script() {
  echo "[INFO] Start to execute the run_script: ${run_script}"
  process_id=$(pidof ${application_name})
  if [[ -n ${process_id} ]]; then
    echo "[ERROR] Application: ${application_name} process already exists. The run_script will not be executed. Please check!"
    exit 1
  fi
  is_file_exist ${run_script} "run_script"
  /bin/bash ${run_script} >>${log_file} 2>&1 &
  application_pid=$!
  is_success $?
}

# 探测应用进程是否存在
function detect_process() {
  echo "[INFO] Start to detect whether process ${application_name} is started."
  detect_time=0
  while [ -z $(pidof ${application_name}) ]; do
    sleep 1
    ((detect_time = ${detect_time} + 1))
    echo "[INFO] Finding ${application_name}."
    if [[ ${detect_time} -gt ${max_waiting_time} ]]; then
      echo "[ERROR] Process ${application_name} is not found after ${max_waiting_time}. Please check!"
      exit 1
    fi
  done
  echo "[INFO] Found Process ${application_name}: $(pidof ${application_name})"
}

# 使用新的编译选项编译，同时判断优化结果.
# 需注意此检查依赖wrapper中编译器后添加第一个编译选项，
# 因此需保证编译器后添加第一个编译选项为优化选项而非通用选项
function second_compilation() {
  echo "[INFO] Try compiling with the new compilation options."
  if [[ ${check_success} -eq 1 ]]; then
    $bear_prefix /bin/bash ${build_script} >>${log_file} 2>&1 & build_id=$!
    echo "[INFO] Found build id: ${build_id}"
    add_opt=$(cat ${compiler_wrapper}/${c_compiler} | awk -F " " '{print $2}')
    build_status=$(ps -p ${build_id} | grep -c ${build_id})
    opt_success=0
    while [[ ${build_status} -ne 0 ]]; do
      if [[ ${opt_success} -eq 0 ]]; then
        # 使用:1去除编译选项左边的'-'
        if [[ $(ps aux | grep -c "${add_opt:1}") -gt 1 ]]; then
          opt_success=1
          break
        fi
      fi
      build_status=$(ps -p ${build_id} | grep -c ${build_id})
    done
    wait
  else
    $bear_prefix /bin/bash ${build_script} >>${log_file} 2>&1
  fi
  echo "[INFO] Finish compiling with new compilation options."
  is_file_exist ${bin_file}
  split_option second
  is_success $?
}

# 判断上一步执行是否成功
function is_success() {
  pre_result=$1
  if [[ ${pre_result} -ne 0 ]]; then
    echo "[ERROR] Execution failed, please check the log: ${log_file}"
    exit 1
  fi
}

# 检查配置文件脚本是否存在
function is_file_exist() {
  file=$1
  config_item=$2
  if [[ ! -f ${file} ]]; then
    if [[ -n ${config_item} ]]; then
      echo "[ERROR] The file ${file} does not exist! Check the configuration item: ${config_item}"
    else
      echo "[ERROR] The file ${file} does not exist!"
    fi
    exit 1
  fi
}

#初始化profile文件夹和log文件
function init_profile_and_log() {
  # profile和log所在路径
  now_time=$(date '+%Y%m%d-%H%M%S')
  profile_data_path=${work_path}/${now_time}
  log_file=${work_path}/${now_time}/opt.log
  if [[ ! -d ${profile_data_path} ]]; then
    mkdir -p ${profile_data_path}
  fi
  echo "[INFO] Create profile dir: ${profile_data_path}"

  touch ${log_file}
  echo "[INFO] Init log file: ${log_file}"

  # 设置c和c++编译器
  if [ "$compiler" = "gcc" ]; then
    c_compiler="gcc"
    cxx_compiler="g++"
  else
    c_compiler="clang"
    cxx_compiler="clang++"
  fi  

  # 创建Wrapper所在路径
  compiler_wrapper="${work_path}/${now_time}/compiler_wrapper/"
  mkdir -p ${compiler_wrapper}
}

#检测是否优化成功
function is_opt_success() {
  if [[ ${check_success} -eq 1 ]]; then
    if [[ ${opt_success} -eq 0 ]]; then
      echo "[WARNING] Optimization may fail or the build process is too short, please check!"
      echo "[WARNING] Please try gcc/g++ or clang/clang++ at: ${compiler_wrapper} instead of the original compiler."
      exit 1
    else
      echo "[INFO] Optimization success!"
    fi
  fi
  exit 0
}

function do_optimization() {
  if [[ ${opt_mode} == Auto_kernel_PGO ]]; then
    source ${afot_path}/auto_kernel_pgo.sh
  fi
  check_config_items
  init_profile_and_log
  load_script

  check_dependency
  prepare_env
  first_compilation
  execute_run_script
  detect_process
  profiling 

  prepare_new_env
  second_compilation
  is_opt_success
}

#执行入口，部分函数为加载不同优化脚本中得到
function main() {
  parse_input_params "$@"
  parse_config
  parse_input_params "$@"
  do_optimization
  exit "$?"
}

main "$@"
exit "$?"