#!/bin/bash
# BC-250 Safe Post-Reboot Test Script (v2)
# Tests GPU functionality after KIQ bypass patch installation
# CRITICAL: This script performs INCREMENTAL testing with safety checks
set -e

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

pass() { echo -e "${GREEN}[PASS]${NC} $1"; }
fail() { echo -e "${RED}[FAIL]${NC} $1"; }
warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
info() { echo -e "[INFO] $1"; }

ERRORS=0
TESTS=0

echo "==========================================="
echo " BC-250 KIQ Bypass Patch - Verification"
echo " Date: $(date)"
echo "==========================================="
echo ""

# === STEP 0: Verify patch markers in dmesg ===
((TESTS++))
info "Step 0: Checking dmesg for patch markers..."
if dmesg 2>/dev/null | grep -q "BC-250 KIQ bypass active"; then
    pass "BC-250 KIQ bypass markers found in dmesg"
    dmesg | grep "BC-250" | while read line; do echo "  → $line"; done
else
    warn "No BC-250 bypass markers in dmesg yet (may appear after first GPU use)"
fi

# === STEP 1: Check for KIQ errors at boot ===
((TESTS++))
info "Step 1: Checking for KIQ errors in boot log..."
if dmesg 2>/dev/null | grep -q "timeout waiting for kiq fence"; then
    fail "KIQ fence timeout detected at boot!"
    ((ERRORS++))
    echo "  → GPU is likely in a bad state. Do NOT proceed."
    exit 1
else
    pass "No KIQ fence timeouts at boot"
fi

# === STEP 2: Check GPU is alive ===
((TESTS++))
info "Step 2: Checking GPU hardware..."
if ls /sys/class/drm/card0/device/ &>/dev/null; then
    pass "GPU DRM device present"
else
    fail "No GPU DRM device"
    ((ERRORS++))
    exit 1
fi

# === STEP 3: Check Z-Image-Turbo is disabled ===
((TESTS++))
info "Step 3: Checking Z-Image-Turbo service..."
if systemctl --user is-active zimage.service 2>/dev/null | grep -q "^active"; then
    fail "Z-Image-Turbo is RUNNING! Stop it first: systemctl --user stop zimage"
    ((ERRORS++))
    exit 1
else
    pass "Z-Image-Turbo is not running"
fi

# === STEP 4: rocminfo (safe, no TLB flush) ===
((TESTS++))
info "Step 4: Running rocminfo..."
if rocminfo 2>&1 | grep -q "gfx10"; then
    pass "rocminfo detects GPU"
else
    fail "rocminfo failed"
    ((ERRORS++))
fi

# Check dmesg AGAIN after rocminfo
if dmesg 2>/dev/null | grep -q "timeout waiting for kiq fence"; then
    fail "KIQ error appeared after rocminfo!"
    ((ERRORS++))
    exit 1
fi

echo ""
info "Step 4 complete. Checking dmesg for any BC-250 markers..."
dmesg 2>/dev/null | grep "BC-250" 2>/dev/null || info "(no BC-250 markers yet)"
echo ""

# === STEP 5: hip_probe (uses _exit(0), minimal cleanup) ===
((TESTS++))
info "Step 5: Running hip_probe (minimal GPU test)..."
cd /home/dars/VibeROCm

if [ -f hip_probe ]; then
    timeout 30 ./hip_probe 2>&1
    RESULT=$?
    if [ $RESULT -eq 0 ]; then
        pass "hip_probe passed"
    else
        fail "hip_probe failed (exit code: $RESULT)"
        ((ERRORS++))
    fi
else
    warn "hip_probe binary not found, skipping"
fi

# Check dmesg after hip_probe
sleep 2
if dmesg 2>/dev/null | grep -q "timeout waiting for kiq fence"; then
    fail "KIQ error appeared after hip_probe!"
    ((ERRORS++))
    echo "  → The KIQ bypass patch may not be working. STOP HERE."
    exit 1
fi

info "Checking BC-250 bypass markers after hip_probe..."
dmesg 2>/dev/null | grep "BC-250" 2>/dev/null || info "(no markers)"
echo ""

# === STEP 6: Wait and verify GPU is still alive ===
info "Step 6: Waiting 10 seconds to check GPU stability..."
sleep 10
if dmesg 2>/dev/null | grep -q "timeout waiting for kiq fence\|GPU fault\|GPU hang"; then
    fail "GPU error detected during wait period!"
    ((ERRORS++))
    exit 1
else
    pass "GPU stable after 10-second wait"
fi

# === STEP 7: hip_vector_add (full GPU compute test) ===
((TESTS++))
info "Step 7: Running hip_vector_add (full compute test)..."
if [ -f hip_vector_add ]; then
    timeout 60 ./hip_vector_add 2>&1
    RESULT=$?
    if [ $RESULT -eq 0 ]; then
        pass "hip_vector_add passed"
    else
        fail "hip_vector_add failed (exit code: $RESULT)"
        ((ERRORS++))
    fi
else
    warn "hip_vector_add binary not found, skipping"
fi

# Final dmesg check
sleep 5
if dmesg 2>/dev/null | grep -q "timeout waiting for kiq fence"; then
    fail "KIQ error appeared after hip_vector_add!"
    ((ERRORS++))
fi

echo ""
echo "==========================================="
info "Final dmesg BC-250 markers:"
dmesg 2>/dev/null | grep "BC-250" 2>/dev/null || info "(none)"
echo ""

if [ $ERRORS -eq 0 ]; then
    echo -e "${GREEN}ALL TESTS PASSED${NC} ($TESTS tests)"
    echo ""
    echo "The KIQ bypass patch is working. The GPU survived compute workloads"
    echo "without hanging. You can now safely re-enable Z-Image-Turbo:"
    echo "  systemctl --user enable --now zimage.service"
else
    echo -e "${RED}$ERRORS ERRORS${NC} out of $TESTS tests"
    echo ""
    echo "The patch may need additional work. Check dmesg for details."
fi