nvman

#!/bin/bash

# ================================================================================
# Colorize output.
# ================================================================================
RED=$(tput setaf 1)
GREEN=$(tput setaf 2)
YELLOW=$(tput setaf 3)
CYAN=$(tput setaf 6)
BOLD=$(tput bold)
RESET=$(tput sgr0)
# ================================================================================

# ================================================================================
# nvman variables.
# ================================================================================
version="1.3.2"
config_path="/etc/nvman"
config_file="${config_path}/config"
# ================================================================================

# ================================================================================
# Utility function, run the following command as root, pretty print when needing 
# privelege scalation.
# ================================================================================
sudo_run() {
    if [[ $(id -u) -eq 0 ]]; then
        eval "$@"
        return
    fi

    sudo -n true &> /dev/null 

    [[ $? -eq 1 ]] && 
    echo "${GREEN}==>${RESET} Privilege scalation needed"

    eval "sudo $@"
}

# ================================================================================
# Check if $1 is a valid value for service (bumblebee or optimus).
# ================================================================================
valid_service() {
    case $1 in 
        'bumblebee' | 'optimus')
            echo 1
            ;;
        *)
            echo 0
            ;;
    esac
}

valid_gpu() {
    case $1 in
        'intel' | 'nvidia' | 'auto')
            echo 1
            ;;
        *)
            echo 0
            ;;
    esac
}

# ================================================================================
# Get the service file descriptor:
#   optimus    -> optimus-manager.service
#   bumbleblee -> bumblebeed.service
# ================================================================================
service_descriptor() {
    local service=$1
    case $service in
        'bumblebee')
            service='bumblebeed.service'
            ;;
        'optimus')
            service='optimus-manager.service'
            ;;
    esac

    echo $service
}

# ================================================================================
# Prints active or inactive, depending on the service status.
# ================================================================================
service_status() {
    echo "$(systemctl status $(service_descriptor $1) | awk '/Active/ {print $2}')"
}

# ================================================================================
# Prints enabled or disabled, depending on the service status.
# ================================================================================
service_enabled() {
    local status=$(systemctl status $(service_descriptor $1) | awk '/Loaded: loaded/ { print $4}')
    echo "${status::-1}"
}

# ================================================================================
# Utility function to enable a service, check if willing to enable conflictive 
# services, if that's the case, toggle them.
# ================================================================================
enable_service() {
    local service=$1

    # Toggle service depending on input
    case $service in 
        'bumblebee')
            disable_service optimus
            ;;
        'optimus')
            disable_service bumblebee
            ;;
    esac 

    # If the service is already enabled, exit
    [[ $(service_enabled $service) = 'enabled' ]] && return 0
    sudo_run "systemctl enable $(service_descriptor $service) &> /dev/null"
}

# ================================================================================
# Utility function to disable a service, no special checks, it's okay to have both
# disabled.
# ================================================================================
disable_service() {
    local service=$1

    # If the service is already disabled, exit
    [[ $(service_enabled $service) = 'disabled' ]] && return 0
    sudo_run "systemctl disable $(service_descriptor $service) &> /dev/null"
}

# ================================================================================
# Utility function to start a service, check if willing to start conflictive 
# services, if that's the case, toggle them.
# ================================================================================
start_service() {
    local service=$1

    # If the service is already active, exit
    [[ $(service_status $service) = 'active' ]] && return 0

    # Check if willing to activate conflictive services
    case $service in 
        'bumblebee')
            [[ $(service_status optimus) = 'active' ]] && 
            print_error "Trying to activate conflicting services" && return 1
            ;;
        'optimus')
            [[ $(service_status bumblebee) = 'active' ]] && 
            print_error "Trying to activate conflicting services" && return 1
            ;;
    esac 

    # Start the service
    sudo_run "systemctl start $(service_descriptor $service)"
}

# ================================================================================
# Utility function to stop a service, no special checks, it's okay to have both
# stopped.
# ================================================================================
stop_service() {
    local service=$1

    # If the service is already inactive, exit
    [[ $(service_status $service) = 'inactive' ]] && return 0

    # Stop the service
    sudo_run "systemctl stop $(service_descriptor $service)"
}

# ================================================================================
# Print with 'Error:' 
# ================================================================================
print_error() {
    echo "${RED}${BOLD}Error:${RESET} $@"
}

# ================================================================================
# Utility function to read NVIDIA GPU status, prints on or off
# ================================================================================
get_nvidia_status() {
    local status=$(cat /proc/acpi/bbswitch | awk '{printf $2}')    
    echo "$(echo -n $status | tr '[:upper:]' '[:lower:]')"
}
# ================================================================================

# ================================================================================
# Print status utility functions
# ================================================================================
print_optimus() {
    first=$(echo -n "${BOLD}${WHITE}Optimus :${RESET} ${optimus_color}${optimus_status}${RESET} ")
    second=$(echo -n "${YELLOW}($(service_enabled optimus))${RESET}")

    echo "$first $second"
}

print_bumblebee() {
    first=$(echo -n "${BOLD}${WHITE}Bumblebee :${RESET} ${bumblebee_color}${bumblebee_status}${RESET} ")
    second=$(echo -n "${YELLOW}($(service_enabled bumblebee))${RESET}")

    echo "$first $second"
}

print_nvidia() {
    echo "${BOLD}${WHITE}NVIDIA GPU :${RESET} $(get_nvidia_status)"
}

print_default() {
    echo "${BOLD}${WHITE}Default service :${RESET} $(get_default_service)"
}

print_optimus_mode() {
    echo "${BOLD}${WHITE}Optimus mode :${RESET} $(get_optimus_mode)"
}

print_optimus_startup() {
    echo "${BOLD}${WHITE}Optimus startup :${RESET} $(get_optimus_startup)"
}

print_status() {
    local optimus_color=$RED
    local bumblebee_color=$RED

    local optimus_status=$(service_status optimus)
    local bumblebee_status=$(service_status bumblebee)

    [[ $optimus_status = 'active' ]] && optimus_color=$GREEN
    [[ $bumblebee_status = 'active' ]] && bumblebee_color=$GREEN
    
    # Print optimus overall status
    print_optimus | awk '{printf "%-20s %s %-20s %s\n", $1, $2, $3, $4 }'
    
    # Print bumblebee overall status
    print_bumblebee | awk '{printf "%-20s %s %-20s %s\n", $1, $2, $3, $4 }'
    
    # Print NVIDIA GPU status
    print_nvidia | awk '{printf "%s %-9s %s %s\n", $1, $2, $3, $4 }'

    # Print default service
    print_default | awk '{printf "%s %-8s %s %s\n", $1, $2, $3, $4 }'

    # Print optimus mode (if active)
    [[ $(service_status optimus) = 'active' ]] &&
    print_optimus_mode | awk '{printf "%s %-8s %s %s\n", $1, $2, $3, $4 }'

    # Print optimus startup (if enabled)
    [[ $(get_default_service) = 'optimus' ]] &&
    print_optimus_startup | awk '{printf "%s %-8s %s %s\n", $1, $2, $3, $4 }'
}
# ================================================================================

# ================================================================================
# Write to config_file the default service to be enabled at boot.
# ================================================================================
set_default_service() {
    if [[ ! -d $config_path ]]; then
        sudo_run "mkdir -p $config_path"
    fi

    echo "default = $1" | sudo_run "tee $config_file &> /dev/null"
}

# ================================================================================
# Read from config_file the default service to be enabled at boot.
# If there's no config_file, call set_default_service optimus.
# ================================================================================
get_default_service() {
    if [[ ! -e $config_file ]]; then
        set_default_service optimus
    fi

    echo "$(cat $config_file | grep '\<default\>' | tr '=' ' ' | awk '{print $2}')"
}

# ================================================================================
# Utility function to read optimus mode, prints intel or nvidia 
# ================================================================================
get_optimus_mode() {
    echo "$(optimus-manager --print-mode | awk '/mode/ {print $5}')"
}

# ================================================================================
# Utility function to read optimus startup, prints intel or nvidia 
# ================================================================================
get_optimus_startup() {
    echo "$(optimus-manager --print-startup | awk '/startup/ {print $7}')"
}
# ================================================================================
# Check if using gdm, suggest wiki read
# ================================================================================
check_gdm_prime() {
    local gdm_name=$(pacman -Qq gdm)

    [[ $(service_enabled 'gdm.service') = 'disabled' ]] && return

    local bad_flag=false

    if [[ $gdm_name != 'gdm-prime' ]]; then
        print_error "Seems like you're using GDM, but you don't have gdm-prime installed."
        bad_flag=true
    fi

    local xorg_gdm=$(cat /etc/gdm/custom.conf | grep '\<WaylandEnable\>' | cut -dW -f 1)
    if [[ $xorg_gdm = '#' ]]; then 
        print_error "Please disable wayland for GDM /etc/gdm/custom.conf"
        bad_flag=true
    fi

    $bad_flag && exit 1
}

# ================================================================================

# ================================================================================
# Print usage (help)
# ================================================================================
print_usage() {
    echo "Usage: $0 command"
    echo "Version $version"
    echo ""
    echo "Commands:"
    echo "  nvidia  <on|off>                Turn on/off NVIDIA GPU, useful for CUDA"
    echo "  run     <cmd>                   Run any command with primusrun"
    echo "  switch  <nvidia|intel|auto>     Switch gpu using optimus-manager"
    echo "  startup <nvidia|intel>          Set startup gpu for optimus service"
    echo "  default <bumblebee|optimus>     Set default service at boot"
    echo "  start   <bumblebee|optimus>     Manually start the service"
    echo "  stop    <bumblebee|optimus>     Manually stop the service"
    echo ""
    echo "  help                            Show this help"
    echo "  status                          Show the current status of both services"
}


# ================================================================================
# Main section
# ================================================================================

# Read arguments
cmd=$1
arg=$2
shift
shift

# Check if the script is being run as root, the only exception is if it's using the
# command 'inspector'
if [[ $cmd != 'inspector' ]] && [[ $(id -u) -eq 0 ]]; then
    print_error "Don't run this script as root"
    exit 1
fi

case $cmd in 
    'nvidia')
        # Check if optimus service is active and it's using nvidia, fatal error
        optimus_status=$(service_status optimus)
        if [[ $optimus_status = 'active' ]] && [[ $(get_optimus_mode) = 'nvidia' ]]; then
            print_error "Currently using nvidia with optimus!"
            exit 1
        fi

        case $arg in
            'on')
                sudo_run "tee /proc/acpi/bbswitch <<< ON &> /dev/null"
                ;;
            'off')
                sudo_run "rmmod nvidia_uvm &> /dev/null | true" 
                sudo_run "rmmod nvidia_modeset &> /dev/null | true"
                sudo_run "rmmod nvidia &>/dev/null | true"
                sudo_run "tee /proc/acpi/bbswitch <<< OFF &> /dev/null"
                ;;
            *)
                print_error "Invalid value $arg"
                ;;
        esac

        ;;
    'run')
        # Check if optimus service is active and it's using nvidia, fatal error
        optimus_status=$(service_status optimus)
        if [[ $optimus_status = 'active' ]] && [[ $(get_optimus_mode) = 'nvidia' ]]; then
            print_error "Currently using nvidia with optimus!"
            exit 1
        fi

        echo "${CYAN}==>${RESET} Running $arg $@"

        # Toggle services
        stop_service optimus
        start_service bumblebee

        # Run with primusrun
        eval "primusrun $arg $@"
        ;;

    'switch')
        gpu=$arg
        if [[ $(valid_gpu $gpu) = '0' ]]; then
            print_error "Invalid GPU value '$gpu'!"
            exit 1
        fi

        check_gdm_prime 

        # Toggle services
        stop_service bumblebee
        start_service optimus

        # Switch modes
        echo "${CYAN}==>${RESET} Switching to $gpu"
        sudo_run "true" # Do nothing but prompt message
        optimus-manager --switch $gpu --no-confirm &> /dev/null
        ;;

    'start')
        if [[ $(valid_service $arg) = '0' ]]; then
            print_error "Unknown service '$arg'!"
            exit 1
        fi

        echo "${CYAN}==>${RESET} Starting service $arg"
        start_service $arg 
        ;; 

    'stop')
        if [[ $(valid_service $arg) = '0' ]]; then
            print_error "Unknown service '$arg'!"
            exit 1
        fi

        echo "${CYAN}==>${RESET} Stopping service $arg"
        stop_service $arg 
        ;;

    'status' | '')
        print_status
        ;;

    'default')
        if [[ $(valid_service $arg) = '0' ]]; then
            print_error "Unknown service '$arg'!"
            exit 1
        fi

        echo "${CYAN}==>${RESET} Setting default service to $arg"
        set_default_service $arg
        enable_service $arg
        ;;

    'startup')
        if [[ $(valid_gpu $arg) = '0' ]] || [[ $arg = 'auto' ]]; then
            print_error "Invalid GPU value '$arg'!"
            exit 1
        fi

        optimus_status=$(service_status optimus)
        if [[ $optimus_status = 'inactive' ]]; then 
            stop_service bumblebee
            start_service optimus
        fi

        echo "${CYAN}==>${RESET} Setting default startup to $arg"
        sudo_run "true" # Do nothing but prompt message
        optimus-manager --set-startup $arg &> /dev/null

        if [[ $optimus_status = 'inactive' ]]; then
            stop_service optimus
            start_service bumblebee 
        fi

        echo "${CYAN}==>${RESET} Default startup for ${BOLD}${WHITE}optimus${RESET} is ${BOLD}${WHITE}$(get_optimus_startup)${RESET}"
        ;;

    'inspector')
        if [[ $(service_enabled optimus) = 'enabled' ]] && 
            [[ $(service_enabled bumblebee) = 'enabled' ]]; then
            enable_service $(get_default_service)
            exit 0
        fi

        if [[ $(service_enabled optimus) = 'disabled' ]] && 
            [[ $(service_enabled bumblebee) = 'disabled' ]]; then
            enable_service $(get_default_service)
            exit 0
        fi

        ;;

    'help')
        print_usage
        ;;
    *)
        print_error "Unknown command '$cmd'!"
        print_usage 
        exit 1
        ;;
esac
# ================================================================================