#!/bin/bash

#============================================================================
# NPU Verification Script
# Dragonfire Technologies
#============================================================================

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

echo -e "${BLUE}═══════════════════════════════════════════${NC}"
echo -e "${BLUE}     NPU Verification Tool${NC}"
echo -e "${BLUE}═══════════════════════════════════════════${NC}"
echo

# Test functions
check_pass() {
    echo -e "${GREEN}✓${NC} $1"
}

check_fail() {
    echo -e "${RED}✗${NC} $1"
    FAILED=1
}

check_warn() {
    echo -e "${YELLOW}⚠${NC} $1"
}

FAILED=0

# 1. Kernel Module Check
echo -e "${BLUE}Checking kernel module...${NC}"
if lsmod | grep -q amdxdna; then
    MODULE_INFO=$(lsmod | grep amdxdna | awk '{print "Size: "$2" bytes, Used by: "$3" processes"}')
    check_pass "amdxdna module loaded ($MODULE_INFO)"
else
    check_fail "amdxdna module not loaded"
fi

# 2. Device Node Check
echo -e "\n${BLUE}Checking device node...${NC}"
if [ -e /dev/accel/accel0 ]; then
    DEVICE_INFO=$(ls -la /dev/accel/accel0 | awk '{print "Permissions: "$1", Owner: "$3":"$4}')
    check_pass "/dev/accel/accel0 exists ($DEVICE_INFO)"
else
    check_fail "/dev/accel/accel0 not found"
fi

# 3. PCI Hardware Check
echo -e "\n${BLUE}Checking PCI hardware...${NC}"
if lspci | grep -qi "neural"; then
    NPU_PCI=$(lspci | grep -i "neural" | head -1 | cut -d' ' -f1)
    NPU_NAME=$(lspci | grep -i "neural" | head -1 | cut -d' ' -f5-)
    check_pass "NPU detected at PCI [$NPU_PCI] - $NPU_NAME"
else
    check_warn "No NPU in PCI device list (may be normal on some systems)"
fi

# 4. XRT Installation Check
echo -e "\n${BLUE}Checking XRT installation...${NC}"
if [ -d /opt/xilinx/xrt ]; then
    check_pass "XRT installed at /opt/xilinx/xrt"
    
    if [ -f /opt/xilinx/xrt/setup.sh ]; then
        source /opt/xilinx/xrt/setup.sh 2>/dev/null
        
        if command -v xrt-smi &> /dev/null; then
            XRT_VERSION=$(xrt-smi version 2>/dev/null | grep "Version" | head -1 | awk '{print $3}')
            check_pass "XRT tools available (version: $XRT_VERSION)"
        else
            check_fail "XRT tools not found in PATH"
        fi
    else
        check_fail "XRT setup script not found"
    fi
else
    check_fail "XRT not installed"
fi

# 5. Firmware Check
echo -e "\n${BLUE}Checking NPU firmware...${NC}"
if [ -d /lib/firmware/amdnpu ]; then
    FW_COUNT=$(ls /lib/firmware/amdnpu/*.bin 2>/dev/null | wc -l)
    if [ "$FW_COUNT" -gt 0 ]; then
        check_pass "NPU firmware found ($FW_COUNT files)"
    else
        check_fail "No firmware files in /lib/firmware/amdnpu"
    fi
else
    check_fail "Firmware directory /lib/firmware/amdnpu not found"
fi

# 6. DKMS Check
echo -e "\n${BLUE}Checking DKMS status...${NC}"
if command -v dkms &> /dev/null; then
    if dkms status 2>/dev/null | grep -q "xrt-amdxdna"; then
        DKMS_STATUS=$(dkms status | grep "xrt-amdxdna" | head -1)
        check_pass "DKMS configured: $DKMS_STATUS"
    else
        check_warn "DKMS module not found (driver may not persist after kernel updates)"
    fi
else
    check_warn "DKMS not installed"
fi

# 7. Quick NPU Test
echo -e "\n${BLUE}Running NPU quick test...${NC}"
if [ -f /opt/xilinx/xrt/setup.sh ]; then
    source /opt/xilinx/xrt/setup.sh 2>/dev/null
    if command -v xrt-smi &> /dev/null; then
        NPU_INFO=$(xrt-smi examine 2>/dev/null | grep "NPU" | head -1)
        if [ -n "$NPU_INFO" ]; then
            check_pass "NPU responsive: $NPU_INFO"
        else
            check_warn "NPU not detected by xrt-smi"
        fi
    fi
fi

# Results
echo
echo -e "${BLUE}═══════════════════════════════════════════${NC}"
if [ $FAILED -eq 0 ]; then
    echo -e "${GREEN}     VERIFICATION PASSED${NC}"
    echo -e "${GREEN}     NPU is ready for use!${NC}"
else
    echo -e "${YELLOW}     VERIFICATION INCOMPLETE${NC}"
    echo -e "${YELLOW}     Some checks failed - see above${NC}"
fi
echo -e "${BLUE}═══════════════════════════════════════════${NC}"

exit $FAILED