#!/bin/bash

# AMD NPU Test Script
# Author: NPU Setup Assistant
# Date: 2025-08-05

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
NPU_DIR="$(dirname "$SCRIPT_DIR")"

echo "======================================"
echo "AMD NPU Test Script"
echo "======================================"

# Function to print colored output
print_status() {
    echo -e "\033[1;32m[INFO]\033[0m $1"
}

print_error() {
    echo -e "\033[1;31m[ERROR]\033[0m $1"
}

# Setup XRT environment
source /opt/xilinx/xrt/setup.sh 2>/dev/null || {
    print_error "XRT not installed. Please run install_npu_driver.sh first"
    exit 1
}

# Test 1: Check kernel module
print_status "Checking kernel module..."
if lsmod | grep -q amdxdna; then
    echo "  ✓ amdxdna module loaded"
    lsmod | grep amdxdna
else
    print_error "amdxdna module not loaded"
    exit 1
fi

# Test 2: Check device node
print_status "Checking device node..."
if [ -e /dev/accel/accel0 ]; then
    echo "  ✓ NPU device present"
    ls -la /dev/accel/accel0
else
    print_error "NPU device not found"
    exit 1
fi

# Test 3: Check PCI device
print_status "Checking PCI device..."
if lspci | grep -i "neural"; then
    echo "  ✓ NPU hardware detected"
    lspci -vnn | grep -A 10 "Neural Processing Unit"
fi

# Test 4: Run xrt-smi examine
print_status "Running XRT examination..."
xrt-smi examine

# Test 5: Run example test if available
TEST_BIN="$NPU_DIR/src/xdna-driver/build/example_build/example_noop_test"
TEST_XCLBIN="$NPU_DIR/src/xdna-driver/tools/bins/1502_00/validate.xclbin"

if [ -f "$TEST_BIN" ] && [ -f "$TEST_XCLBIN" ]; then
    print_status "Running example test..."
    cd "$(dirname "$TEST_BIN")"
    ./example_noop_test "$TEST_XCLBIN"
else
    print_status "Example test not available (build with ./build.sh -example to enable)"
fi

print_status "All tests completed successfully!"