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

PROJECT="PRJNA530826"
CONDA_SH="${CONDA_SH:-/home/xia-lab/miniconda3/etc/profile.d/conda.sh}"
CONDA_ENV="${CONDA_ENV:-rnaseq}"
OUTDIR="${OUTDIR:-workspace/rnaseq_downloads/${PROJECT}}"
JOBS="${JOBS:-4}"
CONNECTIONS="${CONNECTIONS:-8}"
SPLITS="${SPLITS:-8}"
MIN_SPLIT_SIZE="${MIN_SPLIT_SIZE:-20M}"
RETRY_WAIT="${RETRY_WAIT:-10}"
MAX_TRIES="${MAX_TRIES:-10}"

usage() {
  cat <<EOF
Usage: bash $0 [options]

Download ${PROJECT} RNA-seq FASTQ files from ENA with aria2c.

Options:
  --outdir DIR        Output directory. Default: ${OUTDIR}
  --jobs N           Concurrent files. Default: ${JOBS}
  --connections N    Connections per file/server. Default: ${CONNECTIONS}
  --splits N         Split count per file. Default: ${SPLITS}
  --env NAME         Conda environment. Default: ${CONDA_ENV}
  --help             Show this help.

Examples:
  bash $0
  bash $0 --jobs 6 --connections 8 --splits 8
  OUTDIR=/data/rnaseq_fastq bash $0
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --outdir)
      OUTDIR="$2"
      shift 2
      ;;
    --jobs)
      JOBS="$2"
      shift 2
      ;;
    --connections)
      CONNECTIONS="$2"
      shift 2
      ;;
    --splits)
      SPLITS="$2"
      shift 2
      ;;
    --env)
      CONDA_ENV="$2"
      shift 2
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

mkdir -p "${OUTDIR}"

if [[ -f "${CONDA_SH}" ]]; then
  # shellcheck source=/dev/null
  source "${CONDA_SH}"
  conda activate "${CONDA_ENV}" || {
    echo "ERROR: failed to activate conda env: ${CONDA_ENV}" >&2
    exit 1
  }
else
  echo "WARN: conda init script not found: ${CONDA_SH}; using current PATH" >&2
fi

if ! command -v aria2c >/dev/null 2>&1; then
  cat >&2 <<EOF
ERROR: aria2c is not available in PATH after activating '${CONDA_ENV}'.

Install it first, for example:
  conda install -n ${CONDA_ENV} -c conda-forge aria2

Then rerun:
  bash $0 --outdir "${OUTDIR}"
EOF
  exit 127
fi

if ! command -v curl >/dev/null 2>&1; then
  echo "ERROR: curl is required to fetch the ENA file list." >&2
  exit 127
fi

META_TSV="${OUTDIR}/${PROJECT}.ena_fastq.tsv"
URLS_TXT="${OUTDIR}/${PROJECT}.aria2c.urls.txt"
MD5_TXT="${OUTDIR}/${PROJECT}.fastq.md5"
LOG_FILE="${OUTDIR}/${PROJECT}.aria2c.log"

ENA_URL="https://www.ebi.ac.uk/ena/portal/api/filereport?accession=${PROJECT}&result=read_run&fields=run_accession,sample_accession,experiment_accession,scientific_name,fastq_ftp,fastq_aspera,fastq_bytes,fastq_md5&format=tsv&download=false&limit=0"

echo "Fetching ENA FASTQ metadata: ${PROJECT}"
curl -L -sS "${ENA_URL}" -o "${META_TSV}"

awk -F '\t' '
  NR > 1 && $5 != "" {
    n = split($5, urls, ";")
    for (i = 1; i <= n; i++) {
      if (urls[i] != "") {
        u = urls[i]
        sub(/^ftp:\/\//, "", u)
        print "https://" u
      }
    }
  }
' "${META_TSV}" > "${URLS_TXT}"

awk -F '\t' '
  NR > 1 && $8 != "" && $5 != "" {
    n = split($5, urls, ";")
    m = split($8, md5s, ";")
    for (i = 1; i <= n; i++) {
      if (urls[i] != "" && md5s[i] != "") {
        file = urls[i]
        sub(/^.*\//, "", file)
        print md5s[i] "  " file
      }
    }
  }
' "${META_TSV}" > "${MD5_TXT}"

URL_COUNT="$(wc -l < "${URLS_TXT}")"
if [[ "${URL_COUNT}" -eq 0 ]]; then
  echo "ERROR: no FASTQ URLs were generated. Check ${META_TSV}" >&2
  exit 1
fi

echo "FASTQ URL count: ${URL_COUNT}"
echo "Output directory: ${OUTDIR}"
echo "aria2c settings: jobs=${JOBS}, connections=${CONNECTIONS}, splits=${SPLITS}"

aria2c \
  --dir="${OUTDIR}" \
  --input-file="${URLS_TXT}" \
  --continue=true \
  --max-concurrent-downloads="${JOBS}" \
  --max-connection-per-server="${CONNECTIONS}" \
  --split="${SPLITS}" \
  --min-split-size="${MIN_SPLIT_SIZE}" \
  --retry-wait="${RETRY_WAIT}" \
  --max-tries="${MAX_TRIES}" \
  --file-allocation=none \
  --auto-file-renaming=false \
  --allow-overwrite=false \
  --summary-interval=60 \
  --console-log-level=notice \
  --log="${LOG_FILE}"

echo "Download finished."
echo "URL list: ${URLS_TXT}"
echo "MD5 list: ${MD5_TXT}"
echo "Log file: ${LOG_FILE}"
echo "To verify MD5 after download:"
echo "  cd '${OUTDIR}' && md5sum -c '${PROJECT}.fastq.md5'"