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

if [[ "$#" -lt 4 || "$#" -gt 5 ]]; then
  echo 'Usage: parallel-download URL OUTPUT EXPECTED_SIZE EXPECTED_SHA [CONNECTIONS]' >&2
  exit 2
fi

readonly source_url="$1"
readonly output_file="$2"
readonly expected_size="$3"
readonly expected_sha="$4"
readonly connections="${5:-64}"

[[ "$expected_size" =~ ^[1-9][0-9]*$ ]]
[[ "$connections" =~ ^[1-9][0-9]*$ ]]

if [[ -f "$output_file" ]] \
  && [[ "$(stat -c '%s' "$output_file")" == "$expected_size" ]] \
  && [[ "$(sha256sum "$output_file" | cut -d' ' -f1)" == "$expected_sha" ]]; then
  echo "ALREADY_OK $output_file"
  exit 0
fi

readonly chunk_size=$(( (expected_size + connections - 1) / connections ))
readonly part_dir="${output_file}.parts-${expected_sha:0:8}"
mkdir -p "$part_dir"

echo "DOWNLOAD_START file=$output_file bytes=$expected_size connections=$connections"
pids=()
parts=()
for ((index = 0; index < connections; index++)); do
  start=$(( index * chunk_size ))
  (( start < expected_size )) || break
  end=$(( start + chunk_size - 1 ))
  (( end < expected_size )) || end=$(( expected_size - 1 ))
  part="$(printf '%s/part-%04d' "$part_dir" "$index")"
  parts+=("$part")
  (
    expected_part_size=$(( end - start + 1 ))
    if [[ -f "$part" ]] && [[ "$(stat -c '%s' "$part")" == "$expected_part_size" ]]; then
      echo "PART_CACHED index=$index bytes=$expected_part_size"
      exit 0
    fi
    partial_part="${part}.partial"
    if [[ ! -f "$partial_part" ]]; then
      shopt -s nullglob
      old_candidates=("${part}.tmp."*)
      shopt -u nullglob
      if (( ${#old_candidates[@]} > 0 )); then
        largest_candidate="${old_candidates[0]}"
        for candidate in "${old_candidates[@]}"; do
          if (( $(stat -c '%s' "$candidate") > $(stat -c '%s' "$largest_candidate") )); then
            largest_candidate="$candidate"
          fi
        done
        mv -f "$largest_candidate" "$partial_part"
      fi
    fi
    if [[ -f "$partial_part" ]] && (( $(stat -c '%s' "$partial_part") > expected_part_size )); then
      mv "$partial_part" "${partial_part}.invalid.$(date +%s)"
    fi

    no_progress_attempts=0
    while :; do
      current_size=0
      [[ ! -f "$partial_part" ]] || current_size="$(stat -c '%s' "$partial_part")"
      if (( current_size == expected_part_size )); then
        break
      fi
      resume_start=$(( start + current_size ))
      segment="${partial_part}.segment.$$"
      curl_status=0
      curl --noproxy '*' -fsSL \
        --connect-timeout 20 --speed-limit 512 --speed-time 120 \
        --range "$resume_start-$end" --output "$segment" "$source_url" \
        || curl_status=$?
      segment_size=0
      [[ ! -f "$segment" ]] || segment_size="$(stat -c '%s' "$segment")"
      if (( segment_size > 0 )); then
        cat "$segment" >> "$partial_part"
        no_progress_attempts=0
      else
        no_progress_attempts=$(( no_progress_attempts + 1 ))
      fi
      rm -f "$segment"
      current_size="$(stat -c '%s' "$partial_part" 2>/dev/null || echo 0)"
      if (( current_size > expected_part_size )); then
        echo "PART_OVERFLOW index=$index expected=$expected_part_size actual=$current_size" >&2
        exit 3
      fi
      if (( no_progress_attempts >= 20 )); then
        echo "PART_STALLED index=$index bytes=$current_size curl_status=$curl_status" >&2
        exit 4
      fi
      if (( curl_status != 0 )); then
        echo "PART_RESUME index=$index bytes=$current_size curl_status=$curl_status"
        sleep 1
      fi
    done
    mv -f "$partial_part" "$part"
    echo "PART_OK index=$index bytes=$expected_part_size"
  ) &
  pids+=("$!")
done

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

assembly="${output_file}.assembling.$$"
: > "$assembly"
for part in "${parts[@]}"; do
  cat "$part" >> "$assembly"
done

[[ "$(stat -c '%s' "$assembly")" == "$expected_size" ]]
printf '%s  %s\n' "$expected_sha" "$assembly" | sha256sum -c -
mv -f "$assembly" "$output_file"
find "$part_dir" -type f -delete
rmdir "$part_dir"
echo "DOWNLOAD_OK file=$output_file bytes=$expected_size sha256=$expected_sha"