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

die() {
  echo "error: $*" >&2
  exit 2
}

usage() {
  cat <<'USAGE'
Usage:
  scripts/rerun/stop-external-run.sh [options]

Stops an external TiPOC pressure-run launch by reading wrapper pid files from:

  <remote-root>/logs/<launch-id>/<node-id>.pid

Options can be supplied as flags or matching environment variables.

Options:
  --launch-id ID          TIPOC_RUN_LAUNCH_ID
  --remote-root PATH      TIPOC_RUN_REMOTE_ROOT
  --node-file PATH        TIPOC_RUN_NODE_FILE
  --ssh-key PATH          TIPOC_RUN_SSH_KEY
  --manifest PATH         TIPOC_RUN_MANIFEST, optional stray tipoc run cleanup
  --ssh-user USER         TIPOC_RUN_SSH_USER, default ec2-user
  --bastion TARGET        TIPOC_RUN_BASTION, required for proxy nodes
  --help

Node inventory format:
  node_id ssh_host ssh_mode private_ip expected_slots expected_sql_pool

ssh_mode must be direct or proxy. The private_ip and expected fields are parsed
for compatibility with start-run-template.sh and are not used by this stop step.
USAGE
}

shell_quote() {
  local value="$1"
  printf "'%s'" "${value//\'/\'\"\'\"\'}"
}

proxy_command_quote() {
  local value="$1"
  if [[ "$value" =~ ^[A-Za-z0-9_./:@%+=,-]+$ ]]; then
    printf '%s' "$value"
    return
  fi
  value="${value//\\/\\\\}"
  value="${value//\"/\\\"}"
  value="${value//\$/\\$}"
  value="${value//\`/\\\`}"
  printf '"%s"' "$value"
}

redact_stream() {
  sed -E \
    -e 's#(root:)[^@[:space:]]+(@tcp)#\1<redacted>\2#g' \
    -e 's#([?&](access-key|secret-access-key|session-token|X-Amz-[A-Za-z0-9_-]+)=)[^&[:space:]"'"'"']+#\1<redacted>#g' \
    -e 's#((password|passwd|pwd|secret|token|access_key|secret_key|session_token)=)[^&[:space:]"'"'"']+#\1<redacted>#Ig' \
    -e 's#(TIDB_DSN=).*#\1<redacted>#g'
}

require_file() {
  local path="$1" label="$2"
  [[ -n "$path" ]] || die "$label is required"
  [[ -s "$path" ]] || die "$label is missing or empty: $path"
}

require_path_value() {
  local value="$1" label="$2"
  [[ -n "$value" ]] || die "$label is required"
}

TIPOC_RUN_LAUNCH_ID="${TIPOC_RUN_LAUNCH_ID:-}"
TIPOC_RUN_REMOTE_ROOT="${TIPOC_RUN_REMOTE_ROOT:-}"
TIPOC_RUN_NODE_FILE="${TIPOC_RUN_NODE_FILE:-}"
TIPOC_RUN_SSH_KEY="${TIPOC_RUN_SSH_KEY:-}"
TIPOC_RUN_MANIFEST="${TIPOC_RUN_MANIFEST:-}"
TIPOC_RUN_SSH_USER="${TIPOC_RUN_SSH_USER:-ec2-user}"
TIPOC_RUN_BASTION="${TIPOC_RUN_BASTION:-}"

while [[ "$#" -gt 0 ]]; do
  case "$1" in
    --launch-id)
      [[ "$#" -ge 2 ]] || die "--launch-id requires a value"
      TIPOC_RUN_LAUNCH_ID="$2"
      shift 2
      ;;
    --remote-root)
      [[ "$#" -ge 2 ]] || die "--remote-root requires a value"
      TIPOC_RUN_REMOTE_ROOT="$2"
      shift 2
      ;;
    --node-file)
      [[ "$#" -ge 2 ]] || die "--node-file requires a value"
      TIPOC_RUN_NODE_FILE="$2"
      shift 2
      ;;
    --ssh-key)
      [[ "$#" -ge 2 ]] || die "--ssh-key requires a value"
      TIPOC_RUN_SSH_KEY="$2"
      shift 2
      ;;
    --manifest)
      [[ "$#" -ge 2 ]] || die "--manifest requires a value"
      TIPOC_RUN_MANIFEST="$2"
      shift 2
      ;;
    --ssh-user)
      [[ "$#" -ge 2 ]] || die "--ssh-user requires a value"
      TIPOC_RUN_SSH_USER="$2"
      shift 2
      ;;
    --bastion)
      [[ "$#" -ge 2 ]] || die "--bastion requires a value"
      TIPOC_RUN_BASTION="$2"
      shift 2
      ;;
    --help)
      usage
      exit 0
      ;;
    *)
      die "unknown option: $1"
      ;;
  esac
done

require_path_value "$TIPOC_RUN_LAUNCH_ID" "launch id"
require_path_value "$TIPOC_RUN_REMOTE_ROOT" "remote root"
require_file "$TIPOC_RUN_NODE_FILE" "node inventory file"
require_file "$TIPOC_RUN_SSH_KEY" "SSH private key"
[[ -n "$TIPOC_RUN_SSH_USER" ]] || die "ssh user is required"

LOG_DIR="${TIPOC_RUN_REMOTE_ROOT}/logs/${TIPOC_RUN_LAUNCH_ID}"
SSH_BASE=(-i "$TIPOC_RUN_SSH_KEY" -o IdentitiesOnly=yes -o StrictHostKeyChecking=no -o BatchMode=yes)

ssh_target_for() {
  local host="$1"
  printf '%s@%s\n' "$TIPOC_RUN_SSH_USER" "$host"
}

bastion_target() {
  if [[ "$TIPOC_RUN_BASTION" == *@* ]]; then
    printf '%s\n' "$TIPOC_RUN_BASTION"
  else
    printf '%s@%s\n' "$TIPOC_RUN_SSH_USER" "$TIPOC_RUN_BASTION"
  fi
}

ssh_node() {
  local host="$1" mode="$2"
  shift 2
  local target rc
  target="$(ssh_target_for "$host")"
  if [[ "$mode" == "direct" ]]; then
    ssh "${SSH_BASE[@]}" "$target" "$@" 2>&1 | redact_stream
    rc="${PIPESTATUS[0]}"
  elif [[ "$mode" == "proxy" ]]; then
    [[ -n "$TIPOC_RUN_BASTION" ]] || die "TIPOC_RUN_BASTION or --bastion is required for proxy nodes"
    local proxy_cmd="ssh"
    local arg
    for arg in "${SSH_BASE[@]}"; do
      proxy_cmd+=" $(proxy_command_quote "$arg")"
    done
    proxy_cmd+=" -W %h:%p $(proxy_command_quote "$(bastion_target)")"
    ssh "${SSH_BASE[@]}" -o "ProxyCommand=$proxy_cmd" "$target" "$@" 2>&1 | redact_stream
    rc="${PIPESTATUS[0]}"
  else
    die "invalid ssh_mode for $host: $mode"
  fi
  return "$rc"
}

stop_node() {
  local node="$1" host="$2" mode="$3"
  local node_q log_dir_q manifest_q
  node_q="$(shell_quote "$node")"
  log_dir_q="$(shell_quote "$LOG_DIR")"
  manifest_q="$(shell_quote "$TIPOC_RUN_MANIFEST")"
  ssh_node "$host" "$mode" \
    "NODE=${node_q} LOG_DIR=${log_dir_q} MANIFEST=${manifest_q} bash -s" <<'SH'
set -euo pipefail

PID_FILE="${LOG_DIR}/${NODE}.pid"
STATUS_FILE="${LOG_DIR}/${NODE}.status"
stopped=0

if [[ ! -d "$LOG_DIR" ]]; then
  echo "${NODE} log directory is missing: ${LOG_DIR}" >&2
  exit 66
fi

kill_pid_tree() {
  local pid="$1"
  [[ "$pid" =~ ^[0-9]+$ ]] || return 0
  if ! kill -0 "$pid" >/dev/null 2>&1; then
    return 0
  fi

  local children child
  children="$(pgrep -P "$pid" 2>/dev/null || true)"
  for child in $children; do
    kill_pid_tree "$child"
  done

  kill "$pid" >/dev/null 2>&1 || true
  sleep 1
  if kill -0 "$pid" >/dev/null 2>&1; then
    kill -KILL "$pid" >/dev/null 2>&1 || true
  fi
}

if [[ -s "$PID_FILE" ]]; then
  pid="$(tr -d '[:space:]' < "$PID_FILE")"
  if [[ "$pid" =~ ^[0-9]+$ ]]; then
    if kill -0 "$pid" >/dev/null 2>&1; then
      kill_pid_tree "$pid"
      stopped=1
    fi
  else
    echo "${NODE} invalid pid file: ${PID_FILE}" >&2
    exit 65
  fi
fi

if [[ -n "$MANIFEST" ]]; then
  while read -r pid; do
    [[ -n "$pid" ]] || continue
    kill_pid_tree "$pid"
    stopped=1
  done < <(
    ps -eo pid=,args= |
      awk -v manifest="$MANIFEST" '
        $0 ~ /[t]ipoc run/ &&
        (index($0, "--manifest " manifest) || index($0, "--manifest=" manifest)) {
          print $1
        }'
  )
fi

if [[ "$stopped" -eq 1 ]]; then
  printf 'stopped\n' > "$STATUS_FILE"
else
  printf 'not_running\n' > "$STATUS_FILE"
fi
printf '%s cleanup=%s\n' "$NODE" "$(cat "$STATUS_FILE")"
SH
}

while read -r node host mode private_ip expected_slots expected_sql_pool extra; do
  [[ -z "${node:-}" || "${node:0:1}" == "#" ]] && continue
  [[ -z "${extra:-}" ]] || die "node inventory record has too many fields: $node"
  [[ -n "${host:-}" && -n "${mode:-}" && -n "${private_ip:-}" && -n "${expected_slots:-}" && -n "${expected_sql_pool:-}" ]] ||
    die "node inventory record has too few fields: $node"
  [[ "$mode" == "direct" || "$mode" == "proxy" ]] || die "node $node has invalid ssh_mode: $mode"
  stop_node "$node" "$host" "$mode"
done < "$TIPOC_RUN_NODE_FILE"