#!/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}"
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}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
OUTDIR="${OUTDIR:-${SCRIPT_DIR}}"
URLS_TXT="${URLS_TXT:-${SCRIPT_DIR}/${PROJECT}.aria2c.urls.txt}"
MD5_TXT="${MD5_TXT:-${SCRIPT_DIR}/${PROJECT}.fastq.md5}"
LOG_FILE="${LOG_FILE:-${SCRIPT_DIR}/${PROJECT}.aria2c.log}"

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

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

Required file in the same directory as this script:
  ${PROJECT}.aria2c.urls.txt

Options:
  --outdir DIR        Output directory. Default: script directory
  --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
USAGE
}

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

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 <<ERR
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
ERR
  exit 127
fi

if [[ ! -f "${URLS_TXT}" ]]; then
  echo "ERROR: URL list not found: ${URLS_TXT}" >&2
  exit 1
fi

mkdir -p "${OUTDIR}"

URL_COUNT="$(wc -l < "${URLS_TXT}")"
if [[ "${URL_COUNT}" -eq 0 ]]; then
  echo "ERROR: URL list is empty: ${URLS_TXT}" >&2
  exit 1
fi

echo "FASTQ URL count: ${URL_COUNT}"
echo "Output directory: ${OUTDIR}"
echo "URL list: ${URLS_TXT}"
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}"
if [[ -f "${MD5_TXT}" ]]; then
  echo "MD5 list: ${MD5_TXT}"
fi
echo "Log file: ${LOG_FILE}"
if [[ -f "${MD5_TXT}" ]]; then
  echo "To verify MD5 after download:"
  echo "  cd '${OUTDIR}' && md5sum -c '${MD5_TXT}'"
fi