#!/bin/bash

## [Goal] Install mamba, conda env, complement conda env, download databases, download test data for AQUAMIS
## [Author] Holger Brendebach

## Shell Behaviour --------------------------------------------------
set -Euo pipefail # to exit immediately, # ADD e
trap cleanup SIGINT SIGTERM ERR EXIT

## Source Helper Functions --------------------------------------------------
script_real=$(realpath -P "${BASH_SOURCE[0]}")
script_path=$(dirname $script_real)
script_name=$(basename $script_real)
repo_path=$(dirname $script_path)
[[ ! -f "$script_path/helper_functions.sh" ]] && echo "ERROR: File not found: $script_path/helper_functions.sh" && exit 1
source "$script_path/helper_functions.sh"

## (1) Argument Logic

usage() {
  cat <<- USE

${bold}Basic usage:${normal}
bash $script_real [OPTIONS]

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

${bold}Options:${normal}
  -m, --mamba                Install the latest version of 'mamba' to the Conda base environment and
                             create the AQUAMIS environment from the git repository recipe
  -b, --busco                Download Augustus and BUSCO databases to <AQUAMIS>/download and extract them in the conda environment
  -d, --databases            Download databases to <AQUAMIS>/download and extract them in <AQUAMIS>/reference_db
  -t, --test_data            Download test data to <AQUAMIS>/download and extract them in <AQUAMIS>/test_data
  -s, --status               Show installation status
  -f, --force                Force overwrite for downloads in <AQUAMIS>/download
  -k, --keep_downloads       Do not remove downloads after extraction
  -v, --verbose              Print script debug info
  -h, --help                 Show this help

USE
  exit 0
}

cleanup() {
  trap - SIGINT SIGTERM ERR EXIT
  [[ -d "$repo_path/download" ]] && [[ -n ${arg_keep_dl+x} ]] && [[ "$arg_keep_dl" == false ]] && rm -R $repo_path/download
}

parse_args() {
  local args=("$@")

  ## default values of variables
  conda_recipe="aquamis.yaml"

  ## default values of switches
  conda_status=false
  mamba_status=false
  bioconda_status=false

  arg_mamba=false
  arg_busco=false
  arg_databases=false
  arg_test_data=false
  arg_status=false
  arg_force=false
  arg_keep_dl=false

  while :; do
    case "${1-}" in
      -h | --help) usage ;;
      -v | --verbose)
        set -x
        shift
        ;;
      -m | --mamba)
        arg_mamba=true
        shift
        ;;
      -b | --busco)
        arg_busco=true
        shift
        ;;
      -d | --databases)
        arg_databases=true
        shift
        ;;
      -t | --test_data)
        arg_test_data=true
        shift
        ;;
      -s | --status)
        arg_status=true
        shift
        ;;
      -f | --force)
        arg_force=true
        shift
        ;;
      -k | --keep_downloads)
        arg_keep_dl=true
        shift
        ;;
      -?*) die "Unknown option: $1. For the manual, type ${blue}bash $script_real --help${normal}" ;;
      *) break ;;
    esac
  done

  ## check required arguments
  [[ ${#args[@]} -eq 0 ]] && die "Missing script arguments. For the manual, type ${blue}bash $script_real --help${normal}"

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

  return 0
}

## Parse Arguments --------------------------------------------------
parse_args "$@"

## (2) Script logic

## Info --------------------------------------------------

show_info() {
  logheader "Parameters:"
  echo "
repo_path:|$repo_path
" | column -t -s "|"
}

## Modules --------------------------------------------------

download_file() {
  # Init
  download_success=false
  download_status=false
  download_hash=''
  
  # Download
  if [[ -f $local_archive ]] && [[ $arg_force == false ]]; then
    logwarn "The file $local_archive already exists. Skipping download."
  else
    logentry "Downloading $remote_archive to $local_archive"
    wget --ca-certificate=$repo_path/resources/seafile-bfr-berlin_certificate.pem --output-document $local_archive $remote_archive && download_status=true
    [[ "$download_status" = true ]] && [[ -s $local_archive ]] && download_hash=$(openssl dgst -r -sha256 $local_archive) && download_success=true
  fi
  
  # Verify Integrity of Download
  [[ -n $download_hash ]] && echo "$download_hash" >> $hashfile

  # Unpack Downloads
  [[ "$download_success" == true ]] && logentry "Extracting archive $(basename $local_archive) to $extraction_directory"
  [[ "$download_success" == true ]] && tar -xzv -f $local_archive -C $extraction_directory
}

setup_mamba() {
  # checks
  [[ "$bioconda_status" == false ]] && logentry "Installing the latest version of \"mamba\" to the Conda base environment."
  [[ "$bioconda_status" == true ]] && die "This is a Bioconda installation. There is no need to setup Mamba or an environment for AQUAMIS."

  source $conda_base/etc/profile.d/conda.sh
  [[ "$mamba_status" == false ]] && conda install --channel conda-forge mamba
  [[ "$mamba_status" == true ]] && logwarn "Skipping Mamba installation. Mamba is already detected."

  # Create Conda Environment for AQUAMIS
  logentry "Creating the AQUAMIS environment with the Conda recipe: $repo_path/envs/${conda_recipe}"
  set +eu  # workaround: see https://github.com/conda/conda/issues/8186
  [[ -z "${conda_env_path}" ]] && { mamba env create -f $repo_path/envs/${conda_recipe} || true; }
  [[ -n "${conda_env_path}" && "$arg_force" == true && ! -d "${conda_base}/envs/${conda_env_name}" ]] && { mamba env create -f $repo_path/envs/${conda_recipe} -p ${conda_base}/envs/${conda_env_name} || true; }  # corner case: there is already a conda env with the same name in another conda_base, e.g. NGSAdmin
  [[ -n "${conda_env_path}" ]] && logwarn "A Conda environment with the name \"$conda_env_name\" is already present. Skipping environment creation."
  # conda activate $conda_env_name
  set -eu

  logentry "${green}Environment installation completed.${normal}"
}

complete_busco() {
  # Checks
  [[ -z "${conda_env_path}" ]] && check_conda
  [[ -z "${conda_env_path}" ]] && die "No Conda environment for AQUAMIS found to install the QUAST modules AUGUSTUS and BUSCO. Please run $(basename "${BASH_SOURCE[0]}") --mamba first. Aborting..."
  [[ ! -d "$conda_env_path/lib/python3.7/site-packages/quast_libs" ]] && die "The Conda environment installation is incomplete. QUAST was not installed. Aborting..."

  # Create Subdirectories
  [[ ! -d "$repo_path/download" ]] && mkdir -p $repo_path/download

  # Download files
  remote_archive="https://gitlab.bfr.berlin/bfr_bioinformatics/aquamis_databases/-/raw/main/augustus.tar.gz" # 139MB or https://databay.bfrlab.de/f/fa13c9eb2625477eb729/?dl=1
  local_archive="$repo_path/download/augustus.tar.gz"
  hashfile="$repo_path/download/reference_db.sha256"
  extraction_directory="$conda_env_path/lib/python3.7/site-packages/quast_libs/"
  download_file    
  
  remote_archive="https://gitlab.bfr.berlin/bfr_bioinformatics/aquamis_databases/-/raw/main/bacteria.tar.gz" # 8.8MB or https://busco.ezlab.org/v2/datasets/bacteria_odb9.tar.gz
  local_archive="$repo_path/download/bacteria.tar.gz"
  hashfile="$repo_path/download/reference_db.sha256"
  extraction_directory="$conda_env_path/lib/python3.7/site-packages/quast_libs/busco/"
  download_file    

  # CleanUp Downloads
  [[ "$arg_keep_dl" == false ]] && cat $repo_path/download/reference_db.sha256 >> $repo_path/reference_db/reference_db.sha256

  logentry "${green}BUSCO installation completed.${normal}"
}

download_databases() {
  # Create Subdirectories
  [[ ! -d "$repo_path/download" ]] && mkdir -p $repo_path/download
  mkdir -p $repo_path/reference_db/mash/genomes
  mkdir -p $repo_path/reference_db/taxonkit

  # Download files
  remote_archive="https://gitlab.bfr.berlin/bfr_bioinformatics/aquamis_databases/-/raw/main/confindr_db.tar.gz" # 153MB for the old db incl. some rMLST db's e.g. Enterococcus and Brucella, use https://seafile.bfr.berlin/f/47cae689eda7440c83bb/?dl=1
  local_archive="$repo_path/download/confindr_db.tar.gz"
  hashfile="$repo_path/download/reference_db.sha256"
  extraction_directory="$repo_path/reference_db/"
  download_file

  remote_archive="https://gitlab.bfr.berlin/bfr_bioinformatics/aquamis_databases/-/raw/main/mashDB.tar.gz" # 216MB or https://databay.bfrlab.de/f/34b8c88945a8439dac64/?dl=1
  local_archive="$repo_path/download/mashDB.tar.gz"
  hashfile="$repo_path/download/reference_db.sha256"
  extraction_directory="$repo_path/reference_db/mash/"
  download_file

  remote_archive="https://gitlab.bfr.berlin/bfr_bioinformatics/aquamis_databases/-/raw/main/minikraken2.tar.gz" # 6.0GB or ftp://ftp.ccb.jhu.edu/pub/data/kraken2_dbs/minikraken_8GB_202003.tgz or https://genome-idx.s3.amazonaws.com/kraken/minikraken_8GB_202003.tgz
  local_archive="$repo_path/download/minikraken2.tar.gz"
  hashfile="$repo_path/download/reference_db.sha256"
  extraction_directory="$repo_path/reference_db/"
  download_file

  remote_archive="https://gitlab.bfr.berlin/bfr_bioinformatics/aquamis_databases/-/raw/main/taxdump.tar.gz" #  54MB or ftp://ftp.ncbi.nih.gov/pub/taxonomy/taxdump.tar.gz
  local_archive="$repo_path/download/taxdump.tar.gz"
  hashfile="$repo_path/download/reference_db.sha256"
  extraction_directory="$repo_path/reference_db/taxonkit/"
  download_file

  # CleanUp Downloads
  [[ "$arg_keep_dl" == false ]] && cat $repo_path/download/reference_db.sha256 >> $repo_path/reference_db/reference_db.sha256

  logentry "${green}Database download finished.${normal}"
}

download_test_data() {
  # Create Subdirectories
  [[ ! -d "$repo_path/download" ]] && mkdir -p $repo_path/download
  [[ ! -d "$repo_path/test_data/fastq" ]] && mkdir -p $repo_path/test_data/fastq

  # Download files
  remote_archive="https://gitlab.bfr.berlin/bfr_bioinformatics/aquamis_databases/-/raw/main/test_data.tar.gz" # 1.1GB
  local_archive="$repo_path/download/test_data.tar.gz"
  hashfile="$repo_path/download/test_data.sha256"
  extraction_directory="$repo_path/test_data/fastq/"
  download_file

  # Create Sample Sheet for Test Data
  [[ "$download_success" == true ]] && logentry "Creating test data sample sheet $repo_path/test_data/samples.tsv"
  [[ "$download_success" == true ]] && bash ${script_path}/create_sampleSheet.sh --mode ncbi --fastxDir $repo_path/test_data/fastq --outDir $repo_path/test_data
  echo

  # CleanUp Downloads
  [[ "$arg_keep_dl" == false ]] && cat $repo_path/download/test_data.sha256 >> $repo_path/test_data/test_data.sha256

  logentry "${green}Test data download finished.${normal}"
}

check_success() {
  [[ -f $repo_path/reference_db/confindr/refseq.msh ]] && status_confindr="${green}OK${normal}"
  [[ -f $repo_path/reference_db/kraken/hash.k2d ]] && status_kraken="${green}OK${normal}"
  [[ -f $repo_path/reference_db/mash/mashDB.msh ]] && status_mash="${green}OK${normal}"
  [[ -f $repo_path/reference_db/taxonkit/names.dmp ]] && status_taxonkit="${green}OK${normal}"
  [[ -f ${conda_env_path:-$CONDA_PREFIX}/lib/python3.7/site-packages/quast_libs/busco/busco.py ]] && status_busco="${green}OK${normal}"
  [[ -f ${conda_env_path:-$CONDA_PREFIX}/lib/python3.7/site-packages/quast_libs/augustus3.2.3/bin/augustus ]] && status_augustus="${green}OK${normal}"
  logheader "Database Status:"
  echo "
status_confindr:|${status_confindr:-"${red}FAIL${normal}"}
status_kraken:|${status_kraken:-"${red}FAIL${normal}"}
status_mash:|${status_mash:-"${red}FAIL${normal}"}
status_taxonkit:|${status_taxonkit:-"${red}FAIL${normal}"}
status_busco:|${status_busco:-"${red}FAIL${normal}"}
status_augustus:|${status_augustus:-"${red}FAIL${normal}"}
" | column -t -s "|"
  echo
}

## (3) Main

logentry "[START] $script_real" # from helper_functions.sh
logentry "Arguments: ${*:-"No arguments provided"}"

## Workflow
conda_recipe_env_name=$(head -n1 $repo_path/envs/${conda_recipe} | cut -d' ' -f2)
check_conda $conda_recipe_env_name
[[ "$arg_mamba" == true ]] && setup_mamba
[[ "$arg_busco" == true ]] && complete_busco
[[ "$arg_databases" == true ]] && download_databases
[[ "$arg_test_data" == true ]] && download_test_data
[[ "$arg_status" == false ]] && check_conda $conda_recipe_env_name  # check changes to conda
show_info
check_success

logentry "[STOP] $script_real"  # from helper_functions.sh
echo "Thank you for installing AQUAMIS."
