#!/usr/bin/env bash
set -Eeuo pipefail

## (1) Argument logic - template from https://betterdev.blog/minimal-safe-bash-script-template/
cleanup() {
  trap - SIGINT SIGTERM ERR EXIT
  [[ -d "${INSTALL_PATH}/download" ]] && [[ -n ${arg_keep_dl+x} ]] && [[ "$arg_keep_dl" == 0 ]] && rm -R ${INSTALL_PATH}/download
  echo "Thank you for installing AQUAMIS."
}

usage() {
  cat <<EOF
Usage: $(basename "${BASH_SOURCE[0]}") [OPTIONS]

Description:
This script completes the installation of the AQUAMIS pipeline. The openssl library is required for hashing the downloaded files.
For AQUAMIS installations via Bioconda use the option set [--databases] or
for AQUAMIS installations via Gitlab   use the option set [--mamba --databases].
For more information, please visit https://gitlab.com/bfr_bioinformatics/AQUAMIS

Options:
  -m, --mamba                Install the latest version of 'mamba' to the Conda base environment and
                             create the AQUAMIS environment from the git repository recipe
  -d, --databases            Download databases to ./AQUAMIS/download and extract them in e.g. ./AQUAMIS/reference_db
  -t, --test_data            Download test data to ./AQUAMIS/download and extract them in ./AQUAMIS/test_data
  -f, --force                Force overwrite for downloads in ${INSTALL_PATH}/download
  -k, --keep_downloads       Do not remove downloads after extraction
  -v, --verbose              Print script debug info
  -h, --help                 Show this help

EOF
  exit
}

msg() {
  echo >&2 -e "${1-}"
}

die() {
  local msg=$1
  local code=${2-1} # default exit status 1
  msg "$msg"
  exit "$code"
}

parse_params() {
  # default values of variables set from params
  param=''
  arg_mamba=0
  arg_databases=0
  arg_test_data=0
  arg_force=0
  arg_keep_dl=0

  while :; do
    case "${1-}" in
    -h | --help) usage ;;
    -v | --verbose) set -x ;;
    -m | --mamba) arg_mamba=1 ;;
    -d | --databases) arg_databases=1 ;;
    -t | --test_data) arg_test_data=1 ;;
    -f | --force) arg_force=1 ;;
    -k | --keep_downloads) arg_keep_dl=1 ;;
    -p | --param) # example named parameter
      param="${2-}"
      shift
      ;;
    -?*) die "Unknown option: $1" ;;
    *) break ;;
    esac
    shift
  done

  ## check required params and arguments
  # args=("$@")
  # [[ -z "${param-}" ]] && die "Please provide a parameter for the argument --param. Try $(basename "${BASH_SOURCE[0]}") --help."
  # [[ ${#args[@]} -eq 0 ]] && die "Missing script arguments. Try $(basename "${BASH_SOURCE[0]}") --help."
  # [[ $# -eq 0 ]] && die "Missing script arguments. Try $(basename "${BASH_SOURCE[0]}") --help."

  return 0
}

# Retrieve script paths
SCRIPT=$(readlink -f $0)
SCRIPT_PATH=$(dirname $SCRIPT)
INSTALL_PATH=$(dirname $SCRIPT_PATH)

# Setup cleanup-on-exit function
trap cleanup SIGINT SIGTERM ERR EXIT

# Parse arguments
echo "For help, type $(basename "${BASH_SOURCE[0]}") --help"
echo ""
parse_params "$@"

## (2) Script logic
print_variables() {
  echo "Checking CONDA installation... Detected"
  printf "Checking MAMBA installation... "
  [[ "$MAMBA" == 0 ]] && echo "Not Detected"
  [[ "$MAMBA" == 1 ]] && echo "Detected"
  printf "Checking APP source........... "
  [[ "$BIOCONDA" == 0 ]] && echo "Gitlab"
  [[ "$BIOCONDA" == 1 ]] && echo "Bioconda"
  echo ""
  echo "System variables:"
  echo "INSTALL_PATH = $INSTALL_PATH"
  echo "CONDA_BASE   = $CONDA_BASE"
  [[ -n "$CONDA_ENV_PATH" ]] && echo "ENV_AQUAMIS  = $CONDA_ENV_PATH"
  echo "ENV_NAME     = $CONDA_ENV_NAME"
  echo "ACTIVE_ENV   = $CONDA_ACTIVE"
}

check_conda() {
  CONDA_INFO=$(conda info --json)
  CONDA_BASE=$(echo "$CONDA_INFO" | grep -Po '(?<="root_prefix": ).*$' | tr -d '",\0')
  CONDA_ACTIVE=$(echo "$CONDA_INFO" | grep -Po '(?<="active_prefix_name": ).*$' | tr -d '",\0')
  CONDA_ENVS=$(echo "$CONDA_INFO" | grep -Pzo '(?s)(?<="envs": \[\n).*?(?=])' | tr -d '",\0' | sed 's/^[[:space:]]*//')
  [[ -x "$(command -v mamba)" ]] && MAMBA=1
  [[ "${SCRIPT_PATH##$CONDA_BASE}" != "$SCRIPT_PATH" ]] && BIOCONDA=1
  [[ "$BIOCONDA" == 0 ]] && CONDA_ENV_NAME=$(head -1 ${INSTALL_PATH}/envs/${CONDA_RECIPE} | cut -d' ' -f2)
  [[ "$BIOCONDA" == 1 ]] && CONDA_ENV_NAME=$(basename "$(dirname "$(dirname ${INSTALL_PATH})")")
  CONDA_ENV_PATH=$(echo "$CONDA_ENVS" | { grep -e "\/${CONDA_ENV_NAME}$" || true; })
  [[ "$CHECK" == 0 ]] && print_variables && CHECK=1
  echo ""  # this line always returns function exit 0
}

# Download Function
download_file() {
  download_success=0
  download_hash=''
  local target=$1
  local source=$2
  if [[ -f $target ]] && [[ $arg_force == 0 ]]; then
    echo "The file $target already exists. Skipping download."
  else
    echo "Downloading $source to $target"
    wget --output-document $target $source
    [[ $? -eq 0 ]] && [[ -s $target ]] && download_hash=$(openssl dgst -r -sha256 $target) && download_success=1
  fi
}

# Init Conda and Install Mamba
setup_mamba() {
  [[ "$BIOCONDA" == 0 ]] && echo "Installing the latest version of 'mamba' to the Conda base environment..."
  [[ "$BIOCONDA" == 1 ]] && die "ERROR: This is a Bioconda installation. There is no need to setup Mamba or an environment for AQUAMIS. Aborting..."

  source $CONDA_BASE/etc/profile.d/conda.sh
  [[ "$MAMBA" == 0 ]] && conda install --channel conda-forge mamba
  [[ "$MAMBA" == 1 ]] && echo "WARNING: Skipping Mamba installation. Mamba is already detected."

  # Create Conda Environment for AQUAMIS
  echo "Creating the AQUAMIS environment with the Conda recipe: ${INSTALL_PATH}/envs/${CONDA_RECIPE}"
  set +eu
  [[ -z "${CONDA_ENV_PATH}" ]] && { mamba env create -f ${INSTALL_PATH}/envs/${CONDA_RECIPE} || true; }
  set -eu
  [[ -n "${CONDA_ENV_PATH}" ]] && echo "WARNING: A Conda environment with the name $CONDA_ENV_NAME is already present. Skipping environment creation..."

  # set +eu  # workaround: see https://github.com/conda/conda/issues/8186
  # conda activate $CONDA_ENV_NAME
  # set -eu
  echo ""
}

# Download Databases
download_databases() {
  echo "Downloading databases to ${INSTALL_PATH}/download and extract them in their respective folder, e.g. ${INSTALL_PATH}/reference_db"

  [[ -z "${CONDA_ENV_PATH}" ]] && check_conda
  [[ -z "${CONDA_ENV_PATH}" ]] && die "ERROR: No Conda environment for AQUAMIS found to install the QUAST modules AUGUSTUS and BUSCO. Please run $(basename "${BASH_SOURCE[0]}") --mamba first. Aborting..."

  # Create Subdirectories
  mkdir -p ${INSTALL_PATH}/reference_db/mash
  mkdir -p ${INSTALL_PATH}/reference_db/taxonkit
  mkdir -p ${INSTALL_PATH}/download && cd ${INSTALL_PATH}/download

  # Download files
  [[ ! -d "${CONDA_ENV_PATH}/lib/python3.7/site-packages/quast_libs/" ]] && die "ERROR: The Conda environment installation is incomplete. QUAST was not installed. Aborting..."
  download_file "augustus.tar.gz" "https://seafile.bfr.berlin/f/64cc5034fad74f50a2f0/?dl=1" # 139MB or https://databay.bfrlab.de/f/fa13c9eb2625477eb729/?dl=1
  [[ -n $download_hash ]] && echo "$download_hash" >> reference_db.sha256
  [[ "$download_success" == 1 ]] && tar -xzv -f augustus.tar.gz -C ${CONDA_ENV_PATH}/lib/python3.7/site-packages/quast_libs/

  download_file "bacteria.tar.gz" "https://seafile.bfr.berlin/f/41cf8fdcfe2043d2800e/?dl=1" # 8.8MB or https://busco.ezlab.org/v2/datasets/bacteria_odb9.tar.gz
  [[ -n $download_hash ]] && echo "$download_hash" >> reference_db.sha256
  [[ "$download_success" == 1 ]] && tar -xzv -f bacteria.tar.gz -C ${CONDA_ENV_PATH}/lib/python3.7/site-packages/quast_libs/busco/

  download_file "confindr_db.tar.gz" "https://seafile.bfr.berlin/f/ede87ec860624a0cb406/?dl=1" # 153MB for the old db incl. some rMLST db's e.g. Enterococcus and Brucella, use https://seafile.bfr.berlin/f/47cae689eda7440c83bb/?dl=1
  [[ -n $download_hash ]] && echo "$download_hash" >> reference_db.sha256
  [[ "$download_success" == 1 ]] && tar -xzv -f confindr_db.tar.gz -C ${INSTALL_PATH}/reference_db/

  download_file "mashDB.tar.gz" "https://seafile.bfr.berlin/f/41f804a1eba541788530/?dl=1" # 216MB or https://databay.bfrlab.de/f/34b8c88945a8439dac64/?dl=1
  [[ -n $download_hash ]] && echo "$download_hash" >> reference_db.sha256
  [[ "$download_success" == 1 ]] && tar -xzv -f mashDB.tar.gz -C ${INSTALL_PATH}/reference_db/mash/

  download_file "minikraken2.tar.gz" "https://seafile.bfr.berlin/f/8ca1b4d2c97341498698/?dl=1" # 6.0GB or ftp://ftp.ccb.jhu.edu/pub/data/kraken2_dbs/minikraken_8GB_202003.tgz
  [[ -n $download_hash ]] && echo "$download_hash" >> reference_db.sha256
  [[ "$download_success" == 1 ]] && tar -xzv -f minikraken2.tar.gz -C ${INSTALL_PATH}/reference_db/

  download_file "taxdump.tar.gz" "https://seafile.bfr.berlin/f/1d51700ecfd241e4a6d4/?dl=1" #  54MB or ftp://ftp.ncbi.nih.gov/pub/taxonomy/taxdump.tar.gz
  [[ -n $download_hash ]] && echo "$download_hash" >> reference_db.sha256
  [[ "$download_success" == 1 ]] && tar -xzv -f taxdump.tar.gz -C ${INSTALL_PATH}/reference_db/taxonkit/

  # CleanUp Downloads
  [[ "$arg_keep_dl" == 0 ]] && mv ${INSTALL_PATH}/download/reference_db.sha256 ${INSTALL_PATH}/reference_db/reference_db.sha256
  echo "Download finished."
}

# Download Test Data
download_test_data() {
  echo "Downloading test data to ${INSTALL_PATH}/download and extract them in their respective folder, e.g. ${INSTALL_PATH}/test_data"

  # Create Subdirectories
  mkdir -p ${INSTALL_PATH}/test_data/fastq
  mkdir -p ${INSTALL_PATH}/download && cd ${INSTALL_PATH}/download

  # Download files
  download_file "test_data.tar.gz" "https://seafile.bfr.berlin/f/b8b636dbe6bd4b39801c/?dl=1" # 1.1GB

  # Verify Integrity of Downloads
  [[ -n $download_hash ]] && echo "$download_hash" >test_data.sha256

  # Unpack Downloads
  [[ "$download_success" == 1 ]] && tar -xzv -f test_data.tar.gz -C ${INSTALL_PATH}/test_data/fastq

  # Create Sample Sheet for Test Data
  [[ "$download_success" == 1 ]] && bash ${SCRIPT_PATH}/create_sampleSheet.sh --mode ncbi --fastxDir ${INSTALL_PATH}/test_data/fastq --outDir ${INSTALL_PATH}/test_data

  # CleanUp Downloads
  [[ "$arg_keep_dl" == 0 ]] && mv ${INSTALL_PATH}/download/test_data.sha256 ${INSTALL_PATH}/test_data/test_data.sha256
  echo "Download finished."
}

## Set Variables
CONDA_RECIPE="aquamis.yaml"
CHECK=0
CONDA=0
MAMBA=0
BIOCONDA=0

## Sanity Check
[[ $(basename $SCRIPT_PATH) != "scripts" ]] && die "ERROR: This setup script does not reside in its original installation directory. This is a requirement for proper execution. Aborting..."

## Check Conda
[[ -x "$(command -v conda)" ]] && CONDA=1
[[ "$CONDA" == 0 ]] && die "ERROR: Conda was not detected."
[[ "$CONDA" == 1 ]] && check_conda

## Execute Option Modules
[[ "$arg_mamba" == 1 ]] && setup_mamba
[[ "$arg_databases" == 1 ]] && download_databases
[[ "$arg_test_data" == 1 ]] && download_test_data
