#!/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
  echo "Cleanup."
}

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

Description:
This script downloads MASH references that are not provided by the "offline MASH db".
Activation of the associated AQUAMIS conda env is a prerequisite.
For more information, please visit https://gitlab.com/bfr_bioinformatics/AQUAMIS

Options:
  -a, --accession            String of a best MASH ID
  -m, --mashDBpath           Path to the MASH db directory, e.g. <AQUAMIS>/reference_db/mash/genomes
  -n, --ncbiDB               Choose NCBI reference database, either "nuccore" or "refseq", default is "nuccore"
  -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
  accession=''
  mashDBpath=''
  ncbiDB='nuccore'

  while :; do
    case "${1-}" in
    -h | --help) usage ;;
    -v | --verbose) set -x ;;
    -a | --accession)
      accession="$2"
      shift 2
      ;;
    -m | --mashDBpath)
      mashDBpath="$2"
      shift 2
      ;;
    -n | --ncbiDB)
      ncbiDB="$2"
      shift 2
      ;;
    -?*) die "Unknown option: $1" ;;
    *) break ;;
    esac
  done

  ## check required params and arguments
  [[ -z "${accession-}" ]] && die "Please provide a parameter for the argument --accession. Try $(basename "${BASH_SOURCE[0]}") --help."
  [[ -z "${mashDBpath-}" ]] && die "Please provide a parameter for the argument --mashDBpath. 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

## 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..."

## Find Accession
echo "Using reference repository: $mashDBpath"
reference_dict=${mashDBpath}/accession_dict.tsv
[[ ! -d "${mashDBpath}" ]] && die "The $mashDBpath is not valid."
[[ ! -s "$reference_dict" ]] && die "The $reference_dict is not valid."

## retrieve accession id
echo "Best accession is: $accession"

## query reference repository via dictionary
grep -q "^$accession" "$reference_dict" && die "The reference is already present in the $reference_dict"
grep -q "^$accession\tawaiting_ncbi_retrieval" "$reference_dict" && die "The reference is currently downloaded by another job."

echo "The fasta file is not present in the reference repository. Proceeding with download from NCBI."
echo -e "$accession\tawaiting_ncbi_retrieval" >> "$reference_dict"  # this adds an entry to signal a pending process and log incomplete online queries if a fault ari

## retrieval of ftp path
fna_ftp_path=$(esearch -db assembly -query "${accession}" | esummary | xtract -pattern DocumentSummary -element FtpPath_RefSeq | awk 'BEGIN{FS="/"} END{print $0"/"$NF"_genomic.fna.gz"}')
fna_ftp_file=$(basename $fna_ftp_path)
gff_ftp_file=${fna_ftp_file/.fna.gz/.gff.gz}  # presence of gff file at NCBI is a prerequisite
gff_ftp_path=$(dirname $fna_ftp_path)/$gff_ftp_file

## download with wget to reference repository
echo "Using NCBI DB: $ncbiDB"
if [[ "$ncbiDB" == "nuccore" ]] ; then  # IMPORTANT: these "nuccore" files are single-fasta/gff3 and may only contain a single chromosome of multi-chromosome species
    echo "wget -O - https://www.ncbi.nlm.nih.gov/sviewer/viewer.cgi?db=nuccore\&report=fasta\&id=$accession --tries 10 | gzip > ${mashDBpath}/$fna_ftp_file"
          wget -O - https://www.ncbi.nlm.nih.gov/sviewer/viewer.cgi?db=nuccore\&report=fasta\&id=$accession --tries 10 | gzip > ${mashDBpath}/$fna_ftp_file
    echo "wget -O - https://www.ncbi.nlm.nih.gov/sviewer/viewer.cgi?db=nuccore\&report=gff3\&id=$accession  --tries 10 | gzip > ${mashDBpath}/$gff_ftp_file"
          wget -O - https://www.ncbi.nlm.nih.gov/sviewer/viewer.cgi?db=nuccore\&report=gff3\&id=$accession  --tries 10 | gzip > ${mashDBpath}/$gff_ftp_file
elif [[ "$ncbiDB" == "refseq" ]] ; then # IMPORTANT: these "assembly|refseq" files may include multiple chromosomes and associated plasmids in a multi-fasta/gff3 format
    echo "wget -O ${mashDBpath}/$fna_ftp_file $fna_ftp_path --tries 10"
          wget -O ${mashDBpath}/$fna_ftp_file $fna_ftp_path --tries 10
    echo "wget -O ${mashDBpath}/$gff_ftp_file $gff_ftp_path --tries 10"
          wget -O ${mashDBpath}/$gff_ftp_file $gff_ftp_path --tries 10
fi
awk -v accession=$accession -v ftpfile=$fna_ftp_file 'BEGIN{FS=OFS="\t"} {if ($1 == accession) $2=ftpfile;} 1' "$reference_dict" > tmp.tsv && mv -f tmp.tsv "$reference_dict"
