#!/bin/bash

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


function list_accs()
{
    {
        echo SRR1204085 SRR1204086 SRR1204087 SRR1204088 SRR1204451 SRR1204452
        echo SRR1204453 SRR1204454 SRR1204455 SRR1204456 SRR1204466 SRR1204467
        echo SRR1204468 SRR1204469 SRR1204470 SRR1204471 SRR1204472 SRR1204473
        echo SRR1204480 SRR1204481 SRR1204482 SRR1204483 SRR1204484 SRR1204485
        echo SRR1204486 SRR1204487 SRR1204614 SRR1204615 SRR1204616 SRR1204617
        echo SRR1204618 SRR1204619 SRR1204620 SRR1204621 SRR1204689 SRR1204690
        echo SRR1204691 SRR1204692 SRR1204693 SRR1204694 SRR1204695 SRR1204696
    } | 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 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 ACC;
    do
        download_acc "$ACC"
    done

    wait_jobs
}


main "$@"
