#!/usr/bin/env bash

# XXX: This file isn't *quite* a script.  It is intended to be passed via stdin
# to a node to execute cleanup logic upon deletion.  Currently this is done in
# colo_node_free using the eval-cat trick.  While this gets us what we want,
# care must be taken to ensure variable expansion happens at the right time.
# Any unescaped variable references ($X) in this file will be expanded by eval
# in colo_node_free. Escaped variable references (\$X) will be expanded upon
# execution on the remote node.

RC=false
if [ -f "$SOLANA_LOCK_FILE" ]; then
  exec 9<>"$SOLANA_LOCK_FILE"
  flock -x -n 9 || exit 1
  . "$SOLANA_LOCK_FILE"
  if [ "\$SOLANA_LOCK_USER" = "\$SOLANA_USER" ]; then
    # Begin running process cleanup
    CLEANUP_PID=\$$
    CLEANUP_PIDS=()
    CLEANUP_PPIDS=()
    get_pids() {
      CLEANUP_PIDS=()
      CLEANUP_PPIDS=()
      declare line maybe_ppid maybe_pid
      while read line; do
        read maybe_ppid maybe_pid _ _ _ _ _ _ _ _ <<<"\$line"
        CLEANUP_PIDS+=( \$maybe_pid )
        CLEANUP_PPIDS+=( \$maybe_ppid )
      done < <(ps jxh | sort -rn -k2,2)
    }

    CLEANUP_PROC_CHAINS=()
    resolve_chains() {
      CLEANUP_PROC_CHAINS=()
      declare i pid ppid handled n
      for i in "\${!CLEANUP_PIDS[@]}"; do
        pid=\${CLEANUP_PIDS[\$i]}
        ppid=\${CLEANUP_PPIDS[\$i]}
        handled=false

        for j in "\${!CLEANUP_PROC_CHAINS[@]}"; do
          if grep -q "^\${ppid}\\\\b" <<<"\${CLEANUP_PROC_CHAINS[\$j]}"; then
            CLEANUP_PROC_CHAINS[\$j]="\$pid \${CLEANUP_PROC_CHAINS[\$j]}"
            handled=true
            break
          elif grep -q "\\\\b\${pid}\\$" <<<"\${CLEANUP_PROC_CHAINS[\$j]}"; then
            CLEANUP_PROC_CHAINS[\$j]+=" \$ppid"
            handled=true
            # Don't break, we may be the parent of may proc chains
          fi
        done
        if ! \$handled; then
          n=\${#CLEANUP_PROC_CHAINS[@]}
          CLEANUP_PROC_CHAINS[\$n]="\$pid \$ppid"
        fi
      done
    }

    # Kill screen sessions
    while read SID; do
      screen -S "\$SID" -X quit
    done < <(screen -wipe 2>&1 | sed -e 's/^\s\+\([^[:space:]]\+\)\s.*/\1/;t;d')

    # Kill tmux sessions
    tmux kill-server &> /dev/null

    # Kill other processes
    for SIG in INT TERM KILL; do
      get_pids
      if [[ \${#CLEANUP_PIDS[@]} -eq 0 ]]; then
        break
      else
        resolve_chains
        for p in "\${CLEANUP_PROC_CHAINS[@]}"; do
          if ! grep -q "\b\$CLEANUP_PID\b" <<<"\$p"; then
            read -a TO_KILL <<<"\$p"
            N=\${#TO_KILL[@]}
            ROOT_PPID="\${TO_KILL[\$((N-1))]}"
            if [[ 1 -ne \$ROOT_PPID ]]; then
              LAST_PID_IDX=\$((N-2))
              for I in \$(seq 0 \$LAST_PID_IDX); do
                pid="\${TO_KILL[\$I]}"
                kill -\$SIG \$pid &>/dev/null
              done
            fi
          fi
        done
        get_pids
        if [[ \${#CLEANUP_PIDS[@]} -gt 0 ]]; then
          sleep 5
        fi
      fi
    done
    # End running process cleanup

    # Begin filesystem cleanup
    git clean -qxdff
    rm -f /solana-scratch/* /solana-scratch/.[^.]*
    cat > "\${HOME}/.ssh/authorized_keys" <<EOAK
$("$__colo_here"/add-datacenter-solana-user-authorized_keys.sh 2> /dev/null)
EOAK
    # End filesystem cleanup
    RC=true
  fi
  9>&-
fi
\$RC