#!/bin/bash

# Kernel Update Script for AMD Strix Halo NPU Support
# This script updates the kernel to 6.12 LTS for NPU driver support

set -e  # Exit on any error

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

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

# Check if running as root
if [[ $EUID -eq 0 ]]; then
   error "This script should not be run as root. Run as regular user with sudo privileges."
   exit 1
fi

# Check if sudo is available
if ! command -v sudo &> /dev/null; then
    error "sudo is required but not installed."
    exit 1
fi

log "Starting kernel update for AMD Strix Halo NPU support..."

# Check current kernel version
CURRENT_KERNEL=$(uname -r)
log "Current kernel version: $CURRENT_KERNEL"

# Check if we're already on 6.12+
if [[ "$CURRENT_KERNEL" == "6.12"* ]] || [[ "$CURRENT_KERNEL" == "6.13"* ]] || [[ "$CURRENT_KERNEL" == "6.14"* ]]; then
    log "Kernel 6.12+ already installed. NPU drivers should be available."
    log "Checking for NPU drivers..."
    
    # Check for NPU drivers
    if lsmod | grep -q amdxdna; then
        log "NPU drivers are loaded!"
    else
        warn "NPU drivers not loaded. You may need to install amdxdna-tools."
        info "Try: sudo apt install amdxdna-tools"
    fi
    
    exit 0
fi

# Backup important data
log "Creating backup of important system files..."
sudo cp /etc/fstab /etc/fstab.backup.$(date +%Y%m%d_%H%M%S)
sudo cp /boot/grub/grub.cfg /boot/grub/grub.cfg.backup.$(date +%Y%m%d_%H%M%S)

# Update package lists
log "Updating package lists..."
sudo apt update

# Install required packages
log "Installing required packages..."
sudo apt install -y wget curl

# Check available kernel versions
log "Checking available kernel versions..."
KERNEL_VERSIONS=$(apt list --installed | grep linux-image | grep -E "6\.(12|13|14)" | head -5)
if [[ -z "$KERNEL_VERSIONS" ]]; then
    log "No kernel 6.12+ found in repositories. Installing from Ubuntu mainline..."
    
    # Install mainline kernel installer
    log "Installing mainline kernel installer..."
    sudo add-apt-repository -y ppa:cappelikan/ppa
    sudo apt update
    sudo apt install -y mainline
    
    # Download and install kernel 6.12
    log "Downloading kernel 6.12 LTS..."
    KERNEL_VERSION="6.12.0-061200"
    ARCH="amd64"
    
    # Create temporary directory
    TEMP_DIR=$(mktemp -d)
    cd "$TEMP_DIR"
    
    # Download kernel packages
    log "Downloading kernel packages..."
    wget "https://kernel.ubuntu.com/~kernel-ppa/mainline/v6.12/linux-headers-${KERNEL_VERSION}_all.deb"
    wget "https://kernel.ubuntu.com/~kernel-ppa/mainline/v6.12/linux-headers-${KERNEL_VERSION}-generic_${ARCH}.deb"
    wget "https://kernel.ubuntu.com/~kernel-ppa/mainline/v6.12/linux-image-unsigned-${KERNEL_VERSION}-generic_${ARCH}.deb"
    wget "https://kernel.ubuntu.com/~kernel-ppa/mainline/v6.12/linux-modules-${KERNEL_VERSION}-generic_${ARCH}.deb"
    
    # Install kernel packages
    log "Installing kernel packages..."
    sudo dpkg -i *.deb
    
    # Fix any dependency issues
    sudo apt install -f -y
    
    # Clean up
    cd /
    rm -rf "$TEMP_DIR"
    
else
    log "Kernel 6.12+ found in repositories. Installing..."
    sudo apt install -y linux-image-6.12.0-061200-generic linux-headers-6.12.0-061200-generic
fi

# Update GRUB
log "Updating GRUB bootloader..."
sudo update-grub

# Install NPU tools (if available)
log "Installing NPU tools..."
if apt list --available | grep -q amdxdna-tools; then
    sudo apt install -y amdxdna-tools
    log "NPU tools installed successfully!"
else
    warn "NPU tools not available in repositories yet."
    info "You may need to install them manually when they become available."
fi

# Create NPU test script
log "Creating NPU test script..."
cat > /tmp/test_npu_after_reboot.sh << 'EOF'
#!/bin/bash
echo "=== NPU Status After Kernel Update ==="
echo "Kernel version: $(uname -r)"
echo "NPU devices: $(ls /dev/amdxdna* 2>/dev/null || echo 'No NPU devices found')"
echo "NPU modules: $(lsmod | grep amdxdna || echo 'No NPU modules loaded')"
echo "NPU tools: $(which xrt-smi 2>/dev/null || echo 'NPU tools not found')"
EOF
chmod +x /tmp/test_npu_after_reboot.sh

log "Kernel update completed successfully!"
log "IMPORTANT: You need to reboot your system to use the new kernel."
log ""
warn "Before rebooting:"
info "1. Save all your work"
info "2. Close all applications"
info "3. Run: sudo reboot"
info ""
info "After rebooting, run: /tmp/test_npu_after_reboot.sh"
info ""
log "The new kernel will enable NPU drivers for your AMD Strix Halo NPU!"
log "This will provide 5-100x speedup for AI workloads compared to GPU."

# Ask user if they want to reboot now
read -p "Do you want to reboot now? (y/N): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
    log "Rebooting in 10 seconds... Press Ctrl+C to cancel"
    sleep 10
    sudo reboot
else
    log "Please reboot manually when ready: sudo reboot"
fi