#!/bin/bash

# Exit on error
set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
NC='\033[0m' # No Color

log() {
    echo -e "${GREEN}[+] $1${NC}"
}

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

check_root() {
    if [ "$EUID" -ne 0 ]; then
        error "Please run as root"
    fi
}

install_kernel() {
    log "Installing Proxmox kernel and headers..."
    apt update
    apt install -y proxmox-headers-6.8.12-1-pve proxmox-kernel-6.8.12-1-pve
}

install_build_tools() {
    log "Installing build tools..."
    apt install -y build-essential dkms
}

install_sriov_driver() {
    log "Installing strongtz i915 sriov dkms driver..."
    
    # Clone repository
    git clone https://github.com/strongtz/i915-sriov-dkms /root/i915-sriov-dkms
    cd /root/i915-sriov-dkms
    
    # Install DKMS module
    dkms add .
    dkms install -m i915-sriov-dkms -v $(cat VERSION) --force
}

configure_grub() {
    log "Configuring GRUB..."
    
    # Backup original grub config
    cp /etc/default/grub /etc/default/grub.backup
    
    # Update GRUB configuration
    sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT="\(.*\)"/GRUB_CMDLINE_LINUX_DEFAULT="\1 intel_iommu=on i915.enable_guc=3 i915.max_vfs=7"/' /etc/default/grub
    
    # Update grub
    update-grub
}

configure_sysfs() {
    log "Configuring sysfs for VFs..."
    
    # Install sysfsutils
    apt install -y sysfsutils
    
    # Add sysfs configuration
    echo "devices/pci0000:00/0000:00:02.0/sriov_numvfs = 7" > /etc/sysfs.conf
}

setup_debian_vm() {
    log "Setting up Debian VM requirements..."
    
    # Add backports repository
    echo "deb http://deb.debian.org/debian bookworm-backports main contrib non-free" >> /etc/apt/sources.list
    
    apt update
    
    # Install required packages
    apt install -y \
        linux-image-6.10.11+bpo-amd64 \
        linux-headers-6.10.11+bpo-amd64 \
        firmware-linux \
        git \
        dkms \
        vainfo \
        intel-media-va-driver
    
    # Install i915 DKMS module for VM
    kernel="6.10.11+bpo-amd64"
    
    git clone https://github.com/strongtz/i915-sriov-dkms.git /root/i915-sriov-dkms-vm
    cd /root/i915-sriov-dkms-vm
    rm -rf /usr/src/i915-sriov-dkms-* /var/lib/dkms/i915-sriov-dkms
    dkms add .
    dkms install -m i915-sriov-dkms -v $(cat VERSION) -k $kernel --force
    
    # Configure GRUB for VM
    echo 'GRUB_TIMEOUT=3' > /etc/default/grub.d/15_timeout.cfg
    echo 'GRUB_CMDLINE_LINUX_DEFAULT="i915.enable_guc=3"' > /etc/default/grub.d/i915.cfg
    
    # Update GRUB and initramfs
    update-grub
    update-initramfs -u -k all
}

main() {
    check_root
    
    read -p "Do you want to setup Proxmox host (h) or Debian VM (v)? " choice
    
    case "$choice" in
        h|H)
            log "Starting Proxmox host setup..."
            install_kernel
            install_build_tools
            install_sriov_driver
            configure_grub
            configure_sysfs
            log "Proxmox host setup complete. Please reboot your system."
            ;;
        v|V)
            log "Starting Debian VM setup..."
            setup_debian_vm
            log "Debian VM setup complete. Please reboot your system."
            ;;
        *)
            error "Invalid choice. Please select 'h' for host or 'v' for VM setup."
            ;;
    esac
}

main "$@"