#!/usr/bin/env bash
# Shared parallel-nix-invocation helper for scripts/codex-maintenance.sh.
# Source alongside nix-eval.sh:
#   source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/lib/nix-parallel.sh"
#
# The per-host/per-package `nix eval`/`nix build --dry-run` calls in
# codex-maintenance.sh are independent of each other, so running them one at
# a time leaves most cores idle for most of the sweep -- run_nix_parallel
# fans a batch of them out across up to NIX_PARALLEL_JOBS processes instead.

# NIX_PARALLEL_JOBS: how many `nix` invocations run_nix_parallel runs at
# once. Defaults to core count capped by available memory (~1GB/job,
# floor 1) rather than plain `nproc` -- each concurrent `nix eval` here
# evaluates a whole NixOS system closure from scratch, and on a small/
# memory-constrained CI runner, `nproc` concurrent evals can OOM-kill each
# other (confirmed empirically: on a 4GB/6-core box, 5-6 concurrent evals
# started getting killed while 3-4 ran clean and were still ~2x faster than
# serial). Override via env if a given machine/CI runner has room to spare
# or needs a tighter cap.
default_nix_parallel_jobs() {
  local cores mem_avail_kb mem_cap
  cores="$(nproc 2>/dev/null || echo 4)"
  mem_avail_kb="$(awk '/^MemAvailable:/ {print $2}' /proc/meminfo 2>/dev/null)"
  if [[ -z "$mem_avail_kb" ]]; then
    echo "$cores"
    return
  fi
  mem_cap=$((mem_avail_kb / 1024 / 1024))
  ((mem_cap < 1)) && mem_cap=1
  ((mem_cap < cores)) && echo "$mem_cap" || echo "$cores"
}
NIX_PARALLEL_JOBS="${NIX_PARALLEL_JOBS:-$(default_nix_parallel_jobs)}"

# Separator between a job's label and its flake attr in the arrays
# run_nix_parallel takes -- a control character so it can't collide with
# anything a label or attr path would plausibly contain.
NIX_PARALLEL_SEP=$'\x1f'

# run_nix_parallel <jobs_array_name> <nix subcommand + flags...>
#
# jobs_array_name: name of an already-populated bash array whose entries are
# "<label>${NIX_PARALLEL_SEP}<attr>" pairs, e.g.
#   jobs=("proxmox-docker${NIX_PARALLEL_SEP}.#nixosConfigurations.proxmox-docker...drvPath")
# Remaining args are passed to `nix` before the attr, e.g.:
#   run_nix_parallel jobs eval --raw "${NIX_EVAL_FLAGS[@]}"
#   run_nix_parallel jobs build --dry-run --no-link "${NIX_EVAL_FLAGS[@]}"
#
# Prints "==> <label>" followed by that job's stdout+stderr for every job,
# in submission order (not completion order) so a run stays readable and
# diffable across invocations even though the work itself doesn't finish in
# that order. Returns non-zero if any job failed, only after every job has
# finished and been printed -- same "surface everything, then fail" contract
# a `set -e` caller gets, just parallelized instead of stopping at the first
# failure.
run_nix_parallel() {
  local -n jobs_ref="$1"
  shift
  local -a nix_args=("$@")

  local n=${#jobs_ref[@]}
  [[ $n -eq 0 ]] && return 0

  local tmp_dir
  tmp_dir="$(mktemp -d)"

  local i=0 running=0
  for job in "${jobs_ref[@]}"; do
    local attr="${job#*"${NIX_PARALLEL_SEP}"}"
    printf '%s\n' "${job%%"${NIX_PARALLEL_SEP}"*}" >"${tmp_dir}/${i}.label"
    (
      if nix "${nix_args[@]}" "$attr" >"${tmp_dir}/${i}.out" 2>&1; then
        echo 0 >"${tmp_dir}/${i}.status"
      else
        echo 1 >"${tmp_dir}/${i}.status"
      fi
    ) &
    i=$((i + 1))
    running=$((running + 1))
    if ((running >= NIX_PARALLEL_JOBS)); then
      wait -n
      running=$((running - 1))
    fi
  done
  wait

  local failed=0 j
  for ((j = 0; j < n; j++)); do
    echo "==> $(cat "${tmp_dir}/${j}.label")"
    cat "${tmp_dir}/${j}.out"
    [[ "$(cat "${tmp_dir}/${j}.status")" -ne 0 ]] && failed=1
  done

  rm -rf "$tmp_dir"
  return $failed
}