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

SCRIPT="/workspace/digest/vpn-health.sh"
TMP_ROOT="$(mktemp -d)"
trap 'rm -rf "$TMP_ROOT"' EXIT

MOCK_BIN="${TMP_ROOT}/bin"
MOCK_ACTIONS="${TMP_ROOT}/actions.log"
MOCK_VPN_STATE="${TMP_ROOT}/mock-vpn.env"
mkdir -p "$MOCK_BIN"

cat > "${MOCK_BIN}/sg" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
if [[ "${1:-}" != "nordvpn" || "${2:-}" != "-c" ]]; then
  echo "unexpected sg args: $*" >&2
  exit 1
fi
shift 2
printf 'sg %s\n' "$*" >> "$MOCK_ACTIONS"
PATH="$MOCK_BIN:$PATH" bash -c "$*"
EOF

cat > "${MOCK_BIN}/nordvpn" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail

load_state() {
  VPN_STATUS="disconnected"
  DNS_OK=1
  HTTP_OK=1
  CONNECT_FIX_MODE="full"
  RESTART_FIX_MODE="none"
  CONNECT_COUNT=0
  DISCONNECT_COUNT=0
  RESTART_COUNT=0
  REBOOT_COUNT=0
  if [[ -f "$MOCK_VPN_STATE" ]]; then
    # shellcheck disable=SC1090
    source "$MOCK_VPN_STATE"
  fi
}

save_state() {
  cat > "$MOCK_VPN_STATE" <<STATE
VPN_STATUS=${VPN_STATUS}
DNS_OK=${DNS_OK}
HTTP_OK=${HTTP_OK}
CONNECT_FIX_MODE=${CONNECT_FIX_MODE}
RESTART_FIX_MODE=${RESTART_FIX_MODE}
CONNECT_COUNT=${CONNECT_COUNT}
DISCONNECT_COUNT=${DISCONNECT_COUNT}
RESTART_COUNT=${RESTART_COUNT}
REBOOT_COUNT=${REBOOT_COUNT}
STATE
}

load_state
cmd="${1:-}"
shift || true

case "$cmd" in
  status)
    if [[ "$VPN_STATUS" == "connected" ]]; then
      echo "Status: Connected"
    else
      echo "Status: Disconnected"
    fi
    ;;
  connect)
    CONNECT_COUNT=$((CONNECT_COUNT + 1))
    printf 'nordvpn connect %s\n' "$*" >> "$MOCK_ACTIONS"
    if [[ "$CONNECT_FIX_MODE" == "full" ]]; then
      VPN_STATUS="connected"
      DNS_OK=1
      HTTP_OK=1
    elif [[ "$CONNECT_FIX_MODE" == "vpn-only" ]]; then
      VPN_STATUS="connected"
    fi
    save_state
    echo "Connected"
    ;;
  disconnect)
    DISCONNECT_COUNT=$((DISCONNECT_COUNT + 1))
    printf 'nordvpn disconnect\n' >> "$MOCK_ACTIONS"
    VPN_STATUS="disconnected"
    save_state
    echo "Disconnected"
    ;;
  *)
    echo "unexpected nordvpn command: $cmd $*" >&2
    exit 1
    ;;
esac
EOF

cat > "${MOCK_BIN}/getent" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
if [[ -f "$MOCK_VPN_STATE" ]]; then
  # shellcheck disable=SC1090
  source "$MOCK_VPN_STATE"
fi
if [[ "${DNS_OK:-1}" == "1" ]]; then
  echo "1.2.3.4 $*"
  exit 0
fi
exit 2
EOF

cat > "${MOCK_BIN}/curl" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
if [[ -f "$MOCK_VPN_STATE" ]]; then
  # shellcheck disable=SC1090
  source "$MOCK_VPN_STATE"
fi
if [[ "${HTTP_OK:-1}" == "1" ]]; then
  printf '200'
else
  printf '000'
fi
EOF

cat > "${MOCK_BIN}/systemctl" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
if [[ -f "$MOCK_VPN_STATE" ]]; then
  # shellcheck disable=SC1090
  source "$MOCK_VPN_STATE"
fi
printf 'systemctl %s\n' "$*" >> "$MOCK_ACTIONS"
if [[ "${1:-}" == "restart" && "${2:-}" == "nordvpnd" ]]; then
  RESTART_COUNT=$((RESTART_COUNT + 1))
  if [[ "${RESTART_FIX_MODE:-none}" == "full" ]]; then
    VPN_STATUS="connected"
    DNS_OK=1
    HTTP_OK=1
  fi
  cat > "$MOCK_VPN_STATE" <<STATE
VPN_STATUS=${VPN_STATUS}
DNS_OK=${DNS_OK}
HTTP_OK=${HTTP_OK}
CONNECT_FIX_MODE=${CONNECT_FIX_MODE}
RESTART_FIX_MODE=${RESTART_FIX_MODE}
CONNECT_COUNT=${CONNECT_COUNT}
DISCONNECT_COUNT=${DISCONNECT_COUNT}
RESTART_COUNT=${RESTART_COUNT}
REBOOT_COUNT=${REBOOT_COUNT}
STATE
  exit 0
fi
exit 1
EOF

cat > "${MOCK_BIN}/sudo" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
if [[ "${MOCK_SUDO_FAIL:-0}" == "1" ]]; then
  exit 1
fi
if [[ "${1:-}" == "-n" ]]; then
  shift
fi
exec "$@"
EOF

cat > "${MOCK_BIN}/reboot" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
printf 'reboot\n' >> "$MOCK_ACTIONS"
if [[ -f "$MOCK_VPN_STATE" ]]; then
  # shellcheck disable=SC1090
  source "$MOCK_VPN_STATE"
fi
REBOOT_COUNT=$((REBOOT_COUNT + 1))
cat > "$MOCK_VPN_STATE" <<STATE
VPN_STATUS=${VPN_STATUS}
DNS_OK=${DNS_OK}
HTTP_OK=${HTTP_OK}
CONNECT_FIX_MODE=${CONNECT_FIX_MODE}
RESTART_FIX_MODE=${RESTART_FIX_MODE}
CONNECT_COUNT=${CONNECT_COUNT}
DISCONNECT_COUNT=${DISCONNECT_COUNT}
RESTART_COUNT=${RESTART_COUNT}
REBOOT_COUNT=${REBOOT_COUNT}
STATE
exit 0
EOF

cat > "${MOCK_BIN}/sleep" <<'EOF'
#!/usr/bin/env bash
exit 0
EOF

chmod +x "${MOCK_BIN}"/*

write_mock_state() {
  cat > "$MOCK_VPN_STATE" <<EOF
VPN_STATUS=$1
DNS_OK=$2
HTTP_OK=$3
CONNECT_FIX_MODE=$4
RESTART_FIX_MODE=$5
CONNECT_COUNT=0
DISCONNECT_COUNT=0
RESTART_COUNT=0
REBOOT_COUNT=0
EOF
}

state_var() {
  # shellcheck disable=SC1090
  source "$MOCK_VPN_STATE"
  eval "printf '%s' \"\${$1}\""
}

assert_contains() {
  local file="$1"
  local needle="$2"
  if ! grep -Fq "$needle" "$file"; then
    echo "Assertion failed: expected '$needle' in $file" >&2
    exit 1
  fi
}

assert_not_contains() {
  local file="$1"
  local needle="$2"
  if grep -Fq "$needle" "$file"; then
    echo "Assertion failed: did not expect '$needle' in $file" >&2
    exit 1
  fi
}

run_health() {
  local state_file="$1"
  shift
  PATH="$MOCK_BIN:$PATH" \
    MOCK_BIN="$MOCK_BIN" \
    MOCK_ACTIONS="$MOCK_ACTIONS" \
    MOCK_VPN_STATE="$MOCK_VPN_STATE" \
    STATE_FILE="$state_file" \
    SETTLE_SECONDS=0 \
    REBOOT_COOLDOWN_SECONDS=0 \
    "$SCRIPT" "$@"
}

test_healthy() {
  : > "$MOCK_ACTIONS"
  write_mock_state connected 1 1 full none
  local health_state="${TMP_ROOT}/healthy.state"
  run_health "$health_state" ensure
  assert_not_contains "$MOCK_ACTIONS" "nordvpn connect"
}

test_connect_from_disconnected() {
  : > "$MOCK_ACTIONS"
  write_mock_state disconnected 1 1 full none
  local health_state="${TMP_ROOT}/connect.state"
  run_health "$health_state" ensure
  assert_contains "$MOCK_ACTIONS" "nordvpn connect us"
  [[ "$(state_var VPN_STATUS)" == "connected" ]]
}

test_reconnect_from_broken_egress() {
  : > "$MOCK_ACTIONS"
  write_mock_state connected 0 0 full none
  local health_state="${TMP_ROOT}/reconnect.state"
  run_health "$health_state" ensure
  assert_contains "$MOCK_ACTIONS" "nordvpn disconnect"
  assert_contains "$MOCK_ACTIONS" "nordvpn connect us"
}

test_guarded_reboot() {
  : > "$MOCK_ACTIONS"
  write_mock_state connected 0 0 none none
  local health_state="${TMP_ROOT}/reboot.state"

  set +e
  run_health "$health_state" ensure --allow-reboot
  local first_status=$?
  set -e
  [[ "$first_status" -eq 1 ]]
  assert_not_contains "$MOCK_ACTIONS" "reboot"

  set +e
  REBOOT_THRESHOLD=2 run_health "$health_state" ensure --allow-reboot
  local second_status=$?
  set -e
  [[ "$second_status" -eq 2 ]]
  assert_contains "$MOCK_ACTIONS" "reboot"
}

test_healthy
test_connect_from_disconnected
test_reconnect_from_broken_egress
test_guarded_reboot

echo "test-vpn-health.sh: PASS"
