#!/bin/bash
# run_cycle.sh - one instrumented CUDA bring-up cycle with full capture.
#
# Recipe (PART 70 + PART 71): restart VM/module, push tools, install kprobes
# (uvmt ON, uvsv/uvsvr OFF, 16 MB ring), start pin_gr_stream + semfix, launch
# cdt (strace -f -e ioctl) holding on failure, snapshot host+guest dmesg.
#
# Usage: tools/guest-diag/run_cycle.sh [hold-seconds] [strace|nostrace]
set -uo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
REPO="$HERE/../.."
VMEXEC="$REPO/scripts/vgpu_vm.py"
OUT=/tmp/opencode
HOLD="${1:-900}"
TRACE="${2:-strace}"

mkdir -p "$OUT"
echo "[*] vmrestart"
timeout 180 "$HERE/vmrestart.sh" || { echo "vmrestart failed"; exit 1; }

echo "[*] waiting for guest"
for i in $(seq 1 24); do
    if timeout 10 "$VMEXEC" exec "echo UP" 2>/dev/null | grep -q UP; then
        echo "    guest up"
        break
    fi
    sleep 5
done

echo "[*] setting module params"
sudo sh -c 'echo Y > /sys/module/open_iov_vfio/parameters/sema_fake_rel; echo Y > /sys/module/open_iov_vfio/parameters/sema_skip_ops'

echo "[*] pushing guest tools"
timeout 300 "$HERE/push_guest_tools.sh" >/dev/null 2>&1 || { echo "push failed"; exit 1; }

echo "[*] loading nvidia modules in guest (kprobe symbols need them)"
timeout 90 "$VMEXEC" exec "nvidia-smi -L >/dev/null 2>&1; lsmod | grep -q nvidia_uvm && echo UVM_LOADED" 2>/dev/null | grep -q UVM_LOADED || echo "    WARN: nvidia_uvm not loaded"

echo "[*] installing guest kprobes (uvmt ON; uvsv/uvsvr OFF; 16 MB ring)"
timeout 60 "$VMEXEC" exec "sh /tmp/kprobes.sh" 2>/dev/null | grep -E "ok|FAIL" | head -8
timeout 30 "$VMEXEC" exec "TR=/sys/kernel/tracing; [ -d \$TR ] || TR=/sys/kernel/debug/tracing; echo 16384 > \$TR/buffer_size_kb; echo 0 > \$TR/events/kprobes/uvsv/enable; echo 0 > \$TR/events/kprobes/uvsvr/enable; echo 1 > \$TR/events/kprobes/uvmt/enable; echo 1 > \$TR/events/kprobes/uvtw/enable; echo 1 > \$TR/events/kprobes/uvbp/enable; echo 1 > \$TR/events/kprobes/uvch/enable" >/dev/null 2>&1

echo "[*] starting host dmesg capture"
sudo dmesg -w > "$OUT/hostdmesg.log" 2>&1 &
DMESG_PID=$!

echo "[*] starting pin_gr_stream + semfix"
( setsid nohup env SEMFIX_SKIP_PB=1 SEMFIX_COUNTER_RETRY=10 \
    python3 "$HERE/pin_gr_stream.py" > "$OUT/pin_gr_stream.log" 2>&1 & )
sleep 1
( setsid nohup env SEMFIX_SKIP_PB=1 SEMFIX_COUNTER_RETRY=10 \
    python3 "$HERE/semfix_uvmt.py" 0.5 > "$OUT/semfix.log" 2>&1 & )
sleep 2

echo "[*] launching cdt (HOLD=$HOLD, $TRACE)"
if [ "$TRACE" = "strace" ]; then
    timeout 30 "$VMEXEC" exec \
        "rm -f /tmp/cdt.out /tmp/strace.log; CDT_HOLD=$HOLD nohup strace -f -tt -e trace=ioctl -o /tmp/strace.log /tmp/cdt >/tmp/cdt.out 2>&1 & echo STARTED" \
        >/dev/null 2>&1
elif [ "$TRACE" = "gdb" ]; then
    timeout 30 "$VMEXEC" exec "rm -f /tmp/cdt.out /tmp/bpwait.out; sync; echo CLEARED" >/dev/null 2>&1
    timeout 30 "$VMEXEC" exec \
        "rm -f /tmp/cdt.out /tmp/bpwait.out; CDT_PAUSE=60 CDT_HOLD=$HOLD nohup /tmp/cdt >/tmp/cdt.out 2>&1 & echo STARTED" \
        >/dev/null 2>&1
    PID=""
    for i in $(seq 1 12); do
        PID=$(timeout 20 "$VMEXEC" exec "pgrep -x cdt" 2>/dev/null | grep -E '^[0-9]+$' | head -1)
        [ -n "$PID" ] && break
        sleep 2
    done
    echo "    cdt pid=$PID; arming ${GDBSCRIPT:-bp999v2.py} (999/719 exit, serial+file)"
    ARMED=0
    for a in 1 2 3; do
        if timeout 40 "$VMEXEC" exec \
            "test -f /tmp/${GDBSCRIPT:-bp999v2.py} && nohup sh -c \"gdb -p '$PID' -batch -x /tmp/${GDBSCRIPT:-bp999v2.py} < /dev/null 2>&1 | tee /tmp/bpwait.out >/dev/ttyS0\" >/dev/null 2>&1 & echo ARMED" \
            2>/dev/null | grep -q ARMED; then
            ARMED=1
            break
        fi
        # The launch exec can time out even though nohup gdb started; verify
        # from the guest side before declaring the arm failed.
        if timeout 20 "$VMEXEC" exec "pgrep -f 'gdb -p' >/dev/null && echo GDBOK" 2>/dev/null | grep -q GDBOK; then
            echo "    arm attempt $a: exec timed out but gdb is running; continuing"
            ARMED=1
            break
        fi
        echo "    arm attempt $a failed; retrying"
        sleep 8
    done
    [ "$ARMED" = 1 ] || echo "    WARN: gdb arm failed 3x"
elif [ "$TRACE" = "gdb-launch" ]; then
    # Launch cdt INSIDE gdb so device-init code that runs before the
    # CDT_PAUSE (cuInit/cuDeviceGet) is visible to the breakpoints.
    timeout 30 "$VMEXEC" exec "rm -f /tmp/cdt.out /tmp/cdtgdb.out; sync; echo CLEARED" >/dev/null 2>&1
    timeout 30 "$VMEXEC" exec \
        "rm -f /tmp/cdt.out /tmp/cdtgdb.out; nohup sh -c 'gdb -batch -x /tmp/${GDBSCRIPT:-bp999v2.py} --args /tmp/cdt > /tmp/cdtgdb.out 2>&1' & echo STARTED" \
        >/dev/null 2>&1
    sleep 3
    PID=$(timeout 20 "$VMEXEC" exec "pgrep -x cdt" 2>/dev/null | grep -E '^[0-9]+$' | head -1)
    echo "    launched under gdb; cdt pid=${PID:-?} script=${GDBSCRIPT:-bp999v2.py}"
else
    timeout 30 "$VMEXEC" exec \
        "rm -f /tmp/cdt.out; CDT_HOLD=$HOLD nohup /tmp/cdt >/tmp/cdt.out 2>&1 & echo STARTED" \
        >/dev/null 2>&1
fi

echo "[*] waiting for the run to finish (watch cdt.out)"
POLL_FILE=/tmp/cdt.out
[ "$TRACE" = "gdb-launch" ] && POLL_FILE=/tmp/cdtgdb.out
for i in $(seq 1 90); do
    sleep 2
    OUT_TXT=$(timeout 20 "$VMEXEC" exec "cat $POLL_FILE 2>/dev/null" 2>/dev/null | tail -6)
    if echo "$OUT_TXT" | grep -q "FAILED rc="; then
        echo "    failure observed after ~$((i*2))s"
        break
    fi
done

echo "[*] snapshotting guest dmesg"
timeout 60 "$VMEXEC" exec "dmesg > /tmp/guestdmesg.log; cp /tmp/cdt.out /tmp/cdt.out.snap; [ -f /tmp/strace.log ] && cp /tmp/strace.log /tmp/strace.log.snap; [ -f /tmp/bpwait.out ] && cp /tmp/bpwait.out /tmp/bpwait.out.snap; [ -f /tmp/cdtgdb.out ] && cp /tmp/cdtgdb.out /tmp/bpwait.out.snap; echo DONE" >/dev/null 2>&1
sleep 1
sudo kill "$DMESG_PID" 2>/dev/null

echo "[*] collecting artifacts to $OUT"
for f in cdt.out.snap strace.log.snap bpwait.out.snap guestdmesg.log; do
    timeout 60 "$VMEXEC" exec "cat /tmp/$f" > "$OUT/$f" 2>/dev/null
done
echo "[*] done: $OUT/cdt.out.snap $OUT/strace.log.snap $OUT/guestdmesg.log $OUT/hostdmesg.log"
tail -5 "$OUT/cdt.out.snap" 2>/dev/null