#!/bin/bash

# Debian/Ubuntu mirror optimizer with optional daily cron refresh.
# Debian path prefers netselect-apt for broad mirror testing.
set -euo pipefail

AUTO_MODE=0
QUIET_MODE=0
INSTALL_CRON=ask

for arg in "$@"; do
  case "$arg" in
    --auto) AUTO_MODE=1 ;;
    --quiet) QUIET_MODE=1 ;;
    --install-cron) INSTALL_CRON=yes ;;
    --no-cron) INSTALL_CRON=no ;;
    *)
      echo "Unknown option: $arg"
      echo "Usage: $0 [--auto] [--quiet] [--install-cron|--no-cron]"
      exit 1
      ;;
  esac
done

log() {
  if [ "$QUIET_MODE" -eq 0 ]; then
    echo "$@"
  fi
}

need_cmd() {
  command -v "$1" >/dev/null 2>&1 || {
    echo "Missing required command: $1"
    exit 1
  }
}

need_cmd curl
need_cmd awk
need_cmd sed
need_cmd tee

if [ ! -f /etc/os-release ]; then
  echo "Cannot detect OS: /etc/os-release not found"
  exit 1
fi

# shellcheck disable=SC1091
. /etc/os-release
DISTRO="${ID:-}"
CODENAME="${VERSION_CODENAME:-}"

if [ -z "$CODENAME" ] && [ -n "${VERSION:-}" ]; then
  CODENAME=$(echo "$VERSION" | awk -F'[() ]+' '{for (i=1;i<=NF;i++) if ($i ~ /^[a-z]+$/) {print $i; exit}}')
fi

if [ -z "$CODENAME" ]; then
  echo "Could not determine release codename."
  exit 1
fi

if [ "$DISTRO" != "debian" ] && [ "$DISTRO" != "ubuntu" ]; then
  echo "Unsupported distro: $DISTRO (supported: debian, ubuntu)"
  exit 1
fi

if [ "$DISTRO" = "ubuntu" ]; then
  MIRRORS=(
    "http://archive.ubuntu.com/ubuntu"
    "http://uk.archive.ubuntu.com/ubuntu"
    "http://us.archive.ubuntu.com/ubuntu"
    "http://mirror.math.princeton.edu/pub/ubuntu"
    "http://mirror.csclub.uwaterloo.ca/ubuntu"
  )
  TEST_PATH="/dists/$CODENAME/InRelease"
fi

BACKUP="/etc/apt/sources.list.backup.$(date +%Y%m%d%H%M%S)"
if [ -f /etc/apt/sources.list ]; then
  sudo cp /etc/apt/sources.list "$BACKUP"
  echo "Backup saved: $BACKUP"
fi

BEST_MIRROR=""
BEST_TIME="999"
GENERATED_SOURCES=""

select_mirror_with_netselect_apt() {
  local out_file
  local err_file
  local work_dir
  local generated_file
  out_file=$(mktemp /tmp/aaran-netselect-apt.XXXXXX.out)
  err_file=$(mktemp /tmp/aaran-netselect-apt.XXXXXX.err)
  work_dir=$(mktemp -d /tmp/aaran-netselect-apt.XXXXXX.dir)

  echo "Installing netselect-apt for full Debian mirror testing..."
  sudo apt update
  sudo apt install -y netselect-apt netselect ca-certificates curl

  echo "Running netselect-apt mirror test for Debian (this can take a minute)..."
  if (cd "$work_dir" && netselect-apt >"$out_file" 2>"$err_file"); then
    :
  fi

  generated_file=$(find "$work_dir" -maxdepth 1 -type f -name 'sources.list*' | head -n 1 || true)
  if [ -n "$generated_file" ] && grep -E '^deb\s+(https?|ftp)://' "$generated_file" >/dev/null 2>&1; then
    GENERATED_SOURCES="$generated_file"
    rm -f "$err_file"
    rm -f "$out_file"
    return 0
  fi

  if grep -E '^deb\s+(https?|ftp)://' "$out_file" >/dev/null 2>&1; then
    GENERATED_SOURCES="$out_file"
    rm -f "$err_file"
    rm -rf "$work_dir"
    return 0
  fi

  if [ -s "$err_file" ]; then
    echo "netselect-apt error output:"
    sed 's/^/  /' "$err_file" || true
  fi

  rm -f "$err_file"
  rm -f "$out_file"
  rm -rf "$work_dir"
  return 1
}

select_mirror_from_shortlist() {
  log "Testing mirrors for $DISTRO $CODENAME..."

  for mirror in "${MIRRORS[@]}"; do
    test_url="${mirror}${TEST_PATH}"
    t=$(curl -LfsS --connect-timeout 3 --max-time 8 -o /dev/null -w "%{time_total}" "$test_url" || true)

    if [ -n "$t" ]; then
      log "  $mirror -> ${t}s"
      if awk "BEGIN {exit !($t < $BEST_TIME)}"; then
        BEST_TIME="$t"
        BEST_MIRROR="$mirror"
      fi
    else
      log "  $mirror -> unreachable"
    fi
  done

  if [ -z "$BEST_MIRROR" ]; then
    echo "No reachable mirrors found."
    exit 1
  fi

  echo "Best mirror selected: $BEST_MIRROR (${BEST_TIME}s)"
}

if [ "$DISTRO" = "debian" ]; then
  if select_mirror_with_netselect_apt; then
    echo "Best mirrors selected via netselect-apt."
    sudo cp "$GENERATED_SOURCES" /etc/apt/sources.list
    rm -f "$GENERATED_SOURCES"
  else
    echo "netselect-apt selection failed, falling back to curated mirror latency test."
    MIRRORS=(
      "http://deb.debian.org/debian"
      "http://ftp.uk.debian.org/debian"
      "http://ftp.de.debian.org/debian"
      "http://ftp.us.debian.org/debian"
      "http://ftp.nl.debian.org/debian"
    )
    TEST_PATH="/dists/$CODENAME/InRelease"
    select_mirror_from_shortlist

    cat <<EOF | sudo tee /etc/apt/sources.list >/dev/null
deb $BEST_MIRROR $CODENAME main contrib non-free non-free-firmware
deb $BEST_MIRROR ${CODENAME}-updates main contrib non-free non-free-firmware
deb http://security.debian.org/debian-security ${CODENAME}-security main contrib non-free non-free-firmware
EOF
  fi
else
  select_mirror_from_shortlist

  cat <<EOF | sudo tee /etc/apt/sources.list >/dev/null
deb $BEST_MIRROR $CODENAME main restricted universe multiverse
deb $BEST_MIRROR ${CODENAME}-updates main restricted universe multiverse
deb $BEST_MIRROR ${CODENAME}-backports main restricted universe multiverse
deb http://security.ubuntu.com/ubuntu ${CODENAME}-security main restricted universe multiverse
EOF
fi

echo "Updated /etc/apt/sources.list"
sudo apt update

echo "Mirror optimization applied successfully."

install_cron_job() {
  cron_line='18 4 * * * tmpfile="$(mktemp /tmp/aaran-mirror-auto.XXXXXX.sh)" && curl -fsSL https://aaran.cloud/assets/scripts/optimize-debian-ubuntu-mirrors.sh -o "$tmpfile" && bash "$tmpfile" --auto --quiet --no-cron; rm -f "$tmpfile"'

  existing=$(sudo crontab -l 2>/dev/null || true)
  if echo "$existing" | grep -F "aaran-mirror-auto" >/dev/null 2>&1; then
    echo "Daily mirror cron job already present."
    return
  fi

  printf "%s\n%s\n" "$existing" "$cron_line" | sudo crontab -
  echo "Installed daily cron job at 04:18 server time."
}

if [ "$AUTO_MODE" -eq 1 ]; then
  if [ "$INSTALL_CRON" = "yes" ]; then
    install_cron_job
  fi
  exit 0
fi

if [ "$INSTALL_CRON" = "ask" ]; then
  echo
  echo "Do you want to install a daily cron job to re-check and adjust mirrors automatically?"
  echo "  [y] Yes, install daily cron"
  echo "  [n] No, apply now only (you can run manually later)"
  printf "Choose [y/N]: "
  read -r answer

  case "$answer" in
    y|Y|yes|YES)
      install_cron_job
      ;;
    *)
      echo "No cron job installed."
      ;;
  esac
elif [ "$INSTALL_CRON" = "yes" ]; then
  install_cron_job
else
  echo "No cron job installed."
fi