#!/bin/bash

confdir=/etc/driverctl.d
bus=pci
probe=1
save=1
debug=0

declare -A devclasses
devclasses=(["all"]=""
            ["storage"]="01"
            ["network"]="02"
            ["display"]="03"
            ["multimedia"]="04"
            ["memory"]="05"
            ["bridge"]="06"
            ["communication"]="07"
            ["system"]="08"
            ["input"]="09"
            ["docking"]="0a"
            ["processor"]="0b"
            ["serial"]="0c"
)

function log()
{
    echo driverctl: $* >&2
}

function debug()
{
    [ "$debug" -ne 0 ] && log $*
}

function error()
{
    log $*
    exit 1
}

function usage()
{
    echo "Usage: driverctl [-v] [--noprobe] [--nosave] set-override <device> <driver>"
    echo "       driverctl [-v] [--noprobe] [--nosave] unset-override <device>"
    echo "       driverctl [-v] [--noprobe] load-override <device>"
    echo "       driverctl [-v] list-devices"
    echo "       driverctl [-v] list-overrides"
    exit 1
}

function unbind()
{
    if [ -L /sys/$devpath/driver ]; then
        debug "unbinding previous driver" $(basename $(readlink /sys/$devpath/driver))
        echo "$dev" > /sys/$devpath/driver/unbind
    else
        debug "device $dev not bound" 
    fi
}

function probe_driver()
{
    debug "reprobing driver for $dev"
    echo "$dev" > /sys/bus/$bus/drivers_probe
}

function save_override()
{
    debug "saving driver override for $dev"
    if [ -n "$drv" ]; then
	[ -d $confdir ] || mkdir -p $confdir
        echo "$drv" > $confdir/$sddev
    else
        rm -f $confdir/$sddev
    fi
}

function list_devices()
{
    devices=()
    for d in /sys/bus/$bus/devices/*; do
        if [ -f "$d/driver_override" ]; then
            override="$(cat $d/driver_override)"
            if [ $1 -eq 1 ] && [ "$override" == "(null)" ]; then
                continue
            fi
          
            line="$(basename $d)"
            devices+=($line)

            if [ -n "$2" ]; then
                class="$(cat $d/class)"
                [ "$2" == "${class:2:2}" ] || continue
            fi
            if [ -L "$d/driver" ]; then
                line+=" $(basename $(readlink $d/driver))"
            else
                line+=" (none)"
            fi
            if [ $1 -ne 1 ] && [ "$override" != "(null)" ]; then
                line+=" [*]"
            fi

            if [ $debug -ne 0 ]; then
                line+=" ($(udevadm info -q property $d | grep ID_MODEL_FROM_DATABASE | cut -d= -f2))"
            fi
            echo $line
        fi
    done
    if [ ${#devices[@]} -eq 0 ]; then
        error "No overridable devices found. Kernel too old?"
    fi
}

function set_override()
{
    if [ ! -f /sys/$devpath/driver_override ]; then
        error "device does not support driver override:" $dev
    fi
    if [ -n "$drv" ] && [ "$drv" != "none" ]; then
        debug "setting driver override for $dev: $drv"
        if [ ! -d "/sys/module/$drv" ]; then
            debug "loading driver $drv"
            /sbin/modprobe -q "$drv" || error "no such module: $drv"
        fi
    else
        debug "unsetting driver override for $dev"
    fi
    unbind
    echo "$drv" > /sys/$devpath/driver_override

    if [ "$drv" != "none" ] && [ $probe -ne 0 ]; then 
        probe_driver
        if [ ! -L /sys/$devpath/driver ]; then
            error "failed to bind device $dev to driver $drv"
        fi
    fi
}

while (($# > 0)); do
    case ${1} in
    --noprobe)
        probe=0
        ;;
    --nosave)
        save=0
        ;;
    --debug|--verbose|-v)
        debug=1
        ;;
    -b|--bus)
        bus=${2}
        shift
        ;;
    -h|--help|-*)
        usage
        ;;
    set-override)
        [ $# -ne 3 ] && usage
        cmd=$1
        dev=$2
        drv=$3
        break
        ;;
    load-override|unset-override)
        [ $# -ne 2 ] && usage
        cmd=$1
        dev=$2
        break
        ;;
    list-devices|list-overrides)
        [ $# -gt 2 ] && usage
        if [ -n "$2" ] && [ ! ${devclasses[$2]+_} ]; then
            error "device type must be one of: ${!devclasses[@]}"
        fi
        cmd=$1
        devtype="${devclasses[${2:-all}]}"
        break
        ;;
    *)
        usage
        ;;
    esac
    shift
done

[ -n "$cmd" ] || usage

if [ -n "$dev" ]; then
        if [ -n "$DEVPATH" ]; then
            devpath="$DEVPATH"
        else
            devlink="/sys/bus/$bus/devices/$dev"
            [ -L "$devlink" ] || error "no such device: $dev"
            devpath=$(realpath $devlink | cut -c5-)
        fi
        syspath=/sys/$devpath
        sddev="$bus-$dev"
fi

case ${cmd} in
    load-override)
        if [ -s $confdir/$sddev ]; then
            drv=$(cat $confdir/$sddev)
	    set_override $dev $drv
        else
            exit 1
        fi
        ;;
    list-devices)
        list_devices 0 $devtype
        ;;
    list-overrides)
        list_devices 1 $devtype
        ;;
    set-override)
        set_override $dev $drv
        [ "$save" -ne 0 ] && save_override
        ;;       
    unset-override)
        set_override $dev ""
        [ "$save" -ne 0 ] && save_override
        ;;
esac
