#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
#
# perf.sh - near-native performance harness for the Open-IOV emulated vGPU.
#
# Runs the same perf_workload binary on the host (bare metal, the oracle) and
# in the guest (emulated vGPU, via scripts/vgpu_vm.py), captures the host-side
# vgsp evidence that matches each run, and compares guest/host ratios.
#
# Usage:
#   perf.sh check                 # prerequisites; pings the VM if reachable,
#                                 # never touches GPU/module/QEMU
#   perf.sh host                  # build + oracle runs -> results/<stamp>-host/
#   perf.sh guest                 # push + measured runs -> results/<stamp>-guest/
#   perf.sh all                   # host + guest + compare
#   perf.sh compare [HOST_DIR] [GUEST_DIR]   # default: newest of each
#   perf.sh push                  # push the built binary only
#
# Sequencing note: the bare-metal oracle needs the physical GPU on the stock
# distribution driver, so run `host` while no campaign cycle holds the device
# (before insmod of vgsp, or after release). `guest` runs only in a green VM
# under the supervised campaign rules. See perf/PLAN-AND-HARNESS.md.
#
# Environment:
#   PERF_RUNS      process runs per side (default 5; protocol wants 10)
#   PERF_N_WARM    warm launches per run (default 50; protocol 100)
#   PERF_N_THRU    throughput launches per run (default 100)
#   PERF_TIMEOUT   guest exec timeout in seconds (default 900)
#   PERF_VMEXEC    guest executor (default scripts/vgpu_vm.py)
#   PERF_STAMP     shared stamp for host+guest pairing
#   PERF_ORACLE    host results.json for guest/compare without a fresh host run
set -u

HERE="$(cd "$(dirname "$0")" && pwd)"
REPO="$(cd "$HERE/../../.." && pwd)"
VMEXEC="${PERF_VMEXEC:-$REPO/scripts/vgpu_vm.py}"
RUN_ROOT="${PERF_RUN_ROOT:-$HERE/results}"
BIN="${PERF_BIN:-$HERE/bin/perf_workload}"
SRC="$HERE/perf_workload.cu"
COLLECT="$HERE/collect.py"
STAMP="${PERF_STAMP:-$(date +%Y%m%d-%H%M%S)}"
RUNS="${PERF_RUNS:-5}"
N_WARM="${PERF_N_WARM:-50}"
N_THRU="${PERF_N_THRU:-100}"
GUEST_TIMEOUT="${PERF_TIMEOUT:-900}"
STATE_FILE="/sys/kernel/debug/vgsp/gsp-phys-state"
GR_GREP='vgsp-gr:|work_gr:|sema-gr:|vgsp-sema: sema-svc:'

build() {
	command -v nvcc >/dev/null || {
		echo "nvcc not found (install nvidia-cuda-toolkit)" >&2
		return 2
	}
	mkdir -p "$HERE/bin"
	echo "[*] build $BIN"
	nvcc -O2 -arch=sm_75 --cudart=static -o "$BIN" "$SRC"
}

run_binary() {
	"$BIN" --n-warm="$N_WARM" --n-thru="$N_THRU"
}

meta_common() {
	local out="$1" role="$2"

	{
		echo "stamp: $STAMP"
		echo "role: $role"
		echo "date: $(date -u +%Y-%m-%dT%H:%M:%SZ)"
		echo "uname: $(uname -r)"
		echo "runs: $RUNS n_warm: $N_WARM n_thru: $N_THRU"
		echo "module_ko_sha256: $(sha256sum "$REPO/kernel/vgsp.ko" 2>/dev/null | awk '{print $1}' || echo unavailable)"
	} > "$out"
}

meta_clocks_host() {
	local out="$1"

	{
		echo "--- nvidia-smi -q -d CLOCK (host) ---"
		nvidia-smi -q -d CLOCK 2>&1 || echo "nvidia-smi unavailable"
	} >> "$out"
}

capture_state() {
	local out="$1"

	{
		echo "--- gsp-phys-state $(date -u +%Y-%m-%dT%H:%M:%SZ) ---"
		cat "$STATE_FILE" 2>/dev/null ||
			sudo -n cat "$STATE_FILE" 2>/dev/null ||
			echo "unavailable (debugfs needs root)"
	} > "$out"
}

capture_dmesg() {
	local out="$1"

	{
		dmesg 2>/dev/null || sudo -n dmesg 2>/dev/null || echo "unavailable"
	} | grep -E "$GR_GREP" | tail -n 3000 > "$out"
}

capture_dmesg_marker() {
	local out="$1"

	{
		dmesg 2>/dev/null || sudo -n dmesg 2>/dev/null || echo "unavailable"
	} | tail -n 1 > "$out"
}

run_host() {
	local dir="$RUN_ROOT/$STAMP-host"

	build || return $?
	mkdir -p "$dir"
	meta_common "$dir/meta.txt" host
	meta_clocks_host "$dir/meta.txt"

	if ! nvidia-smi -L >/dev/null 2>&1; then
		echo "ERROR: host GPU not on the distribution driver (claimed or unavailable)." >&2
		echo "       Run 'perf.sh host' while no vgsp cycle holds the device." >&2
		return 2
	fi

	local i rc=0

	for i in $(seq 1 "$RUNS"); do
		echo "[*] host run $i/$RUNS"
		run_binary > "$dir/run-$i.txt" 2>&1 || rc=1
	done
	cat "$dir"/run-[0-9]*.txt > "$dir/run.txt"
	python3 "$COLLECT" --role host --guest "$dir"/run-[0-9]*.txt \
		--meta "$dir/meta.txt" --out "$dir/results.json" || return 1
	if [ "$rc" -ne 0 ]; then
		echo "WARNING: at least one host run failed self-verification" >&2
	fi
	echo "[*] host results: $dir/results.json"
}

push_bin() {
	local b64 chunk off total size

	[ -x "$BIN" ] || build || return 2
	b64=$(base64 -w0 "$BIN")
	total=${#b64}
	off=0
	size=70000
	echo "[*] pushing $(stat -c%s "$BIN") bytes to guest"
	"$VMEXEC" exec "rm -f /tmp/perf_workload.b64" >/dev/null 2>&1
	while [ "$off" -lt "$total" ]; do
		chunk=${b64:$off:$size}
		"$VMEXEC" exec "echo -n '$chunk' >> /tmp/perf_workload.b64" >/dev/null 2>&1
		off=$((off + size))
	done
	"$VMEXEC" exec "base64 -d /tmp/perf_workload.b64 > /tmp/perf_workload && chmod +x /tmp/perf_workload && rm -f /tmp/perf_workload.b64 && echo PUSH_OK"
}

run_guest() {
	local dir="$RUN_ROOT/$STAMP-guest"

	build || return $?
	mkdir -p "$dir"
	meta_common "$dir/meta.txt" guest

	local ping

	ping="$("$VMEXEC" exec "echo VMPING" 2>&1 || true)"
	if [[ "$ping" != *VMPING* ]]; then
		echo "ERROR: guest not reachable via $VMEXEC (is the VM running?)" >&2
		return 2
	fi

	push_bin || return 1
	capture_state "$dir/host-state-before.txt"
	capture_dmesg_marker "$dir/dmesg-marker.txt"

	{
		echo "--- guest driver/GPU (via guest agent) ---"
		"$VMEXEC" exec "nvidia-smi --query-gpu=name,driver_version --format=csv" 2>&1
		"$VMEXEC" exec "nvidia-smi -q -d CLOCK" 2>&1
	} > "$dir/meta-guest.txt"

	local i rc=0

	for i in $(seq 1 "$RUNS"); do
		echo "[*] guest run $i/$RUNS (timeout ${GUEST_TIMEOUT}s)"
		"$VMEXEC" exec --timeout "$GUEST_TIMEOUT" \
			"/tmp/perf_workload --n-warm=$N_WARM --n-thru=$N_THRU" \
			> "$dir/run-$i.txt" 2>&1
		[ $? -ne 0 ] && rc=1
	done
	cat "$dir"/run-[0-9]*.txt > "$dir/run.txt"

	capture_state "$dir/host-state-after.txt"
	capture_dmesg "$dir/host-gr-dmesg.txt"
	python3 "$COLLECT" --role guest --guest "$dir"/run-[0-9]*.txt \
		--host-gr "$dir/host-gr-dmesg.txt" \
		--host-gr-marker "$dir/dmesg-marker.txt" \
		--state-before "$dir/host-state-before.txt" \
		--state-after "$dir/host-state-after.txt" \
		--meta "$dir/meta.txt" \
		--out "$dir/results.json" || return 1
	if [ "$rc" -ne 0 ]; then
		echo "WARNING: at least one guest run failed or timed out" >&2
	fi
	echo "[*] guest results: $dir/results.json"
}

latest_dir() {
	ls -dt "$RUN_ROOT"/*-"$1" 2>/dev/null | head -n1
}

run_compare() {
	local hdir="${1:-}" gdir="${2:-}" hostjson guestjson

	if [ -z "$hdir" ]; then
		hdir=$(latest_dir host)
		[ -n "$PERF_ORACLE" ] && hdir=$(dirname "$PERF_ORACLE")
	fi
	[ -z "$gdir" ] && gdir=$(latest_dir guest)
	if [ -n "$PERF_ORACLE" ]; then
		hostjson="$PERF_ORACLE"
	else
		hostjson="$hdir/results.json"
	fi
	guestjson="$gdir/results.json"

	[ -f "$hostjson" ] || {
		echo "missing host oracle: $hostjson" >&2
		return 2
	}
	[ -f "$guestjson" ] || {
		echo "missing guest results: $guestjson" >&2
		return 2
	}
	local out="$gdir/compare.json"
	local -a extra=()

	if [ -n "${PERF_THRESHOLDS:-}" ]; then
		extra=(--thresholds "$PERF_THRESHOLDS")
	fi
	python3 "$COLLECT" --compare "$hostjson" "$guestjson" --out "$out" \
		"${extra[@]}"
	local rc=$?

	echo "[*] compare: $out"
	return $rc
}

check() {
	echo "== perf harness check =="
	command -v nvcc >/dev/null && echo "nvcc: $(command -v nvcc)" || echo "nvcc: MISSING"
	[ -f "$VMEXEC" ] && echo "vmexec: $VMEXEC" || echo "vmexec: MISSING ($VMEXEC)"
	echo "runs=$RUNS n_warm=$N_WARM n_thru=$N_THRU timeout=$GUEST_TIMEOUT"
	[ -r "$STATE_FILE" ] && echo "gsp-phys-state: readable" ||
		echo "gsp-phys-state: not readable as user (sudo -n is used when available)"
	if nvidia-smi -L >/dev/null 2>&1; then
		echo "host GPU: available (oracle can run now)"
	else
		echo "host GPU: NOT available (claimed/unavailable; run oracle first/after release)"
	fi
	ping=""
	if [ -f "$VMEXEC" ]; then
		ping="$("$VMEXEC" exec "echo VMPING" 2>&1 || true)"
	fi
	[[ "$ping" == *VMPING* ]] && echo "guest: reachable" || echo "guest: not reachable"
}

case "${1:-all}" in
check)		check ;;
build)		build ;;
host)		run_host ;;
guest)		run_guest ;;
compare)	shift; run_compare "$@" ;;
all)
	run_host || echo "WARNING: host oracle unavailable; continuing with guest" >&2
	run_guest || exit 1
	run_compare || exit 1
	;;
push)		push_bin ;;
*)		echo "usage: $0 [check|build|host|guest|compare|all|push]" >&2
		exit 2 ;;
esac