#!/usr/bin/env bash
# Bounded poller for agent-driven async waits. Polls a command until its output
# satisfies a success condition, a wall-clock timeout elapses, or a non-success
# value repeats enough times to be a stuck finding. It never loops unbounded and
# never blocks past the timeout, even when a single poll command hangs: each poll
# is capped at the remaining wall-clock budget and killed if it overruns.
#
# Usage:
#   wait-for.sh --timeout <sec> --interval <sec> [--expect <regex>] [--stuck-after <n>] -- <poll-cmd...>
#
# Success: with --expect, the poll's trimmed stdout matches the extended regex;
# without --expect, the poll command exits 0. --stuck-after requires --expect,
# because a stuck finding is about a value stuck at an unexpected reading.
#
# Exit codes:
#   0   met     -> prints "met: <laststate>"
#   2   timeout -> prints "timeout after <elapsed>s, last state: <laststate>"
#   3   stuck   -> prints "stuck: <laststate> repeated <n>x"
#   64  usage   -> usage line on stderr
set -u

usage() {
  echo "usage: wait-for.sh --timeout <sec> --interval <sec> [--expect <regex>] [--stuck-after <n>] -- <poll-cmd...>" >&2
  exit 64
}

# Run "$@" with a hard cap of $1 seconds. Sets POLL_OUT (stdout) and POLL_RC.
# A poll that overruns the cap is killed and reported with POLL_RC=124, so a
# hung poll cannot block past the caller's wall-clock budget.
POLL_OUT=""; POLL_RC=0
run_poll() {
  local cap="$1"; shift
  local tmp pid end
  tmp=$(mktemp)
  "$@" >"$tmp" 2>/dev/null &
  pid=$!
  end=$(( $(date +%s) + cap ))
  # Re-check on a sub-second granularity so a fast poll returns promptly while a
  # hung one is still killed at the cap.
  while kill -0 "$pid" 2>/dev/null; do
    if [ "$(date +%s)" -ge "$end" ]; then
      kill -TERM "$pid" 2>/dev/null
      wait "$pid" 2>/dev/null || true
      POLL_RC=124; POLL_OUT="$(cat "$tmp")"; rm -f "$tmp"
      return
    fi
    sleep 0.2
  done
  wait "$pid"; POLL_RC=$?
  POLL_OUT="$(cat "$tmp")"; rm -f "$tmp"
}

TIMEOUT=""; INTERVAL=""; EXPECT=""; EXPECT_SET=0; STUCK_AFTER=""
while [ $# -gt 0 ]; do
  case "$1" in
    --timeout)     TIMEOUT="${2:-}"; shift 2 || usage;;
    --interval)    INTERVAL="${2:-}"; shift 2 || usage;;
    --expect)      EXPECT="${2:-}"; EXPECT_SET=1; shift 2 || usage;;
    --stuck-after) STUCK_AFTER="${2:-}"; shift 2 || usage;;
    --)            shift; break;;
    *)             usage;;
  esac
done

# --timeout, --interval, and a poll command are required; both must be >= 1.
[ -n "$TIMEOUT" ] && [ -n "$INTERVAL" ] && [ $# -gt 0 ] || usage
case "$TIMEOUT" in ''|*[!0-9]*) usage;; esac;  [ "$TIMEOUT" -ge 1 ] || usage
case "$INTERVAL" in ''|*[!0-9]*) usage;; esac; [ "$INTERVAL" -ge 1 ] || usage
# --stuck-after, when present, must be a positive integer and needs --expect.
if [ -n "$STUCK_AFTER" ]; then
  case "$STUCK_AFTER" in *[!0-9]*) usage;; esac
  [ "$STUCK_AFTER" -ge 1 ] || usage
  [ "$EXPECT_SET" -eq 1 ] || usage
fi

start=$(date +%s)
deadline=$((start + TIMEOUT))
laststate=""; prev=""; repeat=0

while :; do
  now=$(date +%s); remaining=$((deadline - now))
  if [ "$remaining" -le 0 ]; then echo "timeout after $((now - start))s, last state: $laststate"; exit 2; fi

  run_poll "$remaining" "$@"
  laststate="$(printf '%s' "$POLL_OUT" | tr '\n' ' ' | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')"

  if [ "$EXPECT_SET" -eq 1 ]; then
    if printf '%s' "$laststate" | grep -Eq -- "$EXPECT"; then echo "met: $laststate"; exit 0; fi
    if [ -n "$STUCK_AFTER" ]; then
      if [ "$laststate" = "$prev" ]; then repeat=$((repeat + 1)); else repeat=1; prev="$laststate"; fi
      if [ "$repeat" -ge "$STUCK_AFTER" ]; then echo "stuck: $laststate repeated ${repeat}x"; exit 3; fi
    fi
  else
    if [ "$POLL_RC" -eq 0 ]; then echo "met: $laststate"; exit 0; fi
  fi

  now=$(date +%s); remaining=$((deadline - now))
  if [ "$remaining" -le 0 ]; then echo "timeout after $((now - start))s, last state: $laststate"; exit 2; fi
  if [ "$INTERVAL" -lt "$remaining" ]; then sleep "$INTERVAL"; else sleep "$remaining"; fi
done