#!/bin/bash

## Copyright (C) 2026 - 2026 ENCRYPTED SUPPORT LLC <adrelanos@whonix.org>
## See the file COPYING for copying conditions.

## AI-Assisted

## Reusable helpers for waiting on background jobs while still
## capturing each job's exit status. Source from a bash script.
##
## Background:
##   `wait pid1 pid2 ...` returns the exit status of the *last* PID
##   only - earlier failures are silently dropped. The helpers below
##   work around that using bash 5.1's `wait -n -p VAR`, which
##   returns when any one child finishes and writes its PID into
##   VAR. `wait -n` for a child that has already exited returns its
##   saved exit status immediately, so total wall-clock cost is
##   bounded by the slowest job regardless of completion order.
##
##   `wait -n -p done_pid PID...` is called with the explicit list
##   of tracked PIDs (not bare `wait -n`) so an unrelated background
##   job started in the caller's shell cannot be reaped here and
##   counted against this batch.
##
## Requirements: bash >= 5.1 (for `wait -n -p`); `local -n` (bash
## >= 4.3) for the array-by-reference parameter.

# shellcheck source=/usr/libexec/helper-scripts/strings.bsh
source "${HELPER_SCRIPTS_PATH:-}"/usr/libexec/helper-scripts/strings.bsh
# shellcheck source=/usr/libexec/helper-scripts/log_run_die.sh
source "${HELPER_SCRIPTS_PATH:-}"/usr/libexec/helper-scripts/log_run_die.sh

## Output array populated by parallel_wait on failure.
parallel_failures=()

## parallel_wait NAME
##
## Wait for every PID in the associative array NAME and return 0
## if all jobs succeeded. Each job's outcome is logged as it is
## reaped: 'log info' on success, 'log error' on failure. The
## failure log line includes the job's PID and exit code. On any
## failure, the global array `parallel_failures` is populated with
## the labels of the failed jobs and the function returns 1.
##
## NAME is the *name* of an associative array { pid => label },
## taken by reference via bash nameref.
##
## Usage:
##   declare -A jobs=()
##   sleep 1 & jobs[$!]="job-A"
##   sleep 2 & jobs[$!]="job-B"
##   if ! parallel_wait jobs ; then
##      log error "failed jobs: ${parallel_failures[*]}"
##      exit 1
##   fi
parallel_wait() {
  local done_pid label pid_iter wait_rc
  local -a remaining_pids new_remaining

  check_variable_name "${1}" || return 1
  local -n _parallel_jobs="${1}"

  parallel_failures=()
  remaining_pids=( "${!_parallel_jobs[@]}" )

  if [ "${#remaining_pids[@]}" -le 0 ]; then
    log debug "parallel_wait: no jobs to wait for"
    return 0
  fi

  while [ "${#remaining_pids[@]}" -gt 0 ] ; do
    ## `wait -n -p done_pid` may unset done_pid when no child is
    ## available; reset before each iteration.
    done_pid=""
    wait_rc=0
    ## Pass the tracked PIDs explicitly so bash only considers our
    ## jobs. Without this, `wait -n` could reap an unrelated
    ## background job started in the caller's shell, which would
    ## still decrement `remaining` and cause one of our tracked
    ## PIDs to be silently skipped.
    ##
    ## Capture the exit code via '|| wait_rc=$?' so we can branch
    ## on success vs failure (and surface the exit code in the log
    ## line) without errexit terminating the loop on a failed job.
    wait -n -p done_pid "${remaining_pids[@]}" || wait_rc=$?

    if [ -z "${done_pid:-}" ]; then
      ## Defensive: all tracked PIDs already reaped or otherwise
      ## not waitable. Bash returned without populating done_pid;
      ## the loop's progress condition depends on done_pid removal
      ## from remaining_pids, so without it we'd spin. Mark every
      ## remaining tracked PID as a failure so the caller's
      ## `[ "${#parallel_failures[@]}" -eq 0 ]` final check returns
      ## non-zero - silently dropping unwaitable PIDs would let
      ## "missing observation" look indistinguishable from "all
      ## succeeded".
      log error "${FUNCNAME[0]}: ${#remaining_pids[@]} tracked job(s) could not be waited on (wait_rc=${wait_rc})"
      for pid_iter in "${remaining_pids[@]}"; do
        if [ -n "${_parallel_jobs[${pid_iter}]:-}" ]; then
          parallel_failures+=( "${_parallel_jobs[${pid_iter}]} (unwaitable)" )
        else
          parallel_failures+=( "pid ${pid_iter} (unwaitable)" )
        fi
      done
      break
    fi

    label="${_parallel_jobs[${done_pid}]:-pid ${done_pid}}"

    if [ "${wait_rc}" -eq 0 ]; then
      log info "${FUNCNAME[0]}: job succeeded: '${label}' (pid=${done_pid})"
    else
      log error "${FUNCNAME[0]}: job failed: '${label}' (pid=${done_pid}, exit=${wait_rc})"
      parallel_failures+=( "${label}" )
    fi

    new_remaining=()
    for pid_iter in "${remaining_pids[@]}"; do
      [ "${pid_iter}" != "${done_pid}" ] && new_remaining+=( "${pid_iter}" )
    done
    remaining_pids=( "${new_remaining[@]}" )
  done
  [ "${#parallel_failures[@]}" -eq 0 ] || return 1
}
