#!/bin/bash

#============================================================================
# TUXEDO AMD NPU Driver Installer
# Version: 1.0
# Date: 2025-08-05
# 
# Developed by: Dragonfire Technologies
# For: TUXEDO Computers & AMD Ryzen AI Users
#============================================================================

set -e

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

# Banner
clear
echo -e "${BLUE}"
echo "════════════════════════════════════════════════════════════════"
echo "     TUXEDO AMD NPU Driver Installer v1.0"
echo "     Dragonfire Technologies × TUXEDO Computers"
echo "════════════════════════════════════════════════════════════════"
echo -e "${NC}"

# Functions
print_status() {
    echo -e "${GREEN}[✓]${NC} $1"
}

print_error() {
    echo -e "${RED}[✗]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[!]${NC} $1"
}

print_info() {
    echo -e "${BLUE}[i]${NC} $1"
}

# Check if running as root
if [[ $EUID -eq 0 ]]; then
   print_error "Please run without sudo (installer will request privileges when needed)"
   exit 1
fi

# System requirements check
print_info "Checking system requirements..."

# Check OS
if [ -f /etc/os-release ]; then
    . /etc/os-release
    if [[ "$ID" == "tuxedo" ]] || [[ "$ID" == "ubuntu" ]] || [[ "$ID_LIKE" == *"ubuntu"* ]]; then
        print_status "OS: $PRETTY_NAME"
    else
        print_warning "Untested OS: $PRETTY_NAME (developed for TUXEDO OS/Ubuntu)"
        read -p "Continue anyway? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            exit 1
        fi
    fi
fi

# Check kernel version
KERNEL_VERSION=$(uname -r)
KERNEL_MAJOR=$(echo $KERNEL_VERSION | cut -d. -f1)
KERNEL_MINOR=$(echo $KERNEL_VERSION | cut -d. -f2)

if [ "$KERNEL_MAJOR" -ge 6 ] && [ "$KERNEL_MINOR" -ge 10 ]; then
    print_status "Kernel: $KERNEL_VERSION"
else
    print_error "Kernel $KERNEL_VERSION is too old (need 6.10+)"
    exit 1
fi

# Check for AMD NPU hardware
print_info "Detecting AMD NPU hardware..."
if lspci | grep -qi "neural"; then
    NPU_DEVICE=$(lspci | grep -i "neural" | head -1)
    print_status "NPU found: $NPU_DEVICE"
else
    print_warning "No NPU detected in PCI devices"
    print_info "Continuing installation anyway..."
fi

# Check for existing installation
if lsmod | grep -q amdxdna 2>/dev/null; then
    print_warning "NPU driver already loaded"
    read -p "Reinstall/Update? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        print_info "Installation cancelled"
        exit 0
    fi
fi

# Check packages exist
PACKAGES_DIR="$(dirname "$0")/packages"
if [ ! -d "$PACKAGES_DIR" ]; then
    print_error "Packages directory not found"
    exit 1
fi

XRT_BASE=$(ls $PACKAGES_DIR/xrt_*base.deb 2>/dev/null | head -1)
XRT_PLUGIN=$(ls $PACKAGES_DIR/xrt_plugin*.deb 2>/dev/null | head -1)

if [ ! -f "$XRT_BASE" ] || [ ! -f "$XRT_PLUGIN" ]; then
    print_error "Required packages not found in $PACKAGES_DIR"
    exit 1
fi

# Installation
echo
print_info "Starting installation..."
echo

# Install XRT base
print_info "Installing XRT runtime..."
sudo dpkg -i "$XRT_BASE" 2>/dev/null || {
    print_warning "Fixing dependencies..."
    sudo apt-get install -f -y
}
print_status "XRT runtime installed"

# Install XDNA plugin
print_info "Installing AMD NPU driver (this may take a minute)..."
sudo dpkg -i "$XRT_PLUGIN" 2>/dev/null || {
    print_warning "Fixing dependencies..."
    sudo apt-get install -f -y
}
print_status "AMD NPU driver installed"

# Verify installation
echo
print_info "Verifying installation..."

if lsmod | grep -q amdxdna; then
    print_status "Kernel module loaded"
else
    print_error "Kernel module not loaded"
    print_info "Trying to load manually..."
    sudo modprobe amdxdna || {
        print_error "Failed to load driver"
        exit 1
    }
fi

if [ -e /dev/accel/accel0 ]; then
    print_status "NPU device created: /dev/accel/accel0"
else
    print_warning "NPU device not found (may appear after reboot)"
fi

# Test with XRT
if [ -f /opt/xilinx/xrt/setup.sh ]; then
    source /opt/xilinx/xrt/setup.sh 2>/dev/null
    print_status "XRT environment configured"
    
    # Quick test
    if command -v xrt-smi &> /dev/null; then
        echo
        print_info "NPU Status:"
        xrt-smi examine 2>/dev/null | grep -A5 "Device" || true
    fi
fi

# Success message
echo
echo -e "${GREEN}════════════════════════════════════════════════════════════════${NC}"
echo -e "${GREEN}     NPU Driver Installation Complete!${NC}"
echo -e "${GREEN}════════════════════════════════════════════════════════════════${NC}"
echo
print_info "Next steps:"
echo "  1. Run: ./verify_npu.sh    # To verify NPU functionality"
echo "  2. Run: ./test_npu.sh       # To test NPU performance"
echo "  3. Run: source /opt/xilinx/xrt/setup.sh  # Before using XRT tools"
echo
print_info "Your AMD NPU is ready for AI acceleration!"
echo
echo -e "${BLUE}Powered by Dragonfire Technologies${NC}"
echo -e "${BLUE}Built for TUXEDO Computers${NC}"