#!/bin/bash

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

set -e

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

echo "======================================"
echo "AMD NPU Driver Installation 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"
}

# Check if running as root
if [[ $EUID -eq 0 ]]; then
   print_error "This script should not be run as root (it will ask for sudo when needed)"
   exit 1
fi

# Check if packages exist
if [ ! -d "$PACKAGES_DIR" ] || [ -z "$(ls -A $PACKAGES_DIR/*.deb 2>/dev/null)" ]; then
    print_error "No packages found in $PACKAGES_DIR"
    echo "Please run build_npu_driver.sh first"
    exit 1
fi

# Step 1: Install XRT base package
print_status "Installing XRT base package..."
XRT_BASE=$(ls $PACKAGES_DIR/xrt_*base.deb 2>/dev/null | head -1)
if [ -f "$XRT_BASE" ]; then
    sudo dpkg -i "$XRT_BASE"
    sudo apt-get install -f -y
else
    print_error "XRT base package not found"
    exit 1
fi

# Step 2: Install XDNA plugin
print_status "Installing XDNA driver plugin..."
XDNA_PLUGIN=$(ls $PACKAGES_DIR/xrt_plugin*.deb 2>/dev/null | head -1)
if [ -f "$XDNA_PLUGIN" ]; then
    sudo dpkg -i "$XDNA_PLUGIN"
    sudo apt-get install -f -y
else
    print_error "XDNA plugin package not found"
    exit 1
fi

# Step 3: Verify installation
print_status "Verifying installation..."
if lsmod | grep -q amdxdna; then
    print_status "amdxdna kernel module loaded successfully"
else
    print_error "amdxdna kernel module not loaded"
    exit 1
fi

if [ -e /dev/accel/accel0 ]; then
    print_status "NPU device /dev/accel/accel0 created successfully"
else
    print_error "NPU device not found"
    exit 1
fi

# Step 4: Test with xrt-smi
print_status "Testing NPU with xrt-smi..."
source /opt/xilinx/xrt/setup.sh
xrt-smi examine

print_status "Installation complete!"
echo ""
echo "NPU is ready to use. You can run AI workloads now."
echo "To use XRT tools, always run: source /opt/xilinx/xrt/setup.sh"