#!/bin/bash

# Script to apply WiFi driver corrections to a running system
# This script extracts pre-compiled WiFi modules and firmware to the live system

echo "=== WiFi Driver Correction Script ==="

# Step 0: Ask for sudo privileges early
echo "This script requires root privileges to modify system files."
sudo -v
if [[ $? -ne 0 ]]; then
    echo "ERROR: sudo access required. Exiting."
    exit 1
fi

# Step 0.5: Create Stereolabs module blacklist
echo "=== Step 0.5: Creating Stereolabs module blacklist ==="
BLACKLIST_FILE="/etc/modprobe.d/sl_blacklist.conf"
echo "Creating $BLACKLIST_FILE to blacklist Stereolabs modules"
sudo bash -c "cat > $BLACKLIST_FILE" << 'EOF'
blacklist sl_zedx
blacklist sl_zedxpro
blacklist sl_zedxone_uhd
blacklist sl_max9295
blacklist sl_max9296
EOF
if [[ $? -eq 0 ]]; then
    echo "Stereolabs module blacklist created successfully"
else
    echo "WARNING: Failed to create Stereolabs module blacklist"
fi

# Step 1: Detect L4T version automatically, if nv_tegra_release file exists, otherwise prompt user
TEGRA_RELEASE_FILE="/etc/nv_tegra_release"
if [[ ! -f "$TEGRA_RELEASE_FILE" ]]; then
    echo "WARNING: Could not find $TEGRA_RELEASE_FILE"
    echo "This usually means you're not on a Jetson device or the file is missing."
    echo ""
    echo "Please manually specify your Jetson version:"
    echo "1) 36_4 (JetPack 6.2)"
    echo "2) 36_3 (JetPack 5.0)" 
    echo "3) 35_4 (JetPack 5.1.2)"
    echo ""
    read -p "Enter your choice (1-3): " choice
    
    case $choice in
        1) L4T_VERSION="36_4" ;;
        2) L4T_VERSION="36_3" ;;
        3) L4T_VERSION="35_4" ;;
        *) 
            echo "ERROR: Invalid choice. Please run the script again and select 1, 2, or 3."
            exit 1
            ;;
    esac
    
    echo "Using Tegra version: $L4T_VERSION"
else
	# Extract version info (e.g., "36_4" from R36 REVISION: 4.3)
	L4T_VERSION=$(head -n 1 "$TEGRA_RELEASE_FILE" | sed -n 's/.*R\([0-9]*\).*REVISION: \([0-9]*\)\.[0-9]*.*/\1_\2/p')
	if [[ -z "$L4T_VERSION" ]]; then
	    echo "ERROR: Could not detect L4T version from $TEGRA_RELEASE_FILE"
	    exit 1
	fi

	echo "Detected L4T version: $L4T_VERSION"
fi

# Check if AGX (only for 35.4)
if [[ "$L4T_VERSION" == "35_4" ]] && grep -q "AGX" /proc/device-tree/model 2>/dev/null; then
    L4T_VERSION="${L4T_VERSION}_agx"
    echo "Detected AGX model on 35.4, updated L4T version to: $L4T_VERSION"
fi

# Detect kernel version
KERNEL_VERSION=$(uname -r)
echo "Detected kernel version: $KERNEL_VERSION"

# Check if kernel is RT and set rt variable accordingly
if echo "$KERNEL_VERSION" | grep -q "rt"; then
    RT_SUFFIX="_rt"
    echo "Detected RT kernel"
else
    RT_SUFFIX=""
    echo "Detected standard kernel"
fi

# Step 2: Define archive names based on detected version
IWLWIFI_ARCHIVE="iwlwifi_bpcore94${RT_SUFFIX}_${L4T_VERSION}.tar"
FIRMWARE_ARCHIVE="iwlwifi-fw-2025-03-11.tar"

echo "Looking for WiFi module archive: $IWLWIFI_ARCHIVE"
echo "Looking for firmware archive: $FIRMWARE_ARCHIVE"

# Download archives if not present or if wrong version exists
if [[ ! -f "$IWLWIFI_ARCHIVE" ]]; then
    echo "Downloading WiFi modules archive..."
    wget "https://stereolabs.sfo2.cdn.digitaloceanspaces.com/utils/zed_boxes/$IWLWIFI_ARCHIVE"
    if [[ $? -ne 0 ]]; then
        echo "ERROR: Could not download WiFi modules archive: $IWLWIFI_ARCHIVE"
        exit 1
    fi
else
    echo "WiFi modules archive already exists: $IWLWIFI_ARCHIVE"
fi

if [[ ! -f "$FIRMWARE_ARCHIVE" ]]; then
    echo "Downloading WiFi firmware archive..."
    wget -N "https://stereolabs.sfo2.cdn.digitaloceanspaces.com/utils/zed_boxes/$FIRMWARE_ARCHIVE"
    if [[ $? -ne 0 ]]; then
        echo "ERROR: Could not download WiFi firmware archive: $FIRMWARE_ARCHIVE"
        exit 1
    fi
fi

# Ask for confirmation
echo ""
echo "This script will:"
echo "  - Remove existing WiFi module metadata files"
echo "  - Extract WiFi kernel modules for $KERNEL_VERSION"
echo "  - Extract WiFi firmware files"
echo "  - Add iwlwifi module to autoload"
echo ""
read -p "Continue? (Y/n): " -n 1 -r
echo
if [[ $REPLY =~ ^[Nn]$ ]]; then
    echo "Operation cancelled"
    exit 0
fi

echo "=== Step 1: Cleaning existing WiFi module metadata ==="
# Remove module metadata files that need regeneration
sudo rm -f /lib/modules/${KERNEL_VERSION}/modules.alias
sudo rm -f /lib/modules/${KERNEL_VERSION}/modules.alias.bin
sudo rm -f /lib/modules/${KERNEL_VERSION}/modules.dep
sudo rm -f /lib/modules/${KERNEL_VERSION}/modules.dep.bin
sudo rm -f /lib/modules/${KERNEL_VERSION}/modules.symbols
sudo rm -f /lib/modules/${KERNEL_VERSION}/modules.symbols.bin

echo "=== Step 2: Removing old WiFi modules ==="
# Remove specific old WiFi modules to ensure clean replacement
sudo rm -f /lib/modules/${KERNEL_VERSION}/updates/net/mac80211/mac80211.ko
sudo rm -f /lib/modules/${KERNEL_VERSION}/updates/net/wireless/cfg80211.ko
sudo rm -f /lib/modules/${KERNEL_VERSION}/updates/compat/compat.ko
sudo rm -f /lib/modules/${KERNEL_VERSION}/updates/drivers/net/wireless/intel/iwlwifi/iwlwifi.ko
sudo rm -f /lib/modules/${KERNEL_VERSION}/updates/drivers/net/wireless/intel/iwlwifi/mvm/iwlmvm.ko

echo "=== Step 3: Extracting WiFi kernel modules ==="
# Extract WiFi modules to the live system
sudo tar --skip-old-files -xf "$IWLWIFI_ARCHIVE" -C /lib/modules/${KERNEL_VERSION}/
if [[ $? -ne 0 ]]; then
    echo "ERROR: Could not extract WiFi kernel modules from $IWLWIFI_ARCHIVE"
    exit 1
fi
echo "WiFi kernel modules installed successfully"

echo "=== Step 4: Extracting WiFi firmware files ==="
# Extract firmware files
sudo tar -xf "$FIRMWARE_ARCHIVE" -C /lib/firmware/
if [[ $? -ne 0 ]]; then
    echo "ERROR: Could not extract WiFi firmware files from $FIRMWARE_ARCHIVE"
    exit 1
fi
echo "WiFi firmware files installed successfully"

echo "=== Step 5: Configuring module autoload ==="
# Add iwlwifi module to /etc/modules if not already present
if ! grep -q "^iwlwifi$" /etc/modules 2>/dev/null; then
    echo "Adding iwlwifi to /etc/modules for automatic loading"
    echo iwlwifi | sudo tee -a /etc/modules > /dev/null
else
    echo "iwlwifi already configured for autoload"
fi

echo "=== Step 6: Creating udev rule to disable BT modules ==="
# Create udev rule to disable BT modules on L4T 35.x
if [[ $L4T_VERSION =~ ^35_ ]]; then
    UDEV_RULE_FILE="/etc/udev/rules.d/99-wifi-bt-conflict.rules"
    echo "Creating udev rule to disable BT modules for L4T $L4T_VERSION"
    sudo bash -c "cat > $UDEV_RULE_FILE" << 'EOF'
SUBSYSTEM=="usb", ATTRS{idVendor}=="8087", ATTRS{idProduct}=="0032", ATTR{authorized}="0"
EOF
    if [[ $? -eq 0 ]]; then
        echo "BT disable udev rule created successfully"
        echo "Reloading udev rules..."
        sudo udevadm control --reload-rules
        sudo udevadm trigger
        echo "Udev rules reloaded and triggered"
    else
        echo "WARNING: Failed to create BT disable udev rule"
    fi
else
    echo "Skipping BT disable udev rule creation (not L4T 35.x)"
fi

echo "=== Step 7: Regenerating module dependencies ==="
# Regenerate module dependencies
sudo depmod -a
if [[ $? -ne 0 ]]; then
    echo "WARNING: Failed to regenerate module dependencies"
else
    echo "Module dependencies regenerated successfully"
fi

echo ""
echo "=== Step 8: WiFi Driver Correction Complete ==="
echo ""
echo "Summary:"
echo "  - L4T Version: $L4T_VERSION"
echo "  - Kernel Version: $KERNEL_VERSION"
echo "  - WiFi Modules: Applied from $IWLWIFI_ARCHIVE"
echo "  - Firmware: Applied from $FIRMWARE_ARCHIVE"
echo ""
echo "Changes have been applied to:"
echo "  - Modules: /lib/modules/$KERNEL_VERSION/"
echo "  - Firmware: /lib/firmware/"
echo ""

# Ask for reboot
echo "=== Reboot Required ==="
echo "A reboot is required to ensure all WiFi driver changes take effect."
echo ""
read -p "Reboot now? (Y/n): " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Nn]$ ]]; then
    echo "Rebooting system..."
    sync
    sudo reboot
else
    echo ""
    echo "Please reboot manually when ready: sudo reboot"
    echo ""
    echo "After reboot, verify WiFi functionality with:"
    echo "  - Check module loaded: lsmod | grep iwlmvm"
fi

echo "Done."