#!/bin/bash

# Unofficial Bash Strict Mode (http://redsymbol.net/articles/unofficial-bash-strict-mode/)
set -euo pipefail
IFS=$'\n\t'


function list_accs()
{
    {
        echo SRR4016879 SRR4016880 SRR4016881 SRR4016882 SRR4016884 SRR4016886
        echo SRR4016887 SRR4016890 SRR4016893 SRR4016895 SRR4016899 SRR4016903
        echo SRR4016904 SRR4016905 SRR4016907 SRR4016910 SRR4016912 SRR4016914
        echo SRR4016916 SRR4016921 SRR4016923 SRR4016925 SRR4016927 SRR4016928
        echo SRR4016929 SRR4016930 SRR4016888 SRR4016889 SRR4016891 SRR4016892
        echo SRR4016896 SRR4016898 SRR4016900 SRR4016902 SRR4016906 SRR4016908
        echo SRR4016911 SRR4016913 SRR4016917 SRR4016918 SRR4016919 SRR4016920
        echo SRR4016922 SRR4016924 SRR4016926 SRR4016931 SRR4016933 SRR4016934
        echo SRR4016935 SRR4016936 SRR4016937 SRR4016939 SRR4016940 SRR4016941
        echo SRR4016894 SRR4016915 SRR4016938 SRR4016883 SRR4016897 SRR4016885
        echo SRR4016909 SRR4016932 SRR4016901
    } | tr ' ' '\n'
}


TRACE_URL="https://trace.ncbi.nlm.nih.gov/Traces/sra/?run="


function get_urls_for_acc()
{
    local ACC="$1"

    if [[ -e ".$ACC-downloaded" ]] && (( $(wc -l < ".$ACC-downloaded") > 0 ));
    then
        log "using URLs from cache: .$ACC-downloaded"
    else
        log "fetching URLs from ${TRACE_URL}$ACC"

        curl --silent --show-error "${TRACE_URL}$ACC" | \
            grep -oE 'https://sra-pub-src[^"]+\.bax\.h5(\.1)?' | \
            sed -E 's!https://(.*)\.s3\.amazonaws\.com/(.*)!s3://\1/\2!' \
            > ".$ACC-downloaded"
    fi

    cat ".$ACC-downloaded"
}


function get_state_of()
{
    local URL="$1"
    local FILE="$(get_filename_for_url "$URL")"

    if [[ -e "$FILE" ]];
    then
        echo 'skip' && return
    else
        echo 'download' && return
    fi

    return 1
}


function download_file()
{
    local URL="$1"
    local FILENAME="$2"

    aws s3 cp --no-sign-request "$URL" "$FILENAME" 1>&2
}


function log()
{
    echo "--" "$@"
} >&2


function call_hook()
{
    if [[ "$1" == "--optional" ]];
    then
        local OPTIONAL=yes
        shift
    fi

    if { declare -F "$1" || alias "$1" || which "$1"; } &> /dev/null;
    then
        eval "$@"
    else
        if [[ "${OPTIONAL:-no}" == no ]];
        then
            error "missing definition for hook: $1"
            exit 1
        fi
    fi
}


function get_filename_for_url()
{
    local FILENAME="${1##*/}"
    FILENAME="${FILENAME%.1}"

    echo "$FILENAME"
}


function print_job_status_line()
{
    tail -c80 "$1" >&2
}


function init_job_control()
{
    export JOBS_FILE="$(mktemp .jobs-XXXXXX)"

    trap "kill -SIGTERM \$(job_pids)" err
    trap 'rm -f .job-*.log &> /dev/null; rm -f '"$JOBS_FILE" exit
}


function track_job()
{
    local JOBID="$1"
    local PID="$2"
    local LOG="$3"

    echo -e "$JOBID\t$PID\t$LOG\trunning" >> "$JOBS_FILE"
}


function new_jobid()
{
    awk -F'\t' '{ if ($1 > jobid) jobid = $1 } END { print jobid + 1 }' "$JOBS_FILE"
}


function enqueue_download_job()
{
    wait_for_queue

    local URL="$1"
    local CONTINUE="$2"
    local JOBID="$(new_jobid)"
    local LOG=".job-$JOBID.log"
    local FILENAME="$(call_hook get_filename_for_url "$URL")"

    {
        call_hook download_file "$URL" "$FILENAME" "$CONTINUE" 2> "$LOG" || job_failed "$URL" "$LOG"
    } &

    track_job "$JOBID" "$!" "$LOG"
}


function do_download_job()
{
    local URL="$1"
    local CONTINUE="$2"
    local FILENAME="$(call_hook get_filename_for_url "$URL")"

    call_hook download_file "$URL" "$FILENAME" "$CONTINUE"
}


function log_job_status()
{
    local PID="$1"
    local LOG="$(get_logfile "$PID")"
    local JOBID="${LOG%.log}"
    JOBID="${JOBID##*-}"

    log "job $JOBID status:"
    call_hook print_job_status_line "$LOG" "$JOBID" "$PID" >&2
}


function is_job_running()
{
    [[ -e "/proc/$1" ]]
}


function check_job_states()
{
    local NEW_STATES="$( awk -F'\t' '
        BEGIN {
            OFS = "\t";
        }

        {
            process_exe = "/proc/"$2"/exe";

            if ( (getline dummy < process_exe) > 0 ) {
                $4 = "running";
                close(process_exe)
            } else {
                $4 = "stopped"
            }

            print
        }
    ' "$JOBS_FILE" )";

    echo "$NEW_STATES" > "$JOBS_FILE"
}


function num_jobs()
{
    check_job_states

    awk -F'\t' '
        ($4 == "running") {
            ++num_jobs;
        }

        END {
            print 0 + num_jobs
        }
    ' "$JOBS_FILE"
}


function job_pids()
{
    awk -F'\t' '
        ($4 == "running") {
            print $2;
        }
    ' "$JOBS_FILE"
}


function get_logfile()
{
    awk -F'\t' '
        ($2 == "'"$1"'") {
            print $3;
        }
    ' "$JOBS_FILE"
}


function wait_for_queue()
{
    if parallel;
    then
        while (( $(num_jobs) >= PARALLEL ));
        do
            for PID in $(job_pids);
            do
                log_job_status "$PID"
            done

            sleep "$PARALLEL"
        done
    fi
}


function wait_jobs()
{
    if parallel
    then
        while (( $(num_jobs) > 0 ));
        do
            for PID in $(job_pids);
            do
                log_job_status "$PID"
            done

            sleep "$PARALLEL"
        done
    fi
}


function job_failed()
{
    local URL="$1"
    local LOG="$2"
    local JOBID="${LOG%.log}"
    JOBID="${JOBID##*-}"
    local ERRLOG="error.$JOBID.log"

    mv "$LOG" "$ERRLOG"
    error "failed to download $URL; inspect $ERRLOG for details"
}


function download_acc()
{
    local ACC="$1"
    local URL

    call_hook --optional init_download_acc "$ACC"
    log "syncing $ACC"

    call_hook get_urls_for_acc "$ACC" | while read -r URL;
    do
        local FILE_STATE="$(call_hook get_state_of "$URL")"
        case "$FILE_STATE" in
            skip)
                log "skipping $URL"
                DOWNLOAD=no
                ;;
            resume)
                log "resuming $URL"
                CONTINUE=yes
                DOWNLOAD=yes
                ;;
            replace)
                log "removing corrupted file for $URL"
                dryrun || rm -f "$FILENAME"
                DOWNLOAD=yes
                ;;
            download)
                DOWNLOAD=yes
                ;;
            *)
                error "illegal state: $FILE_STATE"
        esac

        if [[ "${DOWNLOAD:-yes}" != "no" ]];
        then
            log "downloading $URL (job $(num_jobs))"
            if parallel;
            then
                dryrun || enqueue_download_job "$URL" "${CONTINUE:-no}"
            else
                dryrun || do_download_job "$URL" "${CONTINUE:-no}"
            fi
        fi
    done

    call_hook --optional on_acc_downloaded "$ACC"
}


PROG="$0"


function error()
{
    echo "$PROG: error:" "$@"
} >&2


function usage()
{
    echo "USAGE:  $PROG [-hn] [-p<num>]"
} >&2


function parse_args()
{
    while getopts "hnp:" OPTION; do
        case "$OPTION" in
            p)
                PARALLEL="$OPTARG"

                if ! (( OPTARG ));
                then
                    error "-p must have an integer value"
                    exit 1
                elif ! (( OPTARG > 0 ));
                then
                    error "-p must be positive"
                    exit 1
                fi
                ;;
            n)
                DRYRUN='yes'
                ;;
            h)
                usage

                exit
                ;;
            *)
                error "unkown option $OPTION"
                usage

                exit 1
                ;;
        esac
    done
    shift $((OPTIND - 1))
}


function parallel()
{
    [[ -v PARALLEL ]]
}


function dryrun()
{
    [[ -v DRYRUN ]]
}


function main()
{
    call_hook --optional before_parse_args
    parse_args "$@"

    if parallel;
    then
        init_job_control
    fi

    call_hook --optional before_all_downloads
    call_hook list_accs | while read -r ACC;
    do
        download_acc "$ACC"
    done

    wait_jobs
}


main "$@"
