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

## Argument logic - template from https://betterdev.blog/minimal-safe-bash-script-template/
cleanup() {
  trap - SIGINT SIGTERM ERR EXIT
  [[ -d "${AQUAMIS_PATH}/download" ]] && rm -R ${AQUAMIS_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 --keep_downloads] or
for AQUAMIS installations via Gitlab   use the option set [--conda --databases --keep_downloads]
For more information, please visit https://gitlab.com/bfr_bioinformatics/AQUAMIS

Options:
  -c, --conda                Install the latest version of 'mamba' to the conda base environment and
                             create the AQUAMIS environment '$CONDA_ENV_NAME'
  -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 ${AQUAMIS_PATH}/download
  -k, --keep_downloads       Do not remove downloads after extraction
  -v, --verbose              Print script debug info
  -h, --help                 Show this help

System paths detected:
AQUAMIS    = $AQUAMIS_PATH
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_conda=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 ;;
    -c | --conda) arg_conda=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

  args=("$@")

  # check required params and arguments
  # [[ -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."

  return 0
}

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

# Retrieve paths
SCRIPT=$(readlink -f $0)
SCRIPT_PATH=$(dirname $SCRIPT)
AQUAMIS_PATH=$(dirname $SCRIPT_PATH)
CONDA_ENV_NAME=$(head -1 ${AQUAMIS_PATH}/envs/aquamis.yaml | cut -d' ' -f2)

# Parse arguments
parse_params "$@"

## Script logic

# Download Function
dowload_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_conda() {
  echo "Installing the latest version of 'mamba' to the conda base environment and creating the AQUAMIS environment with the recipe: ${AQUAMIS_PATH}/envs/aquamis.yaml"

  CONDA_BASE=$(conda info --base)
  source $CONDA_BASE/etc/profile.d/conda.sh
  ! [[ -x "$(command -v mamba)" ]]  && conda install --channel conda-forge mamba

  # Create Conda Environment for AQUAMIS
  mamba env create -f ${AQUAMIS_PATH}/envs/aquamis.yaml
  conda activate $CONDA_ENV_NAME
}

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

  CONDA_BASE=$(conda info --base)

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

  # Download files
  dowload_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_BASE}/envs/${CONDA_ENV_NAME}/lib/python3.7/site-packages/quast_libs/

  dowload_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_BASE}/envs/${CONDA_ENV_NAME}/lib/python3.7/site-packages/quast_libs/busco/

  dowload_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 ${AQUAMIS_PATH}/reference_db/

  dowload_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 ${AQUAMIS_PATH}/reference_db/mash/

  dowload_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 ${AQUAMIS_PATH}/reference_db/

  dowload_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 ${AQUAMIS_PATH}/reference_db/taxonkit/

  # CleanUp Downloads
  [[ "$arg_keep_dl" == 0 ]] && mv ${AQUAMIS_PATH}/download/reference_db.sha256 ${AQUAMIS_PATH}/reference_db/reference_db.sha256
}

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

  # Create Subdirectories
  mkdir -p ${AQUAMIS_PATH}/test_data/raw
  mkdir -p ${AQUAMIS_PATH}/download && cd ${AQUAMIS_PATH}/download

  # Download files
  dowload_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 ${AQUAMIS_PATH}/test_data/fastq

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

  # CleanUp Downloads
  [[ "$arg_keep_dl" == 0 ]] && mv ${AQUAMIS_PATH}/download/test_data.sha256 ${AQUAMIS_PATH}/reference_db/test_data.sha256
}

## Execute Modules
[[ "$arg_conda" == 1 ]] && setup_conda
[[ "$arg_databases" == 1 ]] && download_databases
[[ "$arg_test_data" == 1 ]] && download_test_data
[[ $# -gt 0 ]] && [[ "$arg_keep_dl" == 0 ]] && cleanup
