#!/usr/bin/env bash
# Manage a Cloudflare WARP (usque) tunnel: connect, disconnect,
# or report Waybar status.

set -euo pipefail

info() { printf '\033[1;34m==>\033[0m %s\n' "$*"; }
warn() { printf '\033[1;33m==> warning:\033[0m %s\n' "$*" >&2; }
error() { printf '\033[1;31m==> error:\033[0m %s\n' "$*" >&2; }
die() {
  error "$*"
  exit 1
}

CONFIG_DIR="$HOME/.cache/usque"
CONFIG="$CONFIG_DIR/config.json"
RUNTIME_DIR="${XDG_RUNTIME_DIR:-/tmp}"
PID_FILE="$RUNTIME_DIR/usque-warp.pid"
STATE_FILE="$RUNTIME_DIR/usque-warp.state"
IFACE_FILE="$RUNTIME_DIR/usque-warp.iface"
LOG_FILE="$RUNTIME_DIR/usque-warp.log"

list_tun_ifaces() {
  ip -o link show 2>/dev/null | awk -F': ' '{print $2}' | grep -E '^tun[0-9]+$' || true
}

detect_iface() {
  if [[ -f $IFACE_FILE ]]; then
    cat "$IFACE_FILE"
    return
  fi
  list_tun_ifaces | head -n1
}

# True if the PID in PID_FILE is a live usque process. Uses /proc rather
# than `sudo kill -0` so it never prompts for a password (matters for
# Waybar's status polling) and can't be fooled by PID reuse.
is_running() {
  [[ -f $PID_FILE ]] || return 1
  local pid
  pid=$(cat "$PID_FILE" 2>/dev/null || true)
  [[ $pid =~ ^[0-9]+$ ]] || return 1
  [[ $(cat "/proc/$pid/comm" 2>/dev/null || true) == usque ]]
}

# True if ANY usque process exists, tracked by the PID file or not.
any_usque_running() {
  pgrep -x usque >/dev/null 2>&1
}

# Stop every usque process (tracked or orphaned): SIGTERM, wait, then SIGKILL.
stop_usque() {
  if any_usque_running; then
    info "Stopping usque..."
    sudo pkill -x usque 2>/dev/null || true
    for _ in {1..25}; do
      any_usque_running || break
      sleep 0.2
    done
    if any_usque_running; then
      warn "usque did not stop on SIGTERM, forcing kill..."
      sudo pkill -9 -x usque 2>/dev/null || true
      sleep 0.5
    fi
  fi
  rm -f "$PID_FILE"
}

ensure_config() {
  # Only register when there is no config; registering on every run
  # creates a brand-new WARP account each time.
  if [[ -f $CONFIG ]]; then
    return
  fi
  info "Creating $CONFIG_DIR..."
  mkdir -p "$CONFIG_DIR"
  info "Registering Cloudflare WARP account..."
  usque -c "$CONFIG" register < <(yes)
  [[ -f $CONFIG ]] || die "Failed to create config file: $CONFIG"
  info "Config created successfully."
}

remove_tun_default_routes() {
  local dev="$1" route
  while ip route show | grep -qE "^default .*dev $dev"; do
    route=$(ip route show | grep -E "^default .*dev $dev" | head -n1)
    info "Removing route: $route"
    # Unquoted on purpose: the route must be split into separate arguments.
    # shellcheck disable=SC2086
    sudo ip route del $route || break
  done
}

CONNECT_OK=false

# EXIT trap for connect: if we leave for any reason (die, a failing command,
# Ctrl-C) before reaching "Connected.", tear everything down again.
cleanup_failed_connect() {
  local rc=$?
  trap - EXIT INT TERM
  if [[ $CONNECT_OK != true ]]; then
    warn "connect failed, cleaning up..."
    disconnect || true
  fi
  exit "$rc"
}

connect() {
  sudo -v

  # Refuse to start if any usque is already running, not just a tracked one.
  # This runs before the cleanup trap is armed, so we never kill a tunnel
  # we didn't start.
  if any_usque_running; then
    die "usque is already running (run '$0 disconnect' first)"
  fi
  rm -f "$PID_FILE" "$LOG_FILE" "$STATE_FILE" "$IFACE_FILE"

  ensure_config

  trap cleanup_failed_connect EXIT
  trap 'exit 130' INT TERM

  info "Saving current default route..."
  DEFAULT_ROUTE=$(ip route show default | grep -vE 'dev tun[0-9]+' | head -n1 || true)
  if [[ -z $DEFAULT_ROUTE ]]; then
    die "Could not determine current default route"
  fi
  GATEWAY=$(echo "$DEFAULT_ROUTE" | awk '{for(i=1;i<=NF;i++) if($i=="via") print $(i+1)}')
  INTERFACE=$(echo "$DEFAULT_ROUTE" | awk '{for(i=1;i<=NF;i++) if($i=="dev") print $(i+1)}')
  if [[ -z $GATEWAY || -z $INTERFACE ]]; then
    die "Cannot determine gateway/interface from: $DEFAULT_ROUTE"
  fi
  echo "$DEFAULT_ROUTE" >"$STATE_FILE"

  info "Recording pre-existing tun interfaces..."
  BEFORE_IFACES=$(list_tun_ifaces)

  info "Starting usque..."
  # sudo does not affect redirects: the outer shell would open "$LOG_FILE"
  # as the unprivileged user. Run the redirection inside sudo (via sh) so the
  # log is opened as root. $$ inside sh is usque's PID (sh exec's usque, keeping
  # the PID stable regardless of how sudo forks internally), written up front.
  sudo sh -c 'echo $$ > "$1"; exec usque nativetun -c "$2" > "$3" 2>&1' \
    _ "$PID_FILE" "$CONFIG" "$LOG_FILE" &

  info "Waiting for MASQUE connection..."
  MASQUE_IP=""
  for _ in {1..30}; do
    # head -n1: usque may log this line more than once; we need exactly one IP.
    MASQUE_IP=$(grep -oP 'MASQUE connection to \K[0-9.]+(?=:443)' "$LOG_FILE" 2>/dev/null | head -n1 || true)
    [[ -n $MASQUE_IP ]] && break
    if [[ -f $PID_FILE ]] && ! is_running; then
      die "usque exited unexpectedly; see $LOG_FILE"
    fi
    sleep 1
  done
  [[ -n $MASQUE_IP ]] || die "Failed to detect MASQUE endpoint; see $LOG_FILE"
  [[ $MASQUE_IP =~ ^[0-9]+(\.[0-9]+){3}$ ]] || die "Unexpected MASQUE endpoint: '$MASQUE_IP'"

  info "Waiting for usque interface..."
  TUN_DEV=""
  for _ in {1..30}; do
    AFTER_IFACES=$(list_tun_ifaces)
    TUN_DEV=$(comm -13 <(echo "$BEFORE_IFACES" | sort) <(echo "$AFTER_IFACES" | sort) | head -n1)
    [[ -n $TUN_DEV ]] && break
    sleep 1
  done
  [[ -n $TUN_DEV ]] || die "Failed to detect usque interface"
  echo "$TUN_DEV" >"$IFACE_FILE"
  info "Detected interface: $TUN_DEV"

  echo "MASQUE_IP=$MASQUE_IP GATEWAY=$GATEWAY INTERFACE=$INTERFACE" >>"$STATE_FILE"

  info "Allowing MASQUE endpoint outside tunnel..."
  sudo ip route replace "$MASQUE_IP" via "$GATEWAY" dev "$INTERFACE"

  info "Removing old tun routes..."
  remove_tun_default_routes "$TUN_DEV"

  info "Switching default route to $TUN_DEV..."
  sudo ip route add default dev "$TUN_DEV" metric 1

  CONNECT_OK=true
  trap - EXIT INT TERM
  info "Connected."
}

disconnect() {
  sudo -v
  info "Disconnecting..."

  local dev=""
  [[ -f $IFACE_FILE ]] && dev=$(cat "$IFACE_FILE")

  # Kill usque FIRST so the kernel tears down the tun device (and every
  # route bound to it) in one step, instead of us racing it by pulling
  # routes out from under a device that's still up. This also kills any
  # orphaned usque processes the PID file doesn't know about.
  stop_usque

  # Give the kernel a moment to remove the device; only touch routes
  # manually if it is not disappearing on its own.
  if [[ -n $dev ]]; then
    info "Waiting for tunnel interface to go down..."
    for _ in {1..25}; do
      list_tun_ifaces | grep -qx "$dev" || break
      sleep 0.2
    done
    if list_tun_ifaces | grep -qx "$dev"; then
      warn "interface $dev is still up, cleaning its routes manually..."
      sudo ip route flush dev "$dev" 2>/dev/null || true
      remove_tun_default_routes "$dev"
    fi
  fi

  if [[ -f $STATE_FILE ]]; then
    local masque_ip original_default
    masque_ip=$(grep -oP 'MASQUE_IP=\K[0-9.]+' "$STATE_FILE" | head -n1 || true)
    if [[ -n $masque_ip ]]; then
      info "Removing MASQUE route: $masque_ip"
      sudo ip route del "$masque_ip" 2>/dev/null || true
    fi

    # Explicitly restore the pre-connect default route rather than
    # assuming it's still intact. 'replace' is idempotent.
    original_default=$(head -n1 "$STATE_FILE")
    if [[ $original_default == default* ]]; then
      info "Restoring original default route..."
      # Unquoted on purpose: the route must be split into separate arguments.
      # shellcheck disable=SC2086
      sudo ip route replace $original_default ||
        warn "could not restore original default route"
    fi
  elif [[ -z $(ip route show default) ]]; then
    warn "no default route and no saved state; restore it manually (ip route add default via <gw> dev <if>)"
  fi

  rm -f "$STATE_FILE" "$IFACE_FILE"
  info "Disconnected."
}

status() {
  local iface running=false
  iface=$(detect_iface)
  is_running && running=true

  if [[ -n $iface && $running == true ]]; then
    printf '{"text":"%s","tooltip":"WARP connected via %s","class":"connected"}\n' "$iface" "$iface"
  elif [[ -n $iface ]]; then
    printf '{"text":"%s","tooltip":"Interface %s up, but usque-warp process not tracked","class":"connected"}\n' "$iface" "$iface"
  elif [[ $running == true ]]; then
    printf '{"text":"connecting","tooltip":"usque starting...","class":"connecting"}\n'
  else
    printf '{"text":"","tooltip":"WARP disconnected","class":"disconnected"}\n'
  fi
}

case "${1:-}" in
connect)
  connect
  ;;
disconnect)
  disconnect
  ;;
status)
  status
  ;;
*)
  echo "Usage: $0 <connect|disconnect|status>"
  echo ""
  echo "Commands:"
  echo "  connect       Start the WARP tunnel"
  echo "  disconnect    Stop the WARP tunnel and restore routes"
  echo "  status        Print Waybar status JSON"
  exit 1
  ;;
esac