#!/usr/bin/env bash

set -efo pipefail

[[ ${DEBUG-} != true ]] || set -x

WORKSPACE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TN_BUILD_ENV_NAME=tn_build
readonly WORKSPACE_DIR TN_BUILD_ENV_NAME

export MAMBA_EXE=${HOME}/.local/bin/micromamba
export MAMBA_ROOT_PREFIX=${HOME}/micromamba

: "${NEED_PREPARE_ENV:=false}"
: "${NEED_ACTIVATE_ENV:=true}"

die() {
  local err=$? err_fmt=
  (( err )) && err_fmt=" (err=$err)" || err=1
  printf >&2 "[ERROR]$err_fmt %s\n" "$*"
  exit $err
}

_prepare_mamba_env(){
   if ! type micromamba >/dev/null 2>&1;then
     HTTPS_PROXY=${PROXY_URL:=${HTTPS_PROXY}} "${SHELL}" <(curl -L micro.mamba.pm/install.sh)
   fi
   _mamba_source
   [[ -z ${NEXUS3_HEADER} ]] || {
       ${MAMBA_EXE} config set --file "${MAMBA_ROOT_PREFIX}/.mambarc" channel_alias "${NEXUS3_HEADER}/conda"
   }
   micromamba create -y -f "${WORKSPACE_DIR}/config/${TN_BUILD_ENV_NAME}.yaml"
   micromamba activate "${TN_BUILD_ENV_NAME}"
}

_mamba_source() {
    [[ -e ${MAMBA_EXE} ]] || { echo "no micromamba exe found, run ./manager prepare_build_env to create env"; exit 1;}
    local __mamba_setup=''
    if __mamba_setup="$("$MAMBA_EXE" shell hook --shell bash --root-prefix "$MAMBA_ROOT_PREFIX" 2> /dev/null)"; then
        eval "$__mamba_setup"
    else
        alias micromamba='$MAMBA_EXE'  # Fallback on help from mamba activate
    fi
    unset __mamba_setup
}

_activate_env() {
    _mamba_source
    micromamba activate ${TN_BUILD_ENV_NAME}
}

_prepare_compile_env() {
    CUR_ENV_PATH=$(ompi_info --parsable --path prefix 2>/dev/null | awk -F":" '{print $NF}')
    export C_INCLUDE_PATH=${CUR_ENV_PATH}/include
    export CPLUS_INCLUDE_PATH=${CUR_ENV_PATH}/include
}

_build_config(){
    CUR_ENV_PATH=$(ompi_info --parsable --path prefix 2>/dev/null | awk -F":" '{print $NF}')
    cd -- "${WORKSPACE_DIR}"; bash configure.sh --openmpi_path "${CUR_ENV_PATH}"
    _prepare_compile_env
}

start_build(){
    [[ ${NEED_PREPARE_ENV} == true ]] && _prepare_mamba_env
    [[ ${NEED_ACTIVATE_ENV} == true ]] && _activate_env
    _build_config
    extra_opts=("$@")
    [[ ${DEBUG-} != true ]] || extra_opts+=(--sandbox_debug)
    bazel build "${extra_opts[@]}" -c opt //core:_pywrap_tn.so
}

only_build(){
    [[ ${NEED_ACTIVATE_ENV} == true ]] && _activate_env
    _prepare_compile_env
    extra_opts=("$@")
    [[ ${DEBUG-} != true ]] || extra_opts+=(--sandbox_debug)
    bazel build "${extra_opts[@]}" -c opt //core:_pywrap_tn.so
}


start_copy_libs(){
    rm -f tensornet/core/_pywrap_tn.so || true
    cp bazel-bin/core/_pywrap_tn.so tensornet/core/_pywrap_tn.so
}

start_test(){
    [[ ${NEED_ACTIVATE_ENV} == true ]] && _activate_env
    export PYTHONPATH=${WORKSPACE_DIR}:${PYTHONPATH}
    MPI_LIB_PATH=$(ompi_info --parsable --path prefix 2>/dev/null | awk -F":" '{print $NF}')
    export LD_LIBRARY_PATH=${MPI_LIB_PATH}/lib:${LD_LIBRARY_PATH}
    cd examples
    rm -rf data model || true
    python gen_example_data.py
    python main.py
}


start_only_upload(){
    [[ ${NEED_ACTIVATE_ENV} == true ]] && _activate_env
    export TWINE_USERNAME=${TWINE_USERNAME:=${NEXUS3_USERNAME}}
    export TWINE_PASSWORD=${TWINE_PASSWORD:=${NEXUS3_PASSWORD}}
    if [[ -z "$TWINE_USERNAME" || -z "$TWINE_PASSWORD" ]];then
        echo "need username/password auth, no env "
        echo "export NEXUS3_USERNAME=xxxx"
        echo "export NEXUS3_PASSWORD=xxxx"
        exit 0
    fi
    [[ -z ${NEXUS3_PYPI_HOST} ]] && { echo "need pypi host address, export NEXUS3_PYPI_HOST=xxx"; exit 0; }
    twine upload --verbose --repository-url "${NEXUS3_PYPI_HOST}" dist/*
}

start_create_dist(){
    [[ ${NEED_PREPARE_ENV} == true ]] && _prepare_mamba_env
    [[ ${NEED_ACTIVATE_ENV} == true ]] && _activate_env
    rm -rf dist/* || true
    PY_VERSION=$(python -c "import sys; print('cp' + ''.join(map(str, sys.version_info[:2])))")
    cat >setup.cfg <<-END
	[bdist_wheel]
	python_tag = $PY_VERSION
	END
    python3 -m build -w -n
    twine check dist/*
}

start_upload(){
    start_copy_libs
    start_create_dist
    start_only_upload
}

bump_dev_version() {
  hash bumpversion >/dev/null || die "cannot find bumpversion command"
  local mode=${1:-patch}
  case "$mode" in
  (major|minor|patch) bumpversion --commit "$mode" ;;
  (build)
    local release_part=''
    release_part=$(bumpversion --allow-dirty build --dry-run --list | awk -vFS=. '/^current_version=/ && $NF ~ "^[a-z]" { print $NF }')
    case "${release_part-}" in
    (dev*) bumpversion --commit "$mode" ;;
    (*) die "use '$0 release ...' to bump the build number of a release version." ;;
    esac
    ;;
  (*) die "Unknown mode ($mode) for bump_dev_version" ;;
  esac
}

bump_release_version() {
  hash bumpversion >/dev/null || die "cannot find bumpversion command"
  local mode=${1:-prod} release_part='' tag_name_option='' tool_deploy=false
  [[ $# -gt 1 && $2 == 'tool' ]] && tool_deploy=true
  release_part=$(bumpversion --allow-dirty build --dry-run --list | awk -vFS=. '/^current_version=/ && $NF ~ "^[a-z]" { print $NF }')

  case "$mode" in
  (rc)
    case "${release_part-}" in
    (dev*) bumpversion --commit --tag release ;;
    (rc*) bumpversion --commit --tag build ;;
    (''|post*) die "An production version cannot do a pre-release" ;;
    (*) die "Unknown release part ($release_part) for bump_release_version" ;;
    esac
    ;;
  (prod)
    case "${release_part-}" in
    (dev*)
      local release_version=''
      release_version=$(bumpversion --allow-dirty release --dry-run --list | grep '^current_version=' | cut -s -d = -f 2 | sed -E 's/\.[^0-9]+([0-9]*)$//')
      [[ $tool_deploy == true ]] && tag_name_option="--tag-name ${release_version}-tool"
      bumpversion --commit --tag ${tag_name_option} --new-version "$release_version" release # prod release, skip rc
      ;;
    (rc*)
      [[ $tool_deploy == true ]] && tag_name_option="--tag-name $(bumpversion --allow-dirty release --dry-run --list | grep '^new_version=' | cut -s -d = -f 2)-tool"
      bumpversion --commit --tag ${tag_name_option} release 
      ;;
    (''|post*) 
      [[ $tool_deploy == true ]] && tag_name_option="--tag-name $(bumpversion --allow-dirty build --dry-run --list | grep '^new_version=' | cut -s -d = -f 2)-tool"
      bumpversion --commit --tag ${tag_name_option} build 
      ;;
    (*) die "Unknown release part ($release_part) for bump_release_version" ;;
    esac
    ;;
  (*) die "Unknown mode ($mode) for bump_release_version" ;;
  esac
}

case "${1-}" in
(prepare_build_env)
  _prepare_mamba_env
  ;;
(build)
  shift 1
  start_build "$@"
  ;;
(only-build)
  shift 1
  only_build "$@"
  ;;
(deploy)
  shift 1
  start_upload
  ;;
(copy-libs)
  start_copy_libs
  ;;
(create_dist)
  shift 1
  start_create_dist
  ;;
(test)
  shift 1
  start_test "$@"
  ;;
(bump-version)
  shift 1
  bump_dev_version "$@"
  ;;
(release)
  shift 1
  bump_release_version "$@"
  ;;
(''|help)
  cmd=$(basename -- "$0")
  cat <<-END
	Usage:
	  $cmd [help] - Print this help.

	  $cmd prepare_build_env - install micromamba environment.

	  $cmd build [args..] - Build tn so file.

	  $cmd only-build [args..] - Build tn so file without config mpi

	  $cmd deploy - deploy tn to pypi

	  $cmd create_dist - create setup dist without upload

	  $cmd bump-version <major|minor|[patch]|build> - bump major/minor/patch/build version, always generates a *-dev version.

	  $cmd release <rc|[prod]> - generate a new (rc) release, and tag the new commit.
	END
  ;;
(*) die Unknown command "$1" ;;
esac