#!/usr/bin/env bash
#
# orthologer - Pipeline for deducing orthology from fasta files.
# Copyright (C) 2005-2023 E. Zdobnov lab: F. Tegenfeldt
# <fredrik.tegenfeldt@unige.ch> based on code by E. Kriventseva
#
# This file is part of orthologer.
#
# orthologer is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# orthologer is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with orthologer.  If not, see <http://www.gnu.org/licenses/>.
#
set +u
CMD_MANAGE=${CMD_MANAGE:=$(basename $0)}

# do not allow unset variables
set -u

# check pipeline root directory - should be set by caller
#TODO: DIR_ORTHOLOGER and DIR_ORTHOLOGER should be different dirs and
#      common_odb.sh should be named just config

# set orthoDB pipeline dir, ie the dir of this script
export DIR_ORTHOLOGER=${DIR_ORTHOLOGER="$(cd -P "$(dirname $(readlink -f "${BASH_SOURCE[0]}"))"/../ && pwd)"}
[[ -d "$DIR_ORTHOLOGER" ]] || { echo "FATAL: unable to determine DIR_ORTHOLOGER! Got <$DIR_ORTHOLOGER>."; exit 255; }
export DIR_ORTHOLOGER="$DIR_ORTHOLOGER"

#--------------------------------------------------------------------------
# Check prerequisites (bash version)
source $DIR_ORTHOLOGER/bin/check_prerequisites.sh

# Check user id - possibly fail if root
source $DIR_ORTHOLOGER/bin/check_runas_root.sh
#--------------------------------------------------------------------------

# project directory
export DIR_PROJECT=${DIR_PROJECT:="$(pwd)"}
[[ -d "$DIR_PROJECT" ]] || { echo "FATAL: no valid DIR_PROJECT! Got <$DIR_PROJECT>."; exit 255; }

if [[ "${LOADED_COMMON_PIPELINE_SH+1}" == "" ]]; then
    if [[ -f orthologer_conf.sh ]]; then
        ##echo "loading local orthologer_conf.sh"
        source orthologer_conf.sh
        debug_msg "loaded local orthologer_conf.sh"
    else
        echo "loading system common_odb.sh"
        [[ -f "$DIR_ORTHOLOGER/bin/common_odb.sh" ]] || { echo "FATAL: unable to read DIR_ORTHOLOGER/bin/common_odb.sh!"; exit 255; }
        source $DIR_ORTHOLOGER/bin/common_odb.sh
        debug_msg "${CMD_MANAGE}: loaded common_odb.sh"
    fi
else
    debug_msg "everything is loaded!"
    source $DIR_ORTHOLOGER/bin/declare_arrays.sh
    op_step_setup
fi

# set API
set_orthodb_api

# get clade defs
odblevels="$DIR_ORTHOLOGER/etc/$URL_ORTHODB_API_VERSION/orthomapper_levels.sh"
[[ -f "$odblevels" ]] || {
    echo "ERROR: orthomapper_levels.sh not found! Is DIR_ORTHOLOGER correctly set?" 1>&2
    echo "Here: $odblevels"
    exit 255
}
source "$odblevels"

# get all orthologer/orthomapper functions
[[ -f "$DIR_ORTHOLOGER/bin/orthomapper_func.sh" ]] || { echo "ERROR: orthomapper_func.sh not found! Is DIR_ORTHOLOGER correctly set?" 1>&2; exit 255; }
source "$DIR_ORTHOLOGER/bin/orthomapper_func.sh"
source "$DIR_ORTHOLOGER/bin/orthomapper_dbi.sh"
source "$DIR_ORTHOLOGER/bin/orthomapper_busco.sh"

[[ -f "$DIR_ORTHOLOGER/bin/common_manage_project.sh" ]] || { echo "FATAL: unable to read DIR_ORTHOLOGER/bin/common_manage_project.sh!"; exit 255; }
source $DIR_ORTHOLOGER/bin/common_manage_project.sh

#-----------------------------------------------
# details on usage
#-----------------------------------------------
usage_setup() {
    echo
    echo -e "${ALWAYS_BOLD}Usage: $CMD_MANAGE [options]${ALWAYS_NORMAL}\n"
    echo
    echo -e "${ALWAYS_BOLD}1. General${ALWAYS_NORMAL}"
    echo " -h         this message"
    echo " -v         version"
    echo
    echo -e "${ALWAYS_BOLD}2. Output${ALWAYS_NORMAL}"
    echo " -d         dry run, just print what will be done but do nothing"
    echo " -p         use color in output messages"
    echo " -V         more verbose output, additive -> increases the verbosity"
    echo
    echo -e "${ALWAYS_BOLD}3. Project setup${ALWAYS_NORMAL}"
    echo " -s         setup empty project, orthologer/orthomapper"
    echo " -B         nr of jobs per batch (OP_NJOBMAX_BATCH)"
    echo " -L         nr of jobs in parallel (OP_NJOBMAX_LOCAL and OP_STEP_NPARALLEL)"
    echo "            if set to zero it will set it to a conservative value"
    echo " -E <path>  path for excutable search (in addition to the system path)"
    echo " -T <tool:path> give explicit path for tool (brhclus/cdhit/blast/diamond/mmseqs)"
    echo " -t <todo>  file containing just fasta labels (todo file)"
    echo
    echo " -f <file>  file containing fasta labels with corresponding raw files (see below)"
    echo " -c         copy file upon import, default is to create soft links (valid for -f and -I)"
    echo " -I <csv>   import from another project source (-S) given by <csv>"
    echo " -S <dir>   project source directory from where to import data (with -I)"
    echo
    echo " -F         force action"
    echo
    echo
    echo -e "${ALWAYS_BOLD}3. Details on options${ALWAYS_NORMAL}"
    echo
    echo -e "${ALWAYS_BOLD}   String <csv> for the import options (-I):${ALWAYS_NORMAL}"
    echo "  ==> requires source dir -S <src> <=="
    echo "  FASTA      -  import fasta files from <src>"
    echo "  SEQUENCES  -  import files from <src>/Sequences"
    echo "  PWC        -  import files from <src>/PWC"
    echo "  CLUSTER    -  import files from <src>/Cluster"
    echo
    echo "  The above strings are not case sensitive and they only have to"
    echo "  match the head (e.g seq will match SEQUENCES)"
    echo
    echo -e "${ALWAYS_BOLD}  Fasta label and files (option -f):${ALWAYS_NORMAL}"
    echo "  The file given with the -f option contains each fasta label plus the"
    echo "  corresponding file name, white space separated:"
    echo
    echo "     AGAMB    /data/agamb_source.fs"
    echo "     +AAEGY    /data/aaegy_source.fs"
    echo "     ..."
    echo
    echo "  The '+' sign in front of AAEGY indicates that the given source file requires"
    echo "  a relabeling of the sequence id's to the pipeline internal id."
    echo
    echo "  Files will be imported into the local Rawdata directory and for those requested"
    echo "  they will also be remapped."
    echo
    echo "${ALWAYS_BOLD}4. Examples${ALWAYS_NORMAL}"
    echo
    echo
    echo "  1. Create a project in current directory"
    echo
    echo "     ./${CMD_MANAGE} -s"
    echo
    echo
    echo "  2. Create a project and import fasta files given by fasta.txt"
    echo
    echo "     ./${CMD_MANAGE} -s -f fasta.txt"
    echo
    echo
    echo "  3. Import fasta files given by fasta.txt into an existing project"
    echo
    echo "     ./${CMD_MANAGE} -f fasta.txt"
    echo
    echo "     For the format of the fasta.txt file - see above."
    echo
    echo
    echo "  4. Import PWC files (by soft links) from another project, assuming project"
    echo "     dir structure is already created"
    echo
    echo "     ./${CMD_MANAGE} -t my.todo -I PWC -S /my/other/project -F"
    echo
    echo "     This command will make soft links to the PWC files defined by the"
    echo "     labels in my.todo and the source directory (/my/other/project)"
    echo
    echo
    echo "  5. Build an import file with files at a given path (data) and given extension (fs)"
    echo
    echo "     ./sbin/build_import_file.sh ./data fs"
    echo
    echo "     See './sbin/build_import_file.sh -h' for more options"
    echo
}

op_step_set_defaults() {
    OP_STEP_ENABLE[SETBRH]=0
    OP_STEP_ENABLE[SETINPAR]=0
}
#-----------------------------------------------
# main function
#-----------------------------------------------
setup_project() {

    # check args
    [[ $# -eq 0 ]] && { usage_setup; return 255; }

    # ok set 'pretty' print style...
    PRETTY_PRINT=1

    isset OPTIND && OPTIND=0

    tmp_setup_trap
    
    #----------------------
    # parse arguments
    #----------------------
    local ec=0

    # group 1 - create flags
    local do_create_project=0         #opt s
    
    # group 2 - import data
    local do_import_using_copy=0      #opt c
    local do_import_project_source=0  #opt S
    local do_import_project=0         #opt I
    local do_import_rawfasta=0        #opt f
    local do_import_xxx=0             #opt c,S,I,f
    
    local do_search_topdir=0          #opt E

    # special
    local do_force=0                  #opt F
    
    # non-argument flags
    local do_if_project_exists=0      # 0 -> make sure project does NOT exist (opts -s)
                                      # 1 -> make sure the project exists (opts -f,-S,-I)
    
    # arguments
    local arg_import_project_source=  #opt S
    local arg_import_project=         #opt I
    local arg_import_rawfasta=        #opt f
    local arg_search_topdir=          #opt E
    local arg_todo=                   #opt t
    local arg_nbatch=100              #opt B
    local arg_nproc=1                 #opt L

    # list of options without/with args
    declare -a opts_noargs=(h v p s F V d c)
    declare -a opts_withargs=(S I f E t L B T)

    # help for options with argument
    declare -A opts_args
    opts_args[S]="<source dir>"
    opts_args[I]="<import option>"
    opts_args[f]="<fasta label to raw file map>"
    opts_args[t]="<todo file>"
    opts_args[E]="<executable path in additional to PATH>"
    opts_args[B]="<nr of jobs per batch>"
    opts_args[L]="<nr of jobs running in parallel>"
    opts_args[T]="<tool:path>"

    # make option string
    local opts_str=":"
    local o=
    for o in ${opts_noargs[@]}; do
        opts_str=${opts_str}$o
    done
    for o in ${opts_withargs[@]}; do
        opts_str=${opts_str}${o}:
    done

    # parse options
    local nopts=0
    local OPT=
    local DRY_RUN=0
    export DIR_EXTERNAL_PATH=""
    local ALL_TOOLS=( BRHCLUS BLAST CDHIT DIAMOND MMSEQS . )
    

    while getopts $opts_str OPT; do
        case "$OPT" in
            # general informative options
	        h)	usage_setup
		        return 0
		        ;;
	        v)  source $DIR_ORTHOLOGER/bin/common_odb.sh
                [[ -f orthologer_conf.sh  ]] && source orthologer_conf.sh
                [[ -f orthomapper_conf.sh ]] && source orthomapper_conf.sh
                print_version
                return 0
		        ;;
            # some general variables
            p)  PRETTY_PRINT=1
                nopts=$((nopts+1))
                ;;
            V)  VERBOSE=$((VERBOSE+1))
                ;;
            d)  DRY_RUN=1
                ;;

            # create directories/scripts
	        s)  do_create_project=1
	            ;;

            # import data
            c)  do_import_using_copy=1
                do_import_xxx=1
                ;;
	        S)  do_import_project_source=1
		        arg_import_project_source="$OPTARG"
                do_import_xxx=1
	            ;;
            I)  do_import_project=1
                arg_import_project="$OPTARG"
                do_import_xxx=1
                ;;
            f)  do_import_rawfasta=1
                arg_import_rawfasta="$OPTARG"
                do_import_xxx=1
                ;;
            B)  arg_nbatch="$OPTARG"
                ;;
            L)  arg_nproc="$OPTARG"
                ;;
            # special files/directories
            t)  arg_todo="$OPTARG"
                ;;
            E)  do_search_topdir=1
                if [[ -n "$OPTARG" ]]; then
                    DIR_EXTERNAL_PATH=$(add_to_path "$DIR_EXTERNAL_PATH" "$OPTARG")
                fi
                ;;
            T)  local tool=($(tokenize ':' "${OPTARG}"))
                declare -u key="${tool[0]}"
                declare -u k=
                for k in ${ALL_TOOLS[@]}; do
                    [[ $k == $key ]] && break
                done
                if [[ "$k" == "." ]]; then
                    error_msg "unknown tool given with -T option : $key"
                else
                    # set e.g DIR_BRHCLUS
                    # these will be used in find_exe_dirs and set_executables calls below
                    local var="DIR_${key}"
                    info_msg "set $var = ${tool[1]}"
                    printf -v "$var" '%s' "${tool[1]}"
                fi
                ;;

            # force action if otherwise not allowed
            F)  do_force=1
                ;;

            # capture bad args
	        \?)	# getopts issues an error message
	            error_msg "argument parse error - for help, run with option -h"
		        return 2
		        ;;
            :)  # missing argument
                local missingarg=${opts_args[$OPTARG]}
                error_msg "missing argument -$OPTARG $missingarg"
                return 2
                ;;
        esac
    done

    find_exe_dirs
    set_executables
    
    # make sure print attributes are ok
    reset_attributes

    #--------------------------------------
    # check options for bad combinations
    #--------------------------------------
    if [[ $do_import_rawfasta -eq 1 ]]; then
        [[ $do_import_project_source -eq 0 ]] || warn_msg "ignoring option -S since -f is set"
        [[ $do_import_project        -eq 0 ]] || warn_msg "ignoring option -I since -f is set"
        #[[ $do_todo                  -eq 0 ]] || warn_msg "ignoring option -t since -f is set"
        
        do_import_project_source=0
        do_import_project=0
        #do_todo=0
    fi

    if [[ $do_import_project -eq 1 ]]; then
        [[ $do_import_project_source -eq 0 ]] && { error_msg "-I missing required option -S"; return 255; }
    fi
    
    #--------------------------------------
    # parse todo option (-t)
    # This will set PL_TODO/1/2
    #--------------------------------------
    if [[ -n "$arg_todo" ]]; then
        parse_todo_spec "$arg_todo" || return 255
    fi
    
    #--------------------------------------
    # setup orthologer environment
    #--------------------------------------
    if is_dryrun; then
        dryrun_info_msg "setup_environment()"
    else
        info_msg "manage_project: setup environment"
        setup_environment nocheck || fatal_msg "failed setup_environment()"

        info_msg "create project dirs"
        if [[ -d PWC ]]; then
            # only create potential new PWC/JobLog dirs
            create_project_dirs || { error_msg "create project dirs failed!"; return 255; }
        fi
        if [[ -n $PL_TODO ]]; then
            info_msg "update todo variables in conf"
            if [[ -f "orthologer_conf.sh" ]]; then
                # update only if config file is there
                file_setvar "orthologer_conf.sh" "PL_TODO"  "$PL_TODO"
                file_setvar "orthologer_conf.sh" "PL_TODO1" "$PL_TODO1"
                file_setvar "orthologer_conf.sh" "PL_TODO2" "$PL_TODO2"
            fi
        else
            info_msg "no todo file given"
        fi
        info_msg "check required executables"
        check_step_reqtools_all
    fi
    

    #--------------------------------------
    # Option -s : setup project
    #--------------------------------------
    # create project directories, will later update orthologer_conf.sh
    if [ $do_create_project -eq 1 ]; then
        if is_dryrun; then
            dryrun_info_msg "creating project"
        else
            info_msg "create orthologer environment"
            create_project_dirs || return 255

            # adjust OP_STEP_XXX according to settings
            set_njobmax "$arg_nproc" "$arg_nbatch"
            op_step_adjust

            # defaults for install
            op_step_set_defaults
            
            # create config file: orthologer_conf.sh
            pl_dump_conf_sh > orthologer_conf.sh  || return 255
            chmod 755 orthologer_conf.sh
            
            pl_dump_pipeline_sh > $CMD_PIPELINE   || return 255
            chmod 755 $CMD_PIPELINE
            
            # create orthomapper scripts
            BUSCO_NCPUS=$OP_NJOBMAX_LOCAL
            info_msg "create orthomapper scripts"
            dump_orthomapper_config || return 255
            {
                cat <<EOF        
#--------------------------------------------------------------------
#!/usr/bin/env bash
export CMD_ODBMAP=orthomapper.sh
./orthologer.sh ODBMAP "\$@"
EOF
            } > orthomapper.sh
            chmod +x orthomapper.sh
        fi
    fi

    #--------------------------------------
    # check if project dir is setup
    #--------------------------------------
    check_project_dirs || {
        error_msg "no project is installed in current directory"
        error_msg "setup using option -s"
        return 255
    }

    #--------------------------------------
    # Option -f : rawfasta import
    #--------------------------------------
    #if [[ $do_import_rawfasta -eq 1 ]]; then
    #    if is_dryrun; then
    #        dryrun_info_msg "import fasta defined in $arg_import_rawfasta"
    #    else
    #        [[ -f "$arg_import_rawfasta" ]] || fatal_msg "rawfasta file is not readable ($arg_import_rawfasta)"
    #        # force creation of todo file
    #        [[ -n $arg_todo ]] || arg_todo=todo/$(file_basename "$arg_import_rawfasta").todo
    #        do_todo=1
    #        # force create project dirs - will only create missing ones
    #        do_create_project_dirs=1
    #        if [[ ! -f "$arg_todo" ]]; then
    #            info_msg "creating todo file $arg_todo from rawfasta file $arg_import_rawfasta"
    #            create_todo "$arg_todo" "$arg_import_rawfasta" || return 255
    #        fi
    #    fi
    #fi
    
    ## check if PL_TODO is set and do_todo is not
    #if [[ $do_todo -eq 0 ]]; then
    #    [[ -f "$PL_TODO" ]] && { do_todo=1; arg_todo="$PL_TODO"; }
    #fi

    # force creation - check for PID file
    chk_pid_file || return 255


    #--------------------------------------
    # always check setup dirs
    #--------------------------------------
    if is_dryrun; then
        dryrun_info_msg "check project directories"
    else
        info_msg "checking project directories"
        check_project_dirs || return 255 # fail if not OK
    fi

    #--------------------------------------
    # Option -f
    #--------------------------------------
    # import raw fasta file(s)
    if [ $do_import_rawfasta -eq 1 ]; then
        if is_dryrun; then
            dryrun_info_msg "Ximporting raw fasta files from <$arg_import_rawfasta>"
        else
            info_msg "importing raw fasta files from <$arg_import_rawfasta>"
            import_fasta_from_source $do_import_using_copy "$arg_import_rawfasta"
            ec=$?
            [[ $ec -eq 0 ]] || { error_msg "failed importing raw fasta files with error code $ec"; return $ec; }
            if (( N_IMPORTED>0 )); then
                # imported some files -> make sure todo file is uptodate
                if [[ -f $arg_todo ]]; then
                    info_msg "updating todo file $arg_todo"
                else
                    info_msg "creating new todo file $arg_todo"
                fi
                # if not arg todo is set
                [[ -n $arg_todo ]] || arg_todo=todo/$(file_basename "$arg_import_rawfasta").todo
                [[ -n $PL_TODO  ]] || PL_TODO="$arg_todo"

                # set todo file in config
                file_setvar "orthologer_conf.sh" "PL_TODO" "$PL_TODO"
                
                # new file(s) imported -> update todo and corresponding directories
                create_todo "$arg_todo" "$arg_import_rawfasta" || return 255
                set_todo_vars   || return 255
                if [[ ${#PL_TODO_LABELS[@]} -eq 0 ]]; then
                    info_msg "no PWC and JobLog dirs are created as no or empty todo file was given"
                else
                    info_msg "creating PWC and JobLog dirs for the new label(s), N=${#PL_TODO_LABELS[@]}"
                    create_joblog_dirs || return 255
                    create_pwc_dirs || return 255
                fi
            fi
        fi
    fi

    #--------------------------------------
    # Option -I, -S
    # Is disabled if option -f is set
    #--------------------------------------
    if [[ $do_import_project -eq 1 ]]; then
        [[ $do_import_project_source -eq 1 ]] || { error_msg "missing -S option required by -I"; return 1; }
        if [[ ! -d "$arg_import_project_source" ]]; then
            error_msg "import data source (-S) is not a readable directory: $arg_import_project_source"
            return 1
        fi

        if is_dryrun; then
            local msg="softlinks"
            [[ $do_import_using_copy -eq 1 ]] && msg="copy"
            dryrun_info_msg "import data ($arg_import_project) using $msg from <$arg_import_project_source>"
        else
            import_data $do_import_using_copy "$arg_import_project" "$arg_import_project_source";
            ec=$?
            [[ $ec -eq 0 ]] || { error_msg "import link returned error code $ec"; return $ec; }
        fi
    fi

    return 0
}

debug_msg "setup project with args <$@>"
debug_msg "target dir: $DIR_PROJECT"

info_msg "------------------------------------------"
info_msg "Setting up project in $DIR_PROJECT"
info_msg "------------------------------------------"
# set up project
cd $DIR_PROJECT
setup_project "$@"
