#!/usr/bin/env bash

# Own and reap one parallel native-emission cohort. Babashka launches this one
# manager instead of waiting on every worker through native-image's Java
# Process implementation. Each worker still has its own subtree supervisor and
# deadline; this manager records every exit only after Bash has reaped it.

set -euo pipefail

pids=()

shutdown_workers() {
  local pid
  for pid in "${pids[@]}"; do
    [[ -n "$pid" ]] || continue
    kill -TERM "$pid" 2>/dev/null || true
  done
  for pid in "${pids[@]}"; do
    [[ -n "$pid" ]] || continue
    wait "$pid" 2>/dev/null || true
  done
}

die() {
  echo "beagle emission manager: $*" >&2
  shutdown_workers
  exit 2
}

[[ $# -eq 12 ]] || die "expected 12 arguments"

supervisor="$1"
worker_timeout="$2"
kill_grace="$3"
bb="$4"
compiled="$5"
backend="$6"
wire_path="$7"
root="$8"
function_count="$9"
worker_count="${10}"
abi_id="${11}"
worker_code="${12}"

[[ "$function_count" =~ ^[1-9][0-9]*$ ]] ||
  die "function count must be positive"
[[ "$worker_count" =~ ^[1-9][0-9]*$ ]] ||
  die "worker count must be positive"
((worker_count <= function_count)) ||
  die "worker count exceeds function count"

# Invoked by the signal trap below.
# shellcheck disable=SC2329
terminate_workers() {
  trap - HUP INT TERM
  shutdown_workers
  exit 2
}

trap terminate_workers HUP INT TERM

# Opt-in phase trace: absolute nanosecond stamps for the manager span and each
# worker span, appended to the file named by BEAGLE_CORE_EMIT_TRACE. Nothing
# outside the trace file changes, so emitted bytes stay identical.
trace_path="${BEAGLE_CORE_EMIT_TRACE:-}"
trace() {
  [[ -n "$trace_path" ]] || return 0
  printf '%s\t%s\t%s\n' "$(date +%s%N)" "$1" "${2:-}" >>"$trace_path"
}

if [[ -n "${BEAGLE_CORE_EMIT_WIRE_COPY:-}" ]]; then
  cp -f -- "$wire_path" "$BEAGLE_CORE_EMIT_WIRE_COPY" 2>/dev/null || true
fi

trace manager-start "backend=$backend functions=$function_count workers=$worker_count"

for ((worker_index = 0; worker_index < worker_count; worker_index++)); do
  worker_dir="$root/$backend-$worker_index"
  mkdir -p "$worker_dir" || die "could not create worker $worker_index directory"
  indexes=""
  for ((function_index = worker_index;
       function_index < function_count;
       function_index += worker_count)); do
    indexes+="${indexes:+,}$function_index"
  done
  trace worker-start "$worker_index"
  BEAGLE_BOUNDED_COMPLETION_RECEIPT="$worker_dir/supervisor.receipt" \
    "$supervisor" "$worker_timeout" "$kill_grace" -- \
    "$bb" -cp "$compiled" -e "$worker_code" \
    "$backend" "$wire_path" "$worker_dir" "$indexes" "$abi_id" \
    "$worker_index" \
    >"$worker_dir/stdout.log" 2>"$worker_dir/stderr.log" &
  pids+=("$!")
done

for ((worker_index = 0; worker_index < worker_count; worker_index++)); do
  if wait "${pids[$worker_index]}"; then
    status=0
  else
    status=$?
  fi
  pids[worker_index]=""
  printf '%s\n' "$status" >"$root/$backend-$worker_index/manager.status" ||
    die "could not record worker $worker_index status"
  trace worker-reaped "$worker_index status=$status"
done

trace manager-end "workers=$worker_count"

receipt="$root/manager.receipt"
temporary="$root/.manager.receipt.$$"
printf 'emission-manager-v0 reaped workers=%s\n' "$worker_count" >"$temporary" ||
  die "could not write completion receipt"
mv -f -- "$temporary" "$receipt" || die "could not publish completion receipt"
trap - HUP INT TERM
exit 0
