#!/usr/bin/env bash
if [[ "${LOADED_ORTHOMAPPER_DBI_SH+1}" == "" ]]; then
    #------------------------------------------------------------------
    # Functions to handle download and install OrthoDB mapping data.
    # Main function dbi_find_and_install_node()
    # Data downloaded (or already present) should be in the directory
    # MAP_ORTHODB_DATA.
    # This variable is set in mapping environment config file (mapping_conf.sh).
    # 
    # OrthoDB mapping data are downloaded from
    #    
    #------------------------------------------------------------------

    # directory where download files are stored
    export DBI_DOWNLOAD=${DBI_DOWNLOAD:=""}

    # flag if the downloaded files are to be removed after install
    export DBI_REMOVE_DOWNLOAD=${DBI_REMOVE_DOWNLOAD:=1}

    # flag to use old FULL data set
    # 0 -> use subsampled data (no children files)
    # 1 -> use full data set with children files -> requires local subsampling
    export DBI_USE_FULL=${DBI_USE_FULL:=0}

    # default sub sampling parameters in downloaded data
    export DBI_SAMPLED_N=${DBI_SAMPLED_N:="8"}
    export DBI_SAMPLED_SCDIR=${DBI_SAMPLED_SCDIR:="ssc"}

    # postfix to be added to taxid part of filenames if subsampled data is used
    ###export DBI_SAMPLED_POSTFIX=${DBI_SAMPLED_POSTFIX:="_n8ssc"}


    #-------------------------------------------------------
    # Variables for URLs and files.
    # They are all set in dbi_set_node()
    # Ex, fetch node 9347, eukaryota for odb11v0
    #-------------------------------------------------------

    # The current node (9347)
    export DBI_NODE=""

    # flag if current node is VIRIDAE
    export DBI_NODE_IS_VIRIDAE=0

    # raw orthodb tar file prefix (euk)
    export DBI_ROOT_PREFIX=""

    # the tar file base name (euk_9347_mapping_odb11v0)
    export DBI_FILE_BASE=""

    # the file containing child nodes; (euk_9347_mapping_odb11v0.children)
    export DBI_FILE_CHILDREN=""

    # the tar file for the given node; (euk_9347_mapping_odb11v0.tgz)
    export DBI_FILE_TARFILE=""

    # tsv file with OG descriptions for the given node
    export DBI_FILE_OGDESC=""

    # URL for child file, constructed from URL_ORTHODB_API and child file
    # (https://data.orthodb.org/current/download/mapping/euk_9347_mapping_odb11v0.children)
    export DBI_URL_CHILDREN=""

    # URL for tar file, constructed like previous
    # (https://data.orthodb.org/current/download/mapping/euk_9347_mapping_odb11v0.tgz)
    export DBI_URL_TARFILE=""

    # URL for tar file, constructed like previous
    # (https://data.orthodb.org/current/download/mapping/og_descriptions_tsv/9347.tsv.gz)
    export DBI_URL_OGDESC=""

    #-------------------------------------------------------
    # readonly variables
    #-------------------------------------------------------
    # directory under MAP_ORTHODB_DATA containing the cluster description files
    declare -xr DBI_DESCRIPTIONS=og_descriptions

    #-------------------------------------------------------
    # functions
    #-------------------------------------------------------

    # check if node is a virus
    #dbi_node_is_viridae() {
    #    [[ $# -eq 1 ]] || return 1
    #    local node=$1
    #    [[ " ${MAP_NODES_VIRIDAE[@]} " =~ " $node " ]] && return 0
    #    return 2
    #}

    # get default label for subsampled taxids
    dbi_sampled_default_label() {
        map_subsample_label "$1" "$DBI_SAMPLED_N" "$DBI_SAMPLED_SCDIR"
    }

    # Make sure DBI_DOWNLOAD is set and is writeable
    dbi_check_download_dir() {
        if [[ -z "$DBI_DOWNLOAD" ]]; then
            # by default dir is /tmp/pipeline_(user)
            tmp_base DBI_DOWNLOAD 
            mkdir -p $DBI_DOWNLOAD
        fi
        [[ -d "$DBI_DOWNLOAD" ]] || { map_error_msg "DBI_DOWNLOAD is not readable at $DBI_DOWNLOAD"; return 255; }
        return 0
    }

    # Set all DBI related variables for the given node id
    # args: <prefix> <node>
    # where prefix is any of arc,bac,euk,vir and must be consistent with node (taxid)
    dbi_set_node() {
        local prefix=$1
        DBI_NODE=$2

        # file bases
        DBI_FILE_BASE="${prefix}_${DBI_NODE}_mapping_${MAP_ORTHODB_VERSION}"
        DBI_FILE_CHILDREN="${DBI_FILE_BASE}.children"
        DBI_FILE_TARFILE="${DBI_FILE_BASE}.tgz"
        DBI_FILE_OGDESC="${DBI_NODE}.tsv"

        # check if node is viridae -> this requires special treatment
        DBI_NODE_IS_VIRIDAE=0
        map_node_is_viridae $DBI_NODE && DBI_NODE_IS_VIRIDAE=1
        if (( $DBI_NODE_IS_VIRIDAE == 1 )); then
            # enforce FULL data set
            DBI_USE_FULL=1
            
            # there are no children
            DBI_URL_CHILDREN=""
            DBI_FILE_CHILDREN=""

            # download file should always be root
            # OGDESC file keeps the given nodeid
            DBI_FILE_BASE="${prefix}_${MAP_NODE_ROOT_VIRIDAE}_mapping_${MAP_ORTHODB_VERSION}"
            DBI_FILE_TARFILE="${DBI_FILE_BASE}.tgz"
        fi
        
        local url=""
        if (( DBI_USE_FULL == 1 )); then
            url="$URL_ORTHODB_API/download/mapping"
            (( $DBI_NODE_IS_VIRIDAE == 0 )) && DBI_URL_CHILDREN="$url/$DBI_FILE_CHILDREN"
        else
            url="$URL_ORTHODB_API/download/mapping/sampled"
            DBI_URL_CHILDREN=""
        fi
        DBI_URL_TARFILE="$url/$DBI_FILE_TARFILE"
        DBI_URL_OGDESC="$URL_ORTHODB_API/download/mapping/og_descriptions_tsv/${DBI_FILE_OGDESC}.gz"
        
        debug_msg "DBGURL: 1. $DBI_URL_CHILDREN"
        debug_msg "DBGURL: 2. $DBI_URL_TARFILE"
        debug_msg "DBGURL: 3. $DBI_URL_OGDESC"
    }

    # try get header to validate URL
    dbi_validate_url() {
        curl --head "$1"  2>&1 | grep 'HTTP/1.1 200 OK' >& /dev/null
    }

    # validate all URLs generated by dbi_set_node
    dbi_validate_node() {
        local silent=1
        [[ $# -eq 1 ]] && silent=0
        if (( DBI_USE_FULL == 1 )) && (( DBI_NODE_IS_VIRIDAE == 0 )); then
            dbi_validate_url "$DBI_URL_CHILDREN" || {
                (( silent == 1 )) || map_error_msg "invalid URL : $DBI_URL_CHILDREN"
                DBI_URL_CHILDREN=""; DBI_FILE_CHILDREN="";
                return 1
            }
        fi
        dbi_validate_url "$DBI_URL_TARFILE"  || {
            (( silent == 1 )) || map_error_msg "invalid URL : $DBI_URL_TARFILE"
            DBI_URL_CHILDREN=""; DBI_URL_TARFILE=""; DBI_FILE_CHILDREN=""; DBI_FILE_TARFILE=""
            return 2
        }
        dbi_validate_url "$DBI_URL_OGDESC"   || {
            (( silent == 1 )) || map_error_msg "invalid URL : $DBI_URL_OGDESC"
            DBI_URL_CHILDREN=""; DBI_URL_TARFILE=""; DBI_FILE_CHILDREN=""; DBI_FILE_TARFILE=""
            DBI_URL_OGDESC=""; DBI_FILE_OGDESC=""
            return 3
        }
        return 0
    }

    # make sure URLs are set, no check if they are valid
    dbi_url_node_are_set() {
        if (( DBI_USE_FULL == 1 )) && (( DBI_NODE_IS_VIRIDAE == 0 )); then
            [[ -n $DBI_URL_CHILDREN ]] || { map_error_msg "children URL is not set"; return 1; }
        fi
        [[ -n $DBI_URL_TARFILE  ]] || { map_error_msg "tarfile URL is not set";  return 2; }
        [[ -n $DBI_URL_OGDESC   ]] || { map_error_msg "ogdesc URL is not set";   return 3; }
        return 0
    }

    # check if map data for given node
    # checks for
    # 1. todo file <node>.todo
    # 2. subsampled todo file g<node>_<postfix>.todo
    dbi_map_exists() {
        # checks for existance of todo file
        local node=$1
        local clade="eukaryota"
        if [[ $DBI_ROOT_PREFIX == bac ]]; then
            clade="bacteria"
        elif [[ $DBI_ROOT_PREFIX == arc ]]; then
	        clade="archaea"
        elif [[ $DBI_ROOT_PREFIX == vir ]]; then
            clade="viridae"
        fi
        # check file!
        [[ -f $MAP_ORTHODB_DATA/$clade/todo/${node}.todo      ]] || return 1
        if (( DBI_USE_FULL == 0 )); then
            local label=$(map_subsample_label $node $DBI_SAMPLED_N $DBI_SAMPLED_SCDIR)
            [[ -f $MAP_ORTHODB_DATA/$clade/todo/${label}.todo ]] || return 1
        fi
        if (( DBI_USE_FULL == 1 )); then
            # full data set - need to check for an unsampled fasta file
            local odbtxid=$(grep -v '#' $MAP_ORTHODB_DATA/$clade/todo/${node}.todo | head -n 1)
            [[ -f $MAP_ORTHODB_DATA/$clade/Sequences/${odbtxid}.fs_selected ]] || return 1
        fi
        [[ -f $MAP_ORTHODB_DATA/$DBI_DESCRIPTIONS/${node}.tsv ]] || return 1
        return 0
    }

    dbi_download_node_check() {
        local fname="$1"
        local furl="$2"

        [[ -f $fname ]] && {
            local f=$(printf "%-60s skip - file exists" "$fname")
            map_info_msg "$f"
            return 0
        }
        local progopt="--silent"
        if test -t 1; then
            progopt=""
        fi
        curl -o "$fname" "$furl" $progopt || { rm -f "$fname"; map_error_msg "curl failed for URL <$furl>"; return 255; }
        return 0
    }

    dbi_download_children() {
        # skip if node is viridae
        (( DBI_NODE_IS_VIRIDAE == 0 )) || return 0
        # skip if not full set
        (( DBI_USE_FULL == 1 )) || return 0
        local saveNodeId=${DBI_NODE}
        local node=""
        while read node; do
            [[ ${node:0:1} == "#" ]] && continue
            dbi_set_node "$DBI_ROOT_PREFIX" "$node"
            dbi_validate_node verbose || return 255
            dbi_check_download_dir    || return 255
            # no need to download any children
            # the list of children includes ALL children
            dbi_download_node_check "$DBI_DOWNLOAD/$DBI_FILE_TARFILE"       "$DBI_URL_TARFILE"  || return 255
            dbi_download_node_check "$DBI_DOWNLOAD/${DBI_FILE_OGDESC}.gz"   "$DBI_URL_OGDESC"   || return 255
        done < "$DBI_DOWNLOAD/$DBI_FILE_CHILDREN"
        
        # reset to root node id
        dbi_set_node "$DBI_ROOT_PREFIX" "$saveNodeId"
    }

    dbi_download_node() {
        map_info_msg "download data files for node $DBI_NODE"
        if  (( DBI_NODE_IS_VIRIDAE == 1 )); then
            map_info_msg "requested node is viridae ($DBI_NODE) - will download FULL data for root node (10239)"
        fi
        
        # check that URL variables are set and that download target dir exists
        dbi_url_node_are_set   || return 255
        dbi_check_download_dir || return 255
        
        if (( DBI_USE_FULL == 1 )) && (( DBI_NODE_IS_VIRIDAE == 0 )); then
            dbi_download_node_check "$DBI_DOWNLOAD/$DBI_FILE_CHILDREN"      "$DBI_URL_CHILDREN" || return 255
        fi
        map_info_msg "download data from url: $DBI_URL_TARFILE"
        dbi_download_node_check "$DBI_DOWNLOAD/$DBI_FILE_TARFILE"       "$DBI_URL_TARFILE"  || return 255
        
        map_info_msg "download cluster annotations from url: $DBI_URL_OGDESC"
        dbi_download_node_check "$DBI_DOWNLOAD/${DBI_FILE_OGDESC}.gz"   "$DBI_URL_OGDESC"   || return 255
        
        if (( DBI_USE_FULL == 1 )) && (( DBI_NODE_IS_VIRIDAE == 0 )); then
            map_info_msg "download children"
            dbi_download_children || return 255
        fi
        return 0
    }

    dbi_untar() {
        local ftar="$DBI_DOWNLOAD/$DBI_FILE_TARFILE"
        [[ -f "$ftar"  ]]    || { map_error_msg "missing tarfile at $ftar"; return 255; }

        # untar mapping data into storage
        [[ -d $MAP_ORTHODB_DATA ]] || { map_error_msg "unable to cd to data directory at $MAP_ORTHODB_DATA"; return 255; }
        ( cd $MAP_ORTHODB_DATA && tar -xzf "$ftar" ) || { map_error_msg "unable to untar file $ftar into $MAP_ORTHODB_DATA"; return 255; }

        # take care of the OG description file
        # this file contains detailed info per OG at the given level
        # it's downloaded as a gzipped tsv file (eg 2.tsv.gz)
        local dirdesc=$MAP_ORTHODB_DATA/$DBI_DESCRIPTIONS
        mkdir -p $dirdesc  || { map_error_msg "unable to create directory $dirdesc"; return 255; }
        
        # decompress data into target file
        gunzip -c "$DBI_DOWNLOAD/${DBI_FILE_OGDESC}.gz" > "$dirdesc/$DBI_FILE_OGDESC"  || {
            map_error_msg "unable to unzip og desc file ${DBI_FILE_OGDESC}.gz"
            return 255
        }
        
        # possibly remove downloaded files
        (( DBI_REMOVE_DOWNLOAD == 1 )) && rm -f "$ftar" "$DBI_DOWNLOAD/${DBI_FILE_OGDESC}.gz"
        verbose_info_msg "installed node <$DBI_NODE> in $MAP_ORTHODB_DATA"
        return 0
    }

    # install current node (set by dbi_set_node())
    dbi_install_node() {
        local saveNodeId=${DBI_NODE}
        map_info_msg "installing node <$DBI_NODE>"
        dbi_check_download_dir || return 255
        dbi_untar              || return 255 


        if (( DBI_USE_FULL == 1 )) && (( DBI_NODE_IS_VIRIDAE == 0 )); then
            # install children
            local fchild="$DBI_DOWNLOAD/$DBI_FILE_CHILDREN"
            [[ -f "$fchild" ]] || {
                map_error_msg "missing children file $fchild"
                return 255
            }
            local cnode=""
            while read cnode; do
                [[ ${cnode:0:1} == "#" ]] && continue
                map_info_msg "installing child node $cnode"
                dbi_set_node "$DBI_ROOT_PREFIX" "$cnode"
                dbi_untar || return 255 
            done < "$fchild"
            (( DBI_REMOVE_DOWNLOAD == 1 )) && rm -f "$fchild"
        fi

        
        # restore DBI_NODE
        dbi_set_node "$DBI_ROOT_PREFIX" "$saveNodeId"
    }

    #------------------------------------------------------------------------------
    # find a node with a given ncbi taxid, download and install
    # 1. check if node is already installed in MAP_ORTHODB_DATA - return if so
    # 2. check if node file exists at MAP_ODBDATA_URL; if not, return with error
    # 3. download tar and children files (if not already done - DBI_DOWNLOAD)
    # 4. download tar files for children
    # 5. install all tar files
    #------------------------------------------------------------------------------
    dbi_find_and_install_node() {
        local node="$1"
        DBI_ROOT_PREFIX=""
        
        for DBI_ROOT_PREFIX in euk arc bac vir; do
            # check if taxid=node exists in given DBI
            # it will check for the existance of the todo file <node>.todo (and the subsampled one)
            dbi_map_exists "$node" && {
                # node has already been install
                map_info_msg "node is already installed : $node"
                return 0
            }
            # set DBI_FILE_XXX for the tgz and children file associated
            # set the corresponding URL's
            # will also set DBI_FILE_BASE
            dbi_set_node "$DBI_ROOT_PREFIX" "$node"

            # make sure the DBI_URL_XXX are valid
            dbi_validate_node && break
            DBI_FILE_BASE=""
        done
        
        if (( DBI_USE_FULL == 1 )); then
            map_info_msg "download node $1 using FULL data"
        else
            map_info_msg "download node $1 using SUBSAMPLED data"
        fi

        # check if a valid node was found
        if [[ -n "$DBI_FILE_BASE" ]]; then
            dbi_download_node || return 255
            dbi_install_node  || return 255
        else
            DBI_ROOT_PREFIX=""
            map_error_msg "did not find file for node $node"
            return 255
        fi
        return 0
    }
    LOADED_ORTHOMAPPER_DBI_SH=0
fi
