#!/bin/bash
# BC-250 KIQ Fix - Post-Reboot Validation Script
# Run this after rebooting with the patched amdgpu module
#
# Tests:
# 1. Module loaded check
# 2. dmesg for KIQ fence timeouts (should be ZERO)
# 3. rocminfo
# 4. hip_probe (with normal exit, no _exit hack)
# 5. hip_vector_add
# 6. Multiple sequential GPU runs (stress test)
# 7. Process exit cleanup (the main fix target)

set +e  # Don't exit on error - we want to see all results

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 "      $1"; }

echo "=========================================="
echo "BC-250 KIQ Fix - Post-Reboot Validation"
echo "=========================================="
echo "Date: $(date)"
echo "Kernel: $(uname -r)"
echo ""

# Test 1: Module loaded
echo "--- Test 1: Module Status ---"
if lsmod | grep -q amdgpu; then
    pass "amdgpu module is loaded"
    MODULE_SIZE=$(lsmod | grep "^amdgpu " | awk '{print $2}')
    info "Module size: ${MODULE_SIZE}K"
else
    fail "amdgpu module not loaded!"
    exit 1
fi

# Test 2: Check for KIQ fence timeouts at boot
echo ""
echo "--- Test 2: KIQ Fence Timeouts (boot) ---"
KIQ_ERRORS=$(dmesg | grep -ci "kiq.*fence\|fence.*kiq\|KIQ.*timeout" 2>/dev/null || echo "0")
if [ "$KIQ_ERRORS" = "0" ]; then
    pass "No KIQ fence timeout errors in dmesg"
else
    warn "Found $KIQ_ERRORS KIQ-related messages (checking if they're errors...)"
    dmesg | grep -i "kiq.*fence\|fence.*kiq\|KIQ.*timeout" | head -5
fi

# Test 3: SDMA check (known issue, cosmetic)
echo ""
echo "--- Test 3: SDMA Status ---"
SDMA_ERRORS=$(dmesg | grep -ci "sdma.*fence\|sdma.*timeout" 2>/dev/null || echo "0")
if [ "$SDMA_ERRORS" = "0" ]; then
    pass "No SDMA fence errors"
else
    warn "SDMA errors present (cosmetic, mitigated by HSA_ENABLE_SDMA=0)"
fi

# Test 4: rocminfo
echo ""
echo "--- Test 4: rocminfo ---"
if timeout 30 rocminfo 2>/dev/null | grep -q "gfx10"; then
    pass "rocminfo detects GPU"
    rocminfo 2>/dev/null | grep "Name:\|Marketing Name:\|Compute Unit:" | head -5
else
    fail "rocminfo failed or timed out"
fi

# Test 5: hip_probe
echo ""
echo "--- Test 5: hip_probe ---"
HIP_PROBE="/home/dars/VibeROCm/hip_probe"
if [ -x "$HIP_PROBE" ]; then
    if timeout 30 "$HIP_PROBE" 2>&1; then
        pass "hip_probe completed successfully"
    else
        fail "hip_probe failed (exit code: $?)"
    fi
    
    # CRITICAL: Check for KIQ errors AFTER process exit
    sleep 3
    POST_KIQ=$(dmesg | tail -20 | grep -ci "kiq.*fence\|fence.*kiq\|KIQ.*timeout" 2>/dev/null || echo "0")
    if [ "$POST_KIQ" = "0" ]; then
        pass "No KIQ errors after hip_probe exit!"
    else
        fail "KIQ errors appeared after hip_probe exit"
        dmesg | tail -10
    fi
else
    warn "hip_probe not found at $HIP_PROBE"
fi

# Test 6: hip_vector_add
echo ""
echo "--- Test 6: hip_vector_add ---"
HIP_VECTOR="/home/dars/VibeROCm/hip_vector_add"
if [ -x "$HIP_VECTOR" ]; then
    if timeout 30 "$HIP_VECTOR" 2>&1; then
        pass "hip_vector_add completed successfully"
    else
        fail "hip_vector_add failed (exit code: $?)"
    fi
    
    sleep 3
    POST_KIQ=$(dmesg | tail -20 | grep -ci "kiq.*fence\|fence.*kiq\|KIQ.*timeout" 2>/dev/null || echo "0")
    if [ "$POST_KIQ" = "0" ]; then
        pass "No KIQ errors after hip_vector_add exit!"
    else
        fail "KIQ errors appeared after hip_vector_add exit"
    fi
else
    warn "hip_vector_add not found at $HIP_VECTOR"
fi

# Test 7: Multiple sequential runs (stress test for process lifecycle)
echo ""
echo "--- Test 7: Sequential GPU Stress (5 rounds) ---"
if [ -x "$HIP_VECTOR" ]; then
    STRESS_PASS=0
    for i in 1 2 3 4 5; do
        if timeout 30 "$HIP_VECTOR" >/dev/null 2>&1; then
            STRESS_PASS=$((STRESS_PASS + 1))
        else
            fail "Round $i failed"
            break
        fi
        sleep 1
    done
    
    if [ "$STRESS_PASS" -eq 5 ]; then
        pass "All 5 sequential GPU runs completed successfully!"
    else
        fail "Only $STRESS_PASS/5 rounds passed"
    fi
    
    # Final KIQ check
    sleep 5
    FINAL_KIQ=$(dmesg | tail -50 | grep -ci "kiq.*fence\|fence.*kiq\|KIQ.*timeout" 2>/dev/null || echo "0")
    if [ "$FINAL_KIQ" = "0" ]; then
        pass "No KIQ errors after stress test!"
    else
        fail "KIQ errors after stress test: $FINAL_KIQ"
        dmesg | tail -20 | grep -i "kiq\|fence"
    fi
else
    warn "Skipping stress test"
fi

# Test 8: rocminfo AFTER all GPU tests (this is the killer - previously hung)
echo ""
echo "--- Test 8: rocminfo After GPU Tests ---"
if timeout 30 rocminfo 2>/dev/null | grep -q "gfx10"; then
    pass "rocminfo still works after GPU tests! (Previously this HUNG)"
else
    fail "rocminfo hung or failed after GPU tests"
fi

echo ""
echo "=========================================="
echo "Validation complete."
echo "=========================================="