#!/bin/bash

# -------------------------------------------------------------------------
# GPU-SWITCH: Discrete Graphics Power & Bus Management
# Target: Razer Blade Stealth 13 (Late 2019) | Void Linux
# -------------------------------------------------------------------------

# Configuration: Static PCI Identifier for Turing dGPU
# We use a static ID because the device disappears from the bus when disabled.
DEVICE_ID="57:00.0"
PCI_PATH="/sys/bus/pci/devices/0000:$DEVICE_ID"
REQUIRED_MODULES=("nvidia_drm" "nvidia_modeset" "nvidia_uvm" "nvidia")

if [[ $EUID -ne 0 ]]; then
   echo "[ERROR] Elevated privileges required. Please execute with sudo."
   exit 1
fi

check_status() {
    echo "--- System Graphics Status ---"
    if [ -d "$PCI_PATH" ]; then
        echo "Hardware State:  [ ONLINE ] at $DEVICE_ID"
        if lsmod | grep -q "nvidia"; then
            echo "Driver State:    [ LOADED ]"
        else
            echo "Driver State:    [ UNLOADED ]"
        fi
    else
        echo "Hardware State:  [ OFFLINE ] (Device removed from PCI bus)"
        echo "Driver State:    [ UNLOADED ]"
    fi
    echo "------------------------------"
}

disable_gpu() {
    echo "[INFO] Initializing dGPU shutdown sequence..."

    # 1. Unload Kernel Modules
    echo "[STEP 1/2] Terminating NVIDIA kernel modules..."
    for mod in "${REQUIRED_MODULES[@]}"; do
        if lsmod | grep -q "$mod"; then
            modprobe -r "$mod" 2>/dev/null
            if [ $? -ne 0 ]; then
                echo "[WARN] Module $mod is currently in use. Ensure Xorg/i3 is not using the dGPU."
            fi
        fi
    done

    # 2. Physical Bus Removal
    if [ -d "$PCI_PATH" ]; then
        echo "[STEP 2/2] Detaching device from PCI bus..."
        echo 1 > "$PCI_PATH/remove"
        echo "[SUCCESS] dGPU has been isolated and powered down."
    else
        echo "[SKIP] Device is already detached from the bus."
    fi
}

enable_gpu() {
    echo "[INFO] Initializing dGPU activation sequence..."

    # 1. Trigger PCI Rescan
    echo "[STEP 1/2] Scanning PCI bridge for hardware changes..."
    echo 1 > /sys/bus/pci/rescan
    
    # Allow firmware-level handshake
    sleep 2

    # 2. Validate and Load Drivers
    if [ -d "$PCI_PATH" ]; then
        echo "[STEP 2/2] Hardware detected at $DEVICE_ID. Loading modules..."
        for mod in "nvidia" "nvidia_uvm" "nvidia_modeset" "nvidia_drm"; do
            modprobe "$mod"
        done
        echo "[SUCCESS] dGPU is now online and available for Prime Offload."
    else
        echo "[FATAL] Hardware failed to initialize at $DEVICE_ID."
        echo "[DEBUG] Check 'dmesg | tail' for PCI link training errors."
    fi
}

case "$1" in
    on)     enable_gpu ;;
    off)    disable_gpu ;;
    status) check_status ;;
    *)      echo "Usage: gpu-switch {on|off|status}" ;;
esac