#!/bin/bash
# =============================================================================
# Step 7: Compile and Run HIP Tests — AMD BC-250
# =============================================================================
# Compiles and runs all 3 HIP test programs, plus a sequential stress test.
#
# Tests:
#   1. hip_probe       — Device query, memory allocation (safe probe)
#   2. hip_minimal_test — Kernel launch, compute verification (32 values)
#   3. hip_vector_add   — Full managed-memory vector addition (65536 elements)
#   4. Sequential stress — Run hip_vector_add 3 times back-to-back
#
# Run as: bash 07_run_hip_tests.sh
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
GUIDE_DIR="$(dirname "$SCRIPT_DIR")"
TESTS_DIR="$GUIDE_DIR/tests"
BUILD_DIR="$GUIDE_DIR/tests/bin"

# Ensure environment
export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HIP_VISIBLE_DEVICES=0
export HSA_ENABLE_SDMA=0
export HSA_TOOLS_LIB=""
export HSA_TOOLS_REPORT_LOAD_FAILURE=0

PASS=0
FAIL=0

echo "============================================"
echo " HIP Test Suite — AMD BC-250"
echo "============================================"
echo ""

# ─────────────────────────────────────────────
# Compile
# ─────────────────────────────────────────────
echo "--- Compiling Tests ---"
mkdir -p "$BUILD_DIR"

TESTS=("hip_probe" "hip_minimal_test" "hip_vector_add")
for TEST in "${TESTS[@]}"; do
    SRC="$TESTS_DIR/${TEST}.cpp"
    BIN="$BUILD_DIR/$TEST"
    
    if [[ ! -f "$SRC" ]]; then
        echo "  [SKIP] $TEST — source not found: $SRC"
        continue
    fi
    
    echo "  Compiling $TEST..."
    if hipcc --offload-arch=gfx1010 -o "$BIN" "$SRC" 2>&1; then
        echo "  [OK] $TEST compiled ($(du -h "$BIN" | cut -f1))"
    else
        echo "  [FAIL] $TEST compilation failed"
        ((FAIL++))
    fi
done

echo ""

# ─────────────────────────────────────────────
# Test 1: hip_probe
# ─────────────────────────────────────────────
echo "--- Test 1: hip_probe ---"
BIN="$BUILD_DIR/hip_probe"
if [[ -x "$BIN" ]]; then
    OUTPUT=$("$BIN" 2>&1) && RC=$? || RC=$?
    echo "$OUTPUT" | head -20
    if echo "$OUTPUT" | grep -q "ALL STEPS COMPLETED"; then
        echo "  [PASS] hip_probe"
        ((PASS++))
    else
        echo "  [FAIL] hip_probe (exit code: $RC)"
        ((FAIL++))
    fi
else
    echo "  [SKIP] Not compiled"
fi

echo ""

# ─────────────────────────────────────────────
# Test 2: hip_minimal_test
# ─────────────────────────────────────────────
echo "--- Test 2: hip_minimal_test ---"
BIN="$BUILD_DIR/hip_minimal_test"
if [[ -x "$BIN" ]]; then
    OUTPUT=$("$BIN" 2>&1) && RC=$? || RC=$?
    echo "$OUTPUT" | head -20
    if echo "$OUTPUT" | grep -q "PASS"; then
        echo "  [PASS] hip_minimal_test"
        ((PASS++))
    else
        echo "  [FAIL] hip_minimal_test (exit code: $RC)"
        ((FAIL++))
    fi
else
    echo "  [SKIP] Not compiled"
fi

echo ""

# ─────────────────────────────────────────────
# Test 3: hip_vector_add
# ─────────────────────────────────────────────
echo "--- Test 3: hip_vector_add ---"
BIN="$BUILD_DIR/hip_vector_add"
if [[ -x "$BIN" ]]; then
    OUTPUT=$("$BIN" 2>&1) && RC=$? || RC=$?
    echo "$OUTPUT" | head -20
    if echo "$OUTPUT" | grep -q "FULLY OPERATIONAL"; then
        echo "  [PASS] hip_vector_add"
        ((PASS++))
    else
        echo "  [FAIL] hip_vector_add (exit code: $RC)"
        ((FAIL++))
    fi
else
    echo "  [SKIP] Not compiled"
fi

echo ""

# ─────────────────────────────────────────────
# Test 4: Sequential Stress
# ─────────────────────────────────────────────
echo "--- Test 4: Sequential Stress (3 rounds) ---"
BIN="$BUILD_DIR/hip_vector_add"
if [[ -x "$BIN" ]]; then
    STRESS_PASS=0
    for i in 1 2 3; do
        OUTPUT=$("$BIN" 2>&1) && RC=$? || RC=$?
        if echo "$OUTPUT" | grep -q "FULLY OPERATIONAL"; then
            echo "  Round $i/3: PASS"
            ((STRESS_PASS++))
        else
            echo "  Round $i/3: FAIL (exit code: $RC)"
        fi
    done
    
    if [[ "$STRESS_PASS" -eq 3 ]]; then
        echo "  [PASS] Sequential stress: $STRESS_PASS/3"
        ((PASS++))
    else
        echo "  [FAIL] Sequential stress: $STRESS_PASS/3"
        ((FAIL++))
    fi
else
    echo "  [SKIP] hip_vector_add not compiled"
fi

echo ""

# ─────────────────────────────────────────────
# Final Check
# ─────────────────────────────────────────────
echo "--- Post-Test Stability ---"
KIQ_COUNT=$(dmesg 2>/dev/null | grep -c "timeout waiting for kiq fence" || echo 0)
echo "  KIQ timeout count: $KIQ_COUNT"

ROCM_OK=$(HSA_OVERRIDE_GFX_VERSION=10.1.0 rocminfo 2>&1 | grep -c "Done" || echo 0)
if [[ "$ROCM_OK" -gt 0 ]]; then
    echo "  rocminfo: Still working"
    ((PASS++))
else
    echo "  rocminfo: FAILED after tests"
    ((FAIL++))
fi

echo ""
echo "============================================"
echo " Results: $PASS passed, $FAIL failed"
echo "============================================"

if [[ "$FAIL" -eq 0 ]]; then
    echo ""
    echo " ROCm is FULLY OPERATIONAL on your BC-250!"
    echo ""
    echo " Your GPU is ready for HIP compute workloads."
    echo " See the README.md 'HIP Programming Guidelines' section"
    echo " for BC-250-specific coding rules."
fi

exit "$FAIL"