#!/usr/bin/env bash

process_identity_token() {
  local pid="$1"
  ps -o lstart= -p "$pid" 2>/dev/null | xargs || true
}

same_process_identity() {
  local pid="$1"
  local expected_token="$2"
  local current_token

  current_token="$(process_identity_token "$pid")"
  [[ -n "$current_token" && "$current_token" == "$expected_token" ]]
}

snapshot_process_tree() {
  local pid="$1"
  local token
  local child
  local child_snapshot
  local descendants=""

  [[ "$pid" =~ ^[0-9]+$ ]] || return 0
  token="$(process_identity_token "$pid")"
  [[ -n "$token" ]] || return 0

  while IFS= read -r child; do
    if [[ -n "$child" ]]; then
      child_snapshot="$(snapshot_process_tree "$child")"
      if [[ -n "$child_snapshot" ]]; then
        descendants="${descendants}${child_snapshot}"$'\n'
      fi
    fi
  done < <(pgrep -P "$pid" 2>/dev/null || true)

  same_process_identity "$pid" "$token" || return 0
  printf '%s' "$descendants"
  printf '%s|%s\n' "$pid" "$token"
}

signal_process_snapshot() {
  local snapshot="$1"
  local signal="${2:-TERM}"
  local pid
  local token

  [[ -n "$snapshot" ]] || return 0
  while IFS='|' read -r pid token; do
    if [[ -n "$pid" ]] && same_process_identity "$pid" "$token"; then
      kill "-$signal" "$pid" >/dev/null 2>&1 || true
    fi
  done <<<"$snapshot"
}

process_snapshot_alive() {
  local snapshot="$1"
  local pid
  local token

  [[ -n "$snapshot" ]] || return 1
  while IFS='|' read -r pid token; do
    if [[ -n "$pid" ]] && same_process_identity "$pid" "$token"; then
      return 0
    fi
  done <<<"$snapshot"
  return 1
}

stop_process_tree() {
  local root_pid="$1"
  local wait_loops="${PROCESS_TREE_TERM_WAIT_LOOPS:-20}"
  local snapshot
  local descendants
  local current
  local all_snapshots

  [[ "$root_pid" =~ ^[0-9]+$ ]] || return 0
  snapshot="$(snapshot_process_tree "$root_pid")"
  [[ -n "$snapshot" ]] || return 0
  all_snapshots="$snapshot"

  # Stop the first child snapshot before the wrapper so it cannot orphan the
  # already-known Java/Vite process when the wrapper exits.
  descendants="$(printf '%s\n' "$snapshot" | sed '$d')"
  signal_process_snapshot "$descendants" TERM

  current="$(snapshot_process_tree "$root_pid")"
  if [[ -n "$current" ]]; then
    all_snapshots="${all_snapshots}"$'\n'"${current}"
    signal_process_snapshot "$current" TERM
  fi

  for _ in $(seq 1 "$wait_loops"); do
    current="$(snapshot_process_tree "$root_pid")"
    if [[ -n "$current" ]]; then
      all_snapshots="${all_snapshots}"$'\n'"${current}"
      signal_process_snapshot "$current" TERM
    fi
    if ! process_snapshot_alive "$all_snapshots"; then
      return 0
    fi
    sleep 0.2
  done

  signal_process_snapshot "$all_snapshots" KILL
}