#!/bin/bash
set -u -o pipefail

# Автоматическое удержание macOS от сна во время выполнения скрипта
if [[ "$(uname)" == "Darwin" ]] && command -v caffeinate >/dev/null 2>&1; then
  if [[ -z "${CAFFEINATED:-}" ]]; then
    export CAFFEINATED=1
    exec caffeinate -dimsu "$0" "$@"
  fi
fi

echo "====================================================="
echo "⚡ Напоминание: Убедитесь, что Amphetamine включен!"
echo "   (Запущен caffeinate для фоновой защиты от сна)"
echo "====================================================="
echo "Запуск быстрой проверки доступности (1 МБ RX/TX + Ping)..."
echo

IPERF_TIMEOUT=4
TEST_BYTES="1M"
DEFAULT_IPERF_PORT_RANGE=(5201 5202 5203 5204 5205 5206 5207 5208 5209)

declare -A SERVERS=(
  ["Москва"]="spd-rudp.hostkey.ru"
  ["Санкт-Петербург"]="st.spb.ertelecom.ru"
  ["Барнаул"]="st.barnaul.ertelecom.ru"
  ["Владивосток"]="vldvst.st.mtsws.net"
  ["Волгоград"]="st.volgograd.ertelecom.ru"
  ["Екатеринбург"]="st.ekat.ertelecom.ru"
  ["Ижевск"]="st.izhevsk.ertelecom.ru"
  ["Иркутск"]="st.irkutsk.ertelecom.ru"
  ["Казань"]="st.kzn.ertelecom.ru"
  ["Киров"]="st.kirov.ertelecom.ru"
  ["Красноярск"]="st.krsk.ertelecom.ru"
  ["Курган"]="st.kurgan.ertelecom.ru"
  ["Курск"]="st.kursk.ertelecom.ru"
  ["Липецк"]="st.lipetsk.ertelecom.ru"
  ["Магнитогорск"]="st.mgn.ertelecom.ru"
  ["Набережные Челны"]="st.chelny.ertelecom.ru"
  ["Нижний Новгород"]="st.nn.ertelecom.ru"
  ["Новосибирск"]="st.nsk.ertelecom.ru"
  ["Норильск"]="norst.st.mtsws.net"
  ["Омск"]="st.omsk.ertelecom.ru"
  ["Оренбург"]="st.oren.ertelecom.ru"
  ["Пенза"]="st.penza.ertelecom.ru"
  ["Пермь"]="st.perm.ertelecom.ru"
  ["Ростов на Дону"]="st.rostov.ertelecom.ru"
  ["Рязань"]="st.ryazan.ertelecom.ru"
  ["Самара"]="st.samara.ertelecom.ru"
  ["Сахалин"]="sahst.st.mtsws.net"
  ["Тула"]="st.tula.ertelecom.ru"
  ["Тюмень"]="st.tmn.ertelecom.ru"
  ["Ульяновск"]="st.ulsk.ertelecom.ru"
  ["Уфа"]="st.ufa.ertelecom.ru"
  ["Челябинск"]="st.chel.ertelecom.ru"
  ["Якутск"]="yktst.st.mtsws.net"
  ["Ярославль"]="st.yar.ertelecom.ru"
)

declare -A SERVER_FIXED_PORTS=(
  ["Якутск"]="3333"
  ["Владивосток"]="3333"
  ["Сахалин"]="3333"
  ["Норильск"]="3333"
)

CITY_ORDER=(
  "Москва"
  "Санкт-Петербург"
  "Барнаул"
  "Владивосток"
  "Волгоград"
  "Екатеринбург"
  "Ижевск"
  "Иркутск"
  "Казань"
  "Киров"
  "Красноярск"
  "Курган"
  "Курск"
  "Липецк"
  "Магнитогорск"
  "Набережные Челны"
  "Нижний Новгород"
  "Новосибирск"
  "Норильск"
  "Омск"
  "Оренбург"
  "Пенза"
  "Пермь"
  "Ростов на Дону"
  "Рязань"
  "Самара"
  "Сахалин"
  "Тула"
  "Тюмень"
  "Ульяновск"
  "Уфа"
  "Челябинск"
  "Якутск"
  "Ярославль"
)

check_ping() {
  local host="$1"
  if [[ "$(uname)" == "Darwin" ]]; then
    ping -c 2 -W 2000 "$host" >/dev/null 2>&1 || return 1
  else
    ping -c 2 -W 2 "$host" >/dev/null 2>&1 || return 1
  fi
}

find_working_port() {
  local city="$1"
  local host="$2"
  local ports=()

  if [[ -n "${SERVER_FIXED_PORTS[$city]:-}" ]]; then
    ports=("${SERVER_FIXED_PORTS[$city]}")
  else
    ports=("${DEFAULT_IPERF_PORT_RANGE[@]}")
  fi

  for port in "${ports[@]}"; do
    if timeout "$IPERF_TIMEOUT" iperf3 -c "$host" -p "$port" -n "$TEST_BYTES" >/dev/null 2>&1; then
      echo "$port"
      return 0
    fi
  done

  return 1
}

test_download() {
  local host="$1"
  local port="$2"
  timeout "$IPERF_TIMEOUT" iperf3 -c "$host" -p "$port" -n "$TEST_BYTES" -R >/dev/null 2>&1
}

main() {
  local total_ok=0
  local total_fail=0
  local start_ts end_ts

  start_ts=$(date +%s)

  for cmd in iperf3 ping timeout; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
      echo "Ошибка: утилита $cmd не найдена в PATH."
      exit 1
    fi
  done

  for city in "${CITY_ORDER[@]}"; do
    local host="${SERVERS[$city]}"
    local port=""
    local status=false

    # Проверка ping -> upload -> download
    if check_ping "$host"; then
      port=$(find_working_port "$city" "$host" || true)
      if [[ -n "$port" ]]; then
        if test_download "$host" "$port"; then
          status=true
        fi
      fi
    fi

    if [[ "$status" == "true" ]]; then
      printf "%-22s 🟢\n" "$city"
      ((total_ok++)) || true
    else
      printf "%-22s 🔴\n" "$city"
      ((total_fail++)) || true
    fi
  done

  end_ts=$(date +%s)
  echo
  echo "Тестирование завершено за $((end_ts - start_ts)) сек."
  echo "Доступно: $total_ok | Недоступно: $total_fail"
}

main "$@"