#!/usr/bin/env bash
# find-fastest-ubuntu-mirrors.sh
set -u

SHORTLIST=50      # shortlist this many lowest-ping mirrors
TOP_N=20          # output this many fastest mirrors
PARALLEL_JOBS=8   # number of mirrors to test in parallel
TEST_SIZE=$((5*1024*1024 - 1)) # 5 MB range request
MIRRORS_PAGE="https://launchpad.net/ubuntu/+archivemirrors"
TMPDIR=$(mktemp -d)
ALL="$TMPDIR/all_mirrors.txt"
SHORT="$TMPDIR/shortlist_mirrors.txt"
SPEEDS="$TMPDIR/mirror_speeds.txt"
OUT="fastest-ubuntu-mirrors.txt"
TEST_FILE="ls-lR.gz"

trap 'rm -rf "$TMPDIR"' EXIT

echo "Fetching mirror list..."
wget -q -O- "$MIRRORS_PAGE" \
  | grep -P -B8 "statusUP" \
  | grep -o -P "(f|ht)tp://[^\"]*" > "$ALL"

if [ ! -s "$ALL" ]; then
  echo "Failed to fetch mirror list." >&2
  exit 1
fi

echo "Selecting $SHORTLIST lowest-latency mirrors with netselect..."
if sudo netselect -s"$SHORTLIST" -t50 $(cat "$ALL") 2>/dev/null \
   | awk '{
       for(i=1;i<=NF;i++) if ($i ~ /^(http|ftp)/) print $i
     }' > "$SHORT"; then
  :
else
  echo "netselect failed or returned no output — falling back to first $SHORTLIST mirrors"
  head -n "$SHORTLIST" "$ALL" > "$SHORT"
fi

test_mirror() {
  mirror="$1"
  mirror="${mirror%/}"
  url="$mirror/$TEST_FILE"

  # Primary test: 5 MB range request
  speed_bytes=$(curl -s -L --connect-timeout 5 -m 30 \
      --range 0-"$TEST_SIZE" -o /dev/null -w '%{speed_download}' "$url" 2>/dev/null || echo 0)

  # fallback without range (in case server doesn’t support it)
  if [ -z "$speed_bytes" ] || [ "$speed_bytes" = "0" ]; then
    speed_bytes=$(curl -s -L --connect-timeout 5 -m 30 \
        -o /dev/null -w '%{speed_download}' "$url" 2>/dev/null || echo 0)
  fi

  if [ -z "$speed_bytes" ] || [ "$speed_bytes" = "0" ]; then
    echo "0 $mirror" >> "$SPEEDS"
    echo "FAILED $mirror"
  else
    speed_int=$(awk -v s="$speed_bytes" 'BEGIN{printf "%d", s}')
    echo "$speed_int $mirror" >> "$SPEEDS"

    if command -v numfmt >/dev/null 2>&1; then
      hr=$(numfmt --to=iec --suffix=B/s "$speed_int")
    else
      hr=$(awk -v b="$speed_int" 'BEGIN{printf "%.1f KB/s", b/1024}')
    fi
    echo "OK     $mirror -> $hr"
  fi
}

export -f test_mirror
export SPEEDS TEST_FILE TEST_SIZE

echo "Testing throughput (~5 MB per mirror) with $PARALLEL_JOBS jobs..."
: > "$SPEEDS"
cat "$SHORT" | xargs -n1 -P"$PARALLEL_JOBS" bash -c 'test_mirror "$@"' _

echo
echo "Picking top $TOP_N mirrors..."
sort -nr "$SPEEDS" | head -n "$TOP_N" > "$OUT"

echo "Top $TOP_N mirrors saved to: $OUT"
nl -ba -w2 -s'. ' <(awk '{if (command="numfmt --to=iec --suffix=B/s " $1 | getline hr) {print hr, $2} else {printf "%.1f KB/s %s\n", $1/1024, $2}}' "$OUT")