#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
#
# soak.sh - CUDA client soak runner for the vGPU guest.
#
# Usage:
#   tools/cuda-ab/soak.sh <seconds> [guest_dir]   # soak for <seconds>
#   tools/cuda-ab/soak.sh --seq [guest_dir]       # one pass over all clients
#
# The client set (victory_mail, smiley, vecadd_loop, streams2,
# memcpy_roundtrip, multikernel plus the victory series: victory_saxpy,
# victory_reduce, victory_matmul, victory_strided, victory_memset2d,
# victory_atomics, victory_streams_events, victory_math) runs in the guest
# one at a time via scripts/ssh_guest.sh. Clients whose binary is missing
# from <guest_dir> (default /tmp) are skipped. A client that does not print
# STATUS OK within SOAK_TIMEOUT seconds (default 240) stops the soak with a
# SOAK FAILED line.
#
# Every run is recorded as a timestamped line per run in
# /tmp/opencode/soak-<stamp>/results.log; the full client output of each run
# is kept next to it as out-r<round>-<client>.txt.
#
# Environment overrides:
#   SOAK_TIMEOUT   per-client seconds in the guest (default 240)
#   SSH_CMD        guest ssh wrapper (default <repo>/scripts/ssh_guest.sh)
#
# Examples:
#   tools/cuda-ab/soak.sh 3600            # 1 h continuous battery
#   tools/cuda-ab/soak.sh --seq           # single pass, quick check
set -u

HERE="$(cd "$(dirname "$0")" && pwd)"
REPO="$(cd "$HERE/../.." && pwd)"
SSH="${SSH_CMD:-$REPO/scripts/ssh_guest.sh}"
TMO="${SOAK_TIMEOUT:-240}"
CLIENTS="${SOAK_CLIENTS:-victory_mail vecadd_loop streams2 memcpy_roundtrip multikernel victory_saxpy victory_reduce victory_matmul victory_strided victory_memset2d victory_atomics victory_streams_events victory_math smiley}"

MODE=soak
if [ "${1:-}" = "--seq" ]; then
	MODE=seq
	SECS=0
	GUEST_DIR="${2:-/tmp}"
else
	SECS="${1:-}"
	GUEST_DIR="${2:-/tmp}"
fi

usage() {
	echo "usage: $0 <seconds> [guest_dir]" >&2
	echo "       $0 --seq [guest_dir]" >&2
	exit 1
}

[ -x "$SSH" ] || { echo "ssh wrapper not found: $SSH" >&2; exit 1; }
if [ "$MODE" = soak ]; then
	case "$SECS" in
	''|*[!0-9]*) usage ;;
	esac
	[ "$SECS" -gt 0 ] || usage
fi

STAMP="$(date +%Y%m%d-%H%M%S)"
LOGDIR="/tmp/opencode/soak-$STAMP"
mkdir -p "$LOGDIR"
RESULTS="$LOGDIR/results.log"
START=$(date +%s)
DEADLINE=$((START + SECS))

log() {
	printf '%s %s\n' "$(date '+%Y-%m-%dT%H:%M:%S%z')" "$*" | tee -a "$RESULTS"
}

# gssh_raw - one guest command, no retry (safe only for idempotent commands).
gssh_raw() {
	timeout 60 "$SSH" "$1" 2>/dev/null
}

# gssh - idempotent guest command, retried for flaky ssh.
gssh() {
	local i out

	for i in 1 2 3; do
		out="$(gssh_raw "$1")" && { printf '%s' "$out"; return 0; }
		sleep 2
	done
	return 1
}

# start_client <name> <round> - launch under remote `timeout`, output to a
# per-round file ending in EXIT=<rc>. The first attempt launches
# unconditionally; retries only re-launch if the output file is absent, so a
# dropped ssh stream cannot cause a double run.
start_client() {
	local name="$1" round="$2" bin="$GUEST_DIR/$1"
	local out="$GUEST_DIR/.soak-$1-r$round.out" i r

	for i in 1 2 3; do
		if [ "$i" = 1 ]; then
			r="$(gssh_raw "rm -f $out; nohup sh -c 'timeout $TMO $bin > $out 2>&1; echo EXIT=\$? >> $out' >/dev/null 2>&1 & echo STARTED")" || r=""
		else
			r="$(gssh_raw "if [ -f $out ]; then echo STARTED; else rm -f $out; nohup sh -c 'timeout $TMO $bin > $out 2>&1; echo EXIT=\$? >> $out' >/dev/null 2>&1 & echo STARTED; fi")" || r=""
		fi
		case "$r" in
		*STARTED*) return 0 ;;
		esac
		sleep 3
	done
	return 1
}

# client_alive <bin> - pgrep without self-match ([/] in the pattern keeps the
# remote check shell - whose argv contains the pattern text - from matching).
client_alive() {
	local pat="[/]${1#/}"

	gssh "pgrep -f '$pat' >/dev/null && echo ALIVE" | grep -q ALIVE
}

# run_one <name> <round> - returns 0 when STATUS OK was seen, 1 otherwise.
run_one() {
	local name="$1" round="$2" bin="$GUEST_DIR/$1"
	local out="$GUEST_DIR/.soak-$1-r$round.out"
	local text="" r="" i t0 tnow dur rc reason

	t0=$(date +%s)
	log "RUN client=$name round=$round timeout=${TMO}s"

	if ! start_client "$name" "$round"; then
		log "FAILED client=$name round=$round reason=ssh-launch-failed"
		return 1
	fi

	while :; do
		sleep 10
		r="$(gssh "cat $out 2>/dev/null")" || true
		[ -n "$r" ] && text="$r"
		printf '%s' "$text" | grep -q 'EXIT=' && break
		tnow=$(date +%s)
		if [ $((tnow - t0)) -gt $((TMO + 120)) ]; then
			if client_alive "$bin"; then
				log "FAILED client=$name round=$round reason=still-running-after-$((tnow - t0))s"
			else
				log "FAILED client=$name round=$round reason=process-gone-without-exit-marker"
			fi
			return 1
		fi
	done

	tnow=$(date +%s)
	dur=$((tnow - t0))

	# Flaky-ssh guard: retry the final status read a few times.
	for i in 1 2 3; do
		r="$(gssh "cat $out 2>/dev/null")" || true
		if printf '%s' "$r" | grep -q 'EXIT='; then
			text="$r"
			break
		fi
		sleep 2
	done

	rc=$(printf '%s' "$text" | sed -n 's/.*EXIT=\([0-9][0-9]*\).*/\1/p' | tail -1)
	[ -n "$rc" ] || rc=?
	printf '%s\n' "$text" > "$LOGDIR/out-r${round}-${name}.txt"

	if printf '%s' "$text" | grep -q 'STATUS OK'; then
		log "OK client=$name rc=$rc dur=${dur}s round=$round"
		return 0
	fi
	reason="$(printf '%s' "$text" | grep -E 'STATUS FAIL|CUDA_ERROR' | head -1)"
	[ -n "$reason" ] || reason="no STATUS OK in output"
	log "FAILED client=$name rc=$rc dur=${dur}s round=$round reason=$reason"
	return 1
}

log "SOAK START mode=$MODE seconds=$SECS guest_dir=$GUEST_DIR timeout=${TMO}s logdir=$LOGDIR"

# Which clients are present in the guest? Retry the listing so a flaky ssh
# cannot silently shrink the battery.
found=""
for i in 1 2 3; do
	found="$(gssh "for b in $CLIENTS; do [ -x $GUEST_DIR/\$b ] && echo \$b; done")" || found=""
	[ -n "$found" ] && break
	sleep 3
done
if [ -z "$found" ]; then
	log "SOAK FAILED reason=no-client-binaries-in-$GUEST_DIR"
	exit 1
fi
log "CLIENTS $(printf '%s' "$found" | tr '\n' ' ')"

round=0
runs=0
oks=0
while :; do
	[ "$MODE" = soak ] && [ "$(date +%s)" -ge "$DEADLINE" ] && break
	round=$((round + 1))
	stopped=0
	for c in $found; do
		if [ "$MODE" = soak ] && [ "$(date +%s)" -ge "$DEADLINE" ]; then
			stopped=1
			break
		fi
		if ! run_one "$c" "$round"; then
			log "SOAK FAILED client=$c round=$round runs=$runs ok=$oks"
			exit 1
		fi
		runs=$((runs + 1))
		oks=$((oks + 1))
	done
	[ "$MODE" = seq ] && break
	[ "$stopped" = 1 ] && break
done

log "SOAK DONE mode=$MODE rounds=$round runs=$runs ok=$oks elapsed=$(( $(date +%s) - START ))s"
exit 0