#!/usr/bin/env bash

#   ./kill_robot.sh                 # kill ROS 2 processes from this workspace
#   ./kill_robot.sh -A              # kill all ROS 2 processes system-wide (dangerous)
#   ./kill_robot.sh -n              # dry run (show PIDs only)
#   ./kill_robot.sh -y              # no confirmation
#   ./kill_robot.sh -v              # verbose logs

set -Eeuo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"
WS_DIR="${WS_DIR:-${SCRIPT_DIR}}"
SCOPE_ALL=false
DRY_RUN=false
ASSUME_YES=false
VERBOSE=false

log() { echo "[kill_ros2] $*"; }
vlog() { $VERBOSE && echo "[kill_ros2][debug] $*" || true; }

usage() {
  cat <<EOF
Usage: $(basename "$0") [options]
Options:
  -A        Kill all ROS 2 processes system-wide (not only this workspace)
  -n        Dry run (list candidate PIDs and commands)
  -y        Assume yes (no confirmation prompt)
  -v        Verbose output
  -h        Show this help and exit
Env:
  WS_DIR    Workspace directory (default: script directory)
EOF
}

while getopts ":Anyvh" opt; do
  case $opt in
    A) SCOPE_ALL=true ;;
    n) DRY_RUN=true ;;
    y) ASSUME_YES=true ;;
    v) VERBOSE=true ;;
    h) usage; exit 0 ;;
    :) echo "Option -$OPTARG requires an argument" >&2; exit 2 ;;
    \?) echo "Unknown option: -$OPTARG" >&2; usage; exit 2 ;;
  esac
done

# Return 0 if PID is alive, else 1
is_alive() { kill -0 "$1" 2>/dev/null; }

# Check if a PID looks like a ROS process by inspecting its environment
is_ros_env() {
  local pid=$1
  [[ -r "/proc/${pid}/environ" ]] || return 1
  tr '\0' '\n' <"/proc/${pid}/environ" | grep -Eq '^(ROS_DISTRO=|RMW_IMPLEMENTATION=)'
}

# List candidate PIDs and their commands (tab-separated)
list_candidates() {
  local ps_out
  # Use full command line; exclude our own shell/grep processes later
  ps_out=$(ps -eo pid=,cmd=)
  while IFS= read -r line; do
    local pid cmd
    pid=$(awk '{print $1}' <<<"$line")
    cmd=${line#* } # everything after first space
    # Skip self
    [[ "$pid" -eq $$ || "$pid" -eq $PPID ]] && continue

    # Scope filter
    local in_scope=false
    if $SCOPE_ALL; then
      in_scope=true
    else
      # Focus on processes started from this workspace or via ros2 run/launch
      if grep -qE "/install/.*/lib/" <<<"$cmd" ||
         grep -qE "\bros2 (run|launch)\b" <<<"$cmd" ||
         grep -qE "python(3)? .*launch" <<<"$cmd"; then
        # If workspace path is known, prefer matches containing it
        if grep -q "${WS_DIR}" <<<"$cmd" || $SCOPE_ALL; then in_scope=true; else in_scope=true; fi
      fi
    fi

    $in_scope || continue

    # Must look like a ROS process (has ROS env vars)
    if is_ros_env "$pid"; then
      printf "%s\t%s\n" "$pid" "$cmd"
    else
      vlog "skip pid=$pid (no ROS env) cmd=$cmd"
    fi
  done <<<"$ps_out"
}

confirm() {
  $ASSUME_YES && return 0
  read -r -p "Kill the above processes? [y/N] " ans || true
  [[ "${ans,,}" == "y" || "${ans,,}" == "yes" ]]
}

kill_with_escalation() {
  local pids=("$@")
  local stages=(INT TERM KILL)
  local stage
  for stage in "${stages[@]}"; do
    local alive=()
    for pid in "${pids[@]}"; do
      if is_alive "$pid"; then alive+=("$pid"); fi
    done
    ((${#alive[@]}==0)) && return 0

    log "Sending SIG${stage} to ${#alive[@]} process(es): ${alive[*]}"
    kill -s "$stage" -- "${alive[@]}" 2>/dev/null || true

    # Wait up to 3s for this stage
    for _ in {1..30}; do
      sleep 0.1
      local still=()
      for pid in "${alive[@]}"; do is_alive "$pid" && still+=("$pid"); done
      ((${#still[@]}==0)) && break
      alive=("${still[@]}")
    done
  done
}

main() {
  log "Workspace: ${WS_DIR} | Scope: $($SCOPE_ALL && echo all || echo workspace)"
  local candidates
  candidates=$(list_candidates || true)
  if [[ -z "$candidates" ]]; then
    log "No ROS 2 processes found to kill."
    exit 0
  fi

  echo "Candidates (PID\tCMD):"
  echo "$candidates" | sed 's/^/  /'

  if $DRY_RUN; then
    log "Dry run only. Exiting."
    exit 0
  fi

  if ! confirm; then
    log "Cancelled by user."
    exit 1
  fi

  # Extract PIDs and kill
  mapfile -t pids < <(awk -F '\t' '{print $1}' <<<"$candidates" | sort -u)
  kill_with_escalation "${pids[@]}"

  # Optionally stop ros2 daemon so it doesn't hold stale graph
  if command -v ros2 >/dev/null 2>&1; then
    ros2 daemon stop >/dev/null 2>&1 || true
  fi

  log "Done."
}

main "$@"