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

MODE="ensure"
ALLOW_REBOOT=0
QUIET=0
CONNECT_TARGET="${CONNECT_TARGET:-us}"
STATE_FILE="${STATE_FILE:-/tmp/digest-vpn-health.state}"
LOG_FILE="${LOG_FILE:-}"
REBOOT_THRESHOLD="${REBOOT_THRESHOLD:-3}"
REBOOT_COOLDOWN_SECONDS="${REBOOT_COOLDOWN_SECONDS:-14400}"
SETTLE_SECONDS="${SETTLE_SECONDS:-5}"
CONNECT_TIMEOUT_SECONDS="${CONNECT_TIMEOUT_SECONDS:-5}"
MAX_TIME_SECONDS="${MAX_TIME_SECONDS:-10}"
USER_AGENT="${USER_AGENT:-Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36}"

usage() {
  cat <<'EOF'
Usage:
  vpn-health.sh [check|ensure] [options]

Modes:
  check                        Only validate VPN + connectivity, no remediation.
  ensure                       Validate and repair if needed (default).

Options:
  --allow-reboot               Allow a guarded reboot after repeated failures.
  --connect-target TARGET      NordVPN location target. Default: us
  --state-file PATH            Persistent state file for failure counters.
  --log-file PATH              Append logs to PATH.
  --quiet                      Suppress stdout logs.
  -h, --help                   Show help.
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    check|ensure)
      MODE="$1"
      shift
      ;;
    --allow-reboot)
      ALLOW_REBOOT=1
      shift
      ;;
    --connect-target)
      CONNECT_TARGET="$2"
      shift 2
      ;;
    --state-file)
      STATE_FILE="$2"
      shift 2
      ;;
    --log-file)
      LOG_FILE="$2"
      shift 2
      ;;
    --quiet)
      QUIET=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

timestamp() {
  date -Is
}

log_line() {
  local level="$1"
  shift
  local line="[$(timestamp)] [$level] $*"

  if [[ "$QUIET" -eq 0 ]]; then
    echo "$line"
  fi

  if [[ -n "$LOG_FILE" ]]; then
    mkdir -p "$(dirname "$LOG_FILE")"
    echo "$line" >> "$LOG_FILE"
  fi
}

info() {
  log_line "INFO" "$@"
}

warn() {
  log_line "WARN" "$@"
}

error() {
  log_line "ERROR" "$@"
}

save_state() {
  mkdir -p "$(dirname "$STATE_FILE")"
  cat > "$STATE_FILE" <<EOF
failure_count=${failure_count}
last_reboot_at=${last_reboot_at}
last_failure_reason=$(printf '%q' "${last_failure_reason}")
EOF
  chmod 666 "$STATE_FILE" 2>/dev/null || true
}

load_state() {
  failure_count=0
  last_reboot_at=0
  last_failure_reason=""

  if [[ -f "$STATE_FILE" ]]; then
    # shellcheck disable=SC1090
    source "$STATE_FILE"
  fi
}

reset_state() {
  failure_count=0
  last_failure_reason=""
  save_state
}

record_failure() {
  local reason="$1"
  failure_count=$((failure_count + 1))
  last_failure_reason="$reason"
  save_state
}

run_nordvpn() {
  local shell_cmd="nordvpn"
  local arg

  for arg in "$@"; do
    shell_cmd+=" $(printf '%q' "$arg")"
  done

  sg nordvpn -c "$shell_cmd"
}

have_privileged_access() {
  if [[ "${EUID}" -eq 0 ]]; then
    return 0
  fi

  sudo -n true >/dev/null 2>&1
}

run_privileged() {
  if [[ "${EUID}" -eq 0 ]]; then
    "$@"
    return $?
  fi

  sudo -n "$@"
}

settle() {
  if [[ "$SETTLE_SECONDS" -gt 0 ]]; then
    sleep "$SETTLE_SECONDS"
  fi
}

nordvpn_status_output() {
  run_nordvpn status 2>&1 || true
}

vpn_is_connected() {
  local output
  output="$(nordvpn_status_output)"
  [[ "$output" == *"Status: Connected"* ]]
}

probe_dns_host() {
  local host="$1"

  if getent ahosts "$host" >/dev/null 2>&1; then
    return 0
  fi

  getent hosts "$host" >/dev/null 2>&1
}

probe_http_url() {
  local url="$1"
  local status

  status="$(
    curl -sS -o /dev/null -w "%{http_code}" \
      -A "$USER_AGENT" \
      --connect-timeout "$CONNECT_TIMEOUT_SECONDS" \
      --max-time "$MAX_TIME_SECONDS" \
      -L \
      "$url" 2>/dev/null || true
  )"

  [[ -n "$status" && "$status" != "000" ]]
}

LAST_HEALTH_REASON=""

health_check() {
  local -a failures=()
  local output

  output="$(nordvpn_status_output)"
  if [[ "$output" != *"Status: Connected"* ]]; then
    failures+=("NordVPN is not connected")
  fi

  local host
  for host in api.openai.com www.reddit.com api.telegram.org; do
    if ! probe_dns_host "$host"; then
      failures+=("DNS probe failed for ${host}")
    fi
  done

  local url
  for url in https://api.openai.com/ https://www.reddit.com/ https://api.telegram.org/; do
    if ! probe_http_url "$url"; then
      failures+=("HTTP probe failed for ${url}")
    fi
  done

  if [[ "${#failures[@]}" -eq 0 ]]; then
    LAST_HEALTH_REASON=""
    return 0
  fi

  LAST_HEALTH_REASON="$(printf '%s; ' "${failures[@]}")"
  LAST_HEALTH_REASON="${LAST_HEALTH_REASON%'; '}"
  warn "Health check failed: ${LAST_HEALTH_REASON}"
  warn "NordVPN status output: ${output//$'\n'/ | }"
  return 1
}

attempt_connect() {
  info "Attempting NordVPN connect to ${CONNECT_TARGET}"
  run_nordvpn connect "$CONNECT_TARGET" >/dev/null
  settle
}

attempt_reconnect() {
  info "Attempting NordVPN reconnect"
  run_nordvpn disconnect >/dev/null 2>&1 || true
  settle
  attempt_connect
}

attempt_daemon_restart_repair() {
  if ! have_privileged_access; then
    warn "Privileged repair unavailable: cannot restart nordvpnd from this context"
    return 1
  fi

  info "Restarting nordvpnd"
  run_privileged systemctl restart nordvpnd >/dev/null
  settle
  attempt_connect
}

should_reboot_now() {
  local now
  now="$(date +%s)"

  if (( failure_count < REBOOT_THRESHOLD )); then
    return 1
  fi

  if (( now - last_reboot_at < REBOOT_COOLDOWN_SECONDS )); then
    return 1
  fi

  return 0
}

trigger_reboot() {
  local now

  if ! have_privileged_access; then
    warn "Reboot requested but privileged access is unavailable"
    return 1
  fi

  now="$(date +%s)"
  last_reboot_at="$now"
  save_state

  error "Triggering reboot after ${failure_count} consecutive failed health repairs"
  run_privileged reboot
}

ensure_health() {
  load_state

  if health_check; then
    info "VPN and outbound connectivity are healthy"
    reset_state
    return 0
  fi

  local initial_reason="$LAST_HEALTH_REASON"

  if vpn_is_connected; then
    attempt_reconnect || true
  else
    attempt_connect || true
  fi

  if health_check; then
    info "Connectivity recovered after NordVPN reconnect"
    reset_state
    return 0
  fi

  attempt_daemon_restart_repair || true

  if health_check; then
    info "Connectivity recovered after restarting nordvpnd"
    reset_state
    return 0
  fi

  record_failure "$LAST_HEALTH_REASON"
  error "Health remediation failed: ${LAST_HEALTH_REASON}"

  if [[ "$ALLOW_REBOOT" -eq 1 ]] && should_reboot_now; then
    trigger_reboot
    return 2
  fi

  warn "Skipping reboot. Failure count=${failure_count}, threshold=${REBOOT_THRESHOLD}, allow_reboot=${ALLOW_REBOOT}"
  warn "Initial failure was: ${initial_reason}"
  return 1
}

main() {
  case "$MODE" in
    check)
      if health_check; then
        info "Check-only mode: healthy"
        exit 0
      fi
      exit 1
      ;;
    ensure)
      ensure_health
      exit $?
      ;;
    *)
      echo "Unsupported mode: $MODE" >&2
      exit 2
      ;;
  esac
}

main
