#!/usr/bin/env bash
set -Eeuo pipefail

readonly RAW_ROOT="${1:-/data/digital-drosophila/datasets/male-cns-v1.0/raw}"
readonly BASE_URL="https://storage.googleapis.com/flyem-male-cns/v1.0/connectome-data/flat-connectome"
readonly LOG_DIR="${RAW_ROOT}/_logs"
readonly MANIFEST_DIR="${RAW_ROOT}/_manifests"

mkdir -p "${RAW_ROOT}/annotations" "${RAW_ROOT}/connectivity" "$LOG_DIR" "$MANIFEST_DIR"

download_one() {
    local subdir="$1"
    local filename="$2"
    local expected_md5="$3"
    local destination="${RAW_ROOT}/${subdir}/${filename}"
    local attempt

    if [[ -f "$destination" ]] && printf '%s  %s\n' "$expected_md5" "$destination" | md5sum --check --status; then
        printf '[OK] already verified: %s/%s\n' "$subdir" "$filename"
        return 0
    fi

    for attempt in 1 2 3; do
        if [[ -f "$destination" && ! -f "${destination}.aria2" ]]; then
            mv -- "$destination" "${destination}.corrupt-$(date +%Y%m%dT%H%M%S)"
        fi

        printf '[DOWNLOAD] attempt %d: %s/%s\n' "$attempt" "$subdir" "$filename"
        aria2c \
            --continue=true \
            --max-tries=0 \
            --retry-wait=10 \
            --connect-timeout=30 \
            --timeout=60 \
            --max-connection-per-server=8 \
            --split=8 \
            --min-split-size=16M \
            --file-allocation=falloc \
            --auto-file-renaming=false \
            --allow-overwrite=false \
            --check-integrity=true \
            --console-log-level=notice \
            --summary-interval=60 \
            --dir="${RAW_ROOT}/${subdir}" \
            --out="$filename" \
            "${BASE_URL}/${filename}" \
            >"${LOG_DIR}/${filename}.log" 2>&1

        if printf '%s  %s\n' "$expected_md5" "$destination" | md5sum --check --status; then
            printf '[VERIFIED] %s/%s\n' "$subdir" "$filename"
            return 0
        fi

        printf '[CHECKSUM FAILED] %s/%s\n' "$subdir" "$filename" >&2
        [[ -f "$destination" ]] && mv -- "$destination" "${destination}.corrupt-$(date +%Y%m%dT%H%M%S)"
        [[ -f "${destination}.aria2" ]] && mv -- "${destination}.aria2" "${destination}.aria2.corrupt-$(date +%Y%m%dT%H%M%S)"
    done

    printf '[FAILED] checksum mismatch after 3 complete downloads: %s/%s\n' "$subdir" "$filename" >&2
    return 1
}

jobs=(
    "annotations|body-annotations-male-cns-v1.0-minconf-0.5.feather|50a7718770c57220f160ba4f431ab89e"
    "annotations|body-neurotransmitters-male-cns-v1.0.feather|3d842b12fe5c49eefade528d7dd24a1f"
    "connectivity|connectome-weights-male-cns-v1.0-minconf-0.5.feather|f30e9dcca25cfd021bf1e7b3d975599e"
)

pids=()
for job in "${jobs[@]}"; do
    IFS='|' read -r subdir filename expected_md5 <<<"$job"
    download_one "$subdir" "$filename" "$expected_md5" &
    pids+=("$!")
done

status=0
for pid in "${pids[@]}"; do
    wait "$pid" || status=1
done

if (( status != 0 )); then
    printf '[FAILED] one or more files could not be downloaded and verified\n' >&2
    exit 1
fi

cd "$RAW_ROOT"
md5sum --check "${MANIFEST_DIR}/checksums.md5"
sha256sum \
    annotations/*.feather \
    connectivity/*.feather \
    >"${MANIFEST_DIR}/checksums.sha256"

find annotations connectivity -maxdepth 1 -type f -name '*.feather' -printf '%P\t%s bytes\n' \
    | sort >"${MANIFEST_DIR}/files.tsv"

printf '[COMPLETE] all MaleCNS v1.0 phase-1 files passed official MD5 verification\n'