#!/usr/bin/env bash
# bc250-cu-health-test.sh - reboot-resuming per-WGP health test for BC-250.

set -euo pipefail

STATEDIR="${BC250_CU_HEALTH_STATE:-/var/lib/bc250-cu-health-test}"
CONF="${BC250_CU_HEALTH_CONF:-/etc/modprobe.d/bc250-cu-health-isolate.conf}"
SERVICE="${BC250_CU_HEALTH_SERVICE:-bc250-cu-health-resume.service}"
VERIFY="${BC250_CU_VERIFY:-$(dirname "$(readlink -f "$0")")/bc250-compute-verify.sh}"
ELEMENTS="${BC250_CU_HEALTH_ELEMENTS:-16777216}"
PASSES="${BC250_CU_HEALTH_PASSES:-2}"
ITERS="${BC250_CU_HEALTH_ITERS:-64}"
REBOOT_DELAY="${BC250_CU_HEALTH_REBOOT_DELAY:-5}"
FINAL_REBOOT="${BC250_CU_HEALTH_FINAL_REBOOT:-1}"

usage() {
	cat <<EOF
Usage: $0 start|resume|quick|status|reset

Commands:
  start   Install a systemd resume service and test all 20 WGPs across reboots.
  resume  Internal command used by systemd after each reboot.
  quick   Run the heavy verifier once in the current CU configuration.
  status  Print saved progress/results.
  reset   Remove state, resume service, and isolation config.

Environment:
  BC250_CU_HEALTH_ELEMENTS=$ELEMENTS
  BC250_CU_HEALTH_PASSES=$PASSES
  BC250_CU_HEALTH_ITERS=$ITERS
  BC250_CU_HEALTH_FINAL_REBOOT=$FINAL_REBOOT

Results:
  $STATEDIR/results.tsv
EOF
}

die() {
	echo "ERROR: $*" >&2
	exit 1
}

need_root() {
	[ "$(id -u)" = "0" ] || die "must run as root"
}

target_to_tuple() {
	local idx="$1"
	local se sh wgp

	se=$((idx / 10))
	sh=$(((idx / 5) % 2))
	wgp=$((idx % 5))
	echo "$se $sh $wgp"
}

disable_cu_for_target() {
	local target="$1"
	local out=()
	local idx se sh wgp tse tsh twgp cu

	read -r tse tsh twgp < <(target_to_tuple "$target")
	for idx in $(seq 0 19); do
		read -r se sh wgp < <(target_to_tuple "$idx")
		if [ "$se" -eq "$tse" ] && [ "$sh" -eq "$tsh" ] && [ "$wgp" -eq "$twgp" ]; then
			continue
		fi
		out+=("$se.$sh.$((wgp * 2))")
		out+=("$se.$sh.$((wgp * 2 + 1))")
	done
	IFS=,
	echo "${out[*]}"
}

write_config_for_target() {
	local target="$1"
	local disable_csv
	local se sh wgp

	read -r se sh wgp < <(target_to_tuple "$target")
	disable_csv="$(disable_cu_for_target "$target")"
	cat >"$CONF" <<EOF
# BC-250 per-WGP health-test isolation config.
# Target: index=$target SE$se SH$sh WGP$wgp CUs $((wgp * 2))-$((wgp * 2 + 1)).
# Generated by $0 on $(date -Iseconds).
options amdgpu bc250_cc_write_mode=3 disable_cu=$disable_csv
EOF
}

write_full_config() {
	cat >"$CONF" <<EOF
# BC-250 full 40-CU config restored after health test.
options amdgpu bc250_cc_write_mode=3
EOF
}

refresh_initramfs() {
	if command -v update-initramfs >/dev/null 2>&1; then
		update-initramfs -u -k "$(uname -r)" || true
	elif command -v dracut >/dev/null 2>&1; then
		dracut -f || true
	fi
}

install_service() {
	local self

	self="$(readlink -f "$0")"
	cat >/etc/systemd/system/"$SERVICE" <<EOF
[Unit]
Description=Resume BC-250 per-WGP CU health test
After=multi-user.target

[Service]
Type=oneshot
ExecStart=/bin/bash $self resume
StandardOutput=journal+console
StandardError=journal+console

[Install]
WantedBy=multi-user.target
EOF
	systemctl daemon-reload
	systemctl enable "$SERVICE"
}

remove_service() {
	systemctl disable "$SERVICE" >/dev/null 2>&1 || true
	rm -f /etc/systemd/system/"$SERVICE"
	systemctl daemon-reload || true
}

current_cu_count() {
	dmesg | grep -o 'active_cu_number [0-9]*' | tail -1 | awk '{print $2}'
}

run_verify_for_target() {
	local target="$1"
	local se sh wgp log status rc cu_count started finished

	read -r se sh wgp < <(target_to_tuple "$target")
	log="$STATEDIR/logs/wgp-${target}-se${se}-sh${sh}-wgp${wgp}.log"
	started="$(date -Iseconds)"
	cu_count="$(current_cu_count || true)"
	status="PASS"
	rc=0

	echo "Testing target index=$target SE$se SH$sh WGP$wgp..."
	set +e
	"$VERIFY" --elements "$ELEMENTS" --passes "$PASSES" --iters "$ITERS" 2>&1 | tee "$log"
	rc=${PIPESTATUS[0]}
	set -e
	[ "$rc" -eq 0 ] || status="FAIL"
	finished="$(date -Iseconds)"

	printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
		"$target" "$se" "$sh" "$wgp" "$status" "$rc" \
		"${cu_count:-unknown}" "$started" "$finished" >>"$STATEDIR/results.tsv"
}

next_reboot() {
	local target="$1"
	local se sh wgp

	read -r se sh wgp < <(target_to_tuple "$target")
	write_config_for_target "$target"
	echo "$target" >"$STATEDIR/current_target"
	refresh_initramfs
	echo "Configured next isolated target: index=$target SE$se SH$sh WGP$wgp."
	echo "Rebooting in $REBOOT_DELAY seconds..."
	sleep "$REBOOT_DELAY"
	reboot
}

finish_test() {
	write_full_config
	refresh_initramfs
	remove_service
	echo "done" >"$STATEDIR/phase"
	echo "Per-WGP health test complete."
	echo "Results: $STATEDIR/results.tsv"
	awk -F'\t' 'BEGIN{pass=0; fail=0} $5=="PASS"{pass++} $5=="FAIL"{fail++} END{printf "PASS WGPs: %d  FAIL WGPs: %d  usable CUs: %d/40\n", pass, fail, 40 - fail * 2}' "$STATEDIR/results.tsv"
	if [ "$FINAL_REBOOT" = "1" ]; then
		echo "Rebooting in $REBOOT_DELAY seconds to reload amdgpu with the restored full-40CU config..."
		sleep "$REBOOT_DELAY"
		reboot
	else
		echo "Full 40-CU config is written for the next boot; current boot remains on the last isolated WGP."
	fi
}

cmd="${1:-}"
case "$cmd" in
	start)
		need_root
		[ -x "$VERIFY" ] || die "verifier not executable: $VERIFY"
		mkdir -p "$STATEDIR/logs"
		cat >"$STATEDIR/results.tsv" <<'EOF'
#idx	se	sh	wgp	status	rc	active_cu	started	finished
EOF
		echo "running" >"$STATEDIR/phase"
		install_service
		next_reboot 0
		;;
	resume)
		need_root
		[ -x "$VERIFY" ] || die "verifier not executable: $VERIFY"
		mkdir -p "$STATEDIR/logs"
		phase="$(cat "$STATEDIR/phase" 2>/dev/null || echo missing)"
		[ "$phase" = "running" ] || die "health test is not running (phase=$phase)"
		target="$(cat "$STATEDIR/current_target" 2>/dev/null || echo 0)"
		run_verify_for_target "$target"
		next=$((target + 1))
		if [ "$next" -lt 20 ]; then
			next_reboot "$next"
		else
			finish_test
		fi
		;;
	quick)
		[ -x "$VERIFY" ] || die "verifier not executable: $VERIFY"
		"$VERIFY" --elements "$ELEMENTS" --passes "$PASSES" --iters "$ITERS"
		;;
	status)
		if [ -f "$STATEDIR/results.tsv" ]; then
			cat "$STATEDIR/results.tsv"
		else
			echo "No health-test results in $STATEDIR"
		fi
		;;
	reset)
		need_root
		remove_service
		rm -rf "$STATEDIR"
		rm -f "$CONF"
		refresh_initramfs
		echo "Removed health-test state, service, and isolation config."
		;;
	-h|--help|"")
		usage
		;;
	*)
		die "unknown command: $cmd"
		;;
esac