#!/bin/bash
# -*- mode: sh; indent-tabs-mode: nil; sh-basic-offset: 4; -*-
# vim: et sts=4 sw=4

#  SPDX-License-Identifier: LGPL-2.1+
#
#  Copyright © 2019-2020 Collabora Ltd.
#  Copyright © 2019-2020 Valve Corporation.
#
#  This file is part of steamos-customizations.
#
#  steamos-customizations is free software; you can redistribute it and/or
#  modify it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation; either version 2.1 of the License,
#  or (at your option) any later version.

set -eu;

export LC_ALL=C;
export LANG=C;

initrd=;
module_list=();
initrd_module_list=();
missing_module_list=();
uptodate_module_list=();
nvidia_module=;

get_initrd ()
{
    local ver="$1";
    local pkgbase="/usr/lib/modules/$(uname -r)/pkgbase";

    initrd=;
    if [ -e "$pkgbase" ];
    then
        initrd="/boot/initramfs-$(<"$pkgbase").img";
    else
        initrd="/boot/initrd.img-$ver";
    fi;
}

get_nvidia_module ()
{
    nvidia_module=;
    if [ -e /etc/arch-release ];
    then
        nvidia_module="nvidia";
    else
        nvidia_module="nvidia-current";
    fi;
}

get_module_list ()
{
    mapfile -t module_list < <(ls -1 "$1");
}

get_initrd_module_list ()
{
    mapfile -t initrd_module_list < <(lsinitrd "$1" | sed -n -E '/\.ko(\.xz)*$/s/.*\/(.*)$/\1/p');
}

get_missing_module_list ()
{
    local initrd="$1"; shift;

    get_initrd_module_list "$initrd";
    mapfile -t missing_module_list < <( \
        comm -13 <(printf "%s\n" "${initrd_module_list[@]}" | sort;) \
                 <(printf "%s\n" "$@"                       | sort;) \
    );
}

get_uptodate_module_list ()
{
    local initrd="$1"; shift;
    local modpath="$1"; shift;

    mapfile -t uptodate_module_list < <( \
        local mod;
        for mod in "$@";
        do
            if [ "$modpath/$mod" -nt "$initrd" ];
            then
                echo "$mod";
            fi;
        done;
    );
}

initrd_is_up_to_date ()
{
    local initrd="$1";  shift;
    local modpath="$1"; shift;

    get_missing_module_list "$initrd" "$@";
    if [ "${#missing_module_list[@]}" -ne 0 ];
    then
        steamos-info "Module list in $initrd incomplete:" "${missing_module_list[@]}";
        return 1;
    fi;

    get_uptodate_module_list "$initrd" "$modpath" "$@";
    if [ "${#uptodate_module_list[@]}" -ne 0 ];
    then
        steamos-info "Module list in $initrd is outdated:" "${uptodate_module_list[@]}";
        return 1;
    fi;
}

get_status ()
{
    local kver=;
    local arch=;
    local words=();
    local status=();

    kver="$(uname -r)";
    arch="$(uname -m)";
    while IFS="$IFS:," read -r -a words;
    do
        # module version status=added
        if [ "${#words[@]}" -eq 3 ] && [ "${words[2]}" = "added" ];
        then
            status=("${words[@]}");
        # module version kernel-version arch status=(built|installed)
        elif [ "${words[2]}" = "$kver" ] && [ "${words[3]}" = "$arch" ];
        then
            status=("${words[@]}");
        fi;
    done < <(dkms status "$nvidia_module");

    if [ "${#status[@]}" -eq 0 ];
    then
        return 1
    fi;

    echo "${status[@]:-}"
}

check_nvidia_dkms_modules ()
{
    local mod=;
    local ver=;
    local krn=;
    local arc=;
    local sta=;
    local modpath=;

    # get the module status
    get_nvidia_module;
    if ! read -r mod ver krn arc sta _ < <(get_status "$nvidia_module");
    then
        # the module is not even added, quit!
        steamos-info "The driver $nvidia_module is not added!";
        return 1;
    fi;
    if [ -z "$sta" ]; then sta="$krn"; krn=; fi;

    # the module is added or built but it is not installed though, install it!
    if [ "$sta" != "installed" ];
    then
        steamos-notice "Installing the driver $nvidia_module...";
        if ! dkms install "$mod/$ver" ||
           ! read -r mod ver krn arc sta _ < <(get_status "$nvidia_module");
        then
            # the module is not installed, quit!
            steamos-error "The driver $nvidia_module cannot be installed!";
            sleep 5s;
            return 1;
        fi;
        if [ -z "$sta" ]; then sta="$krn"; krn=; fi;
    fi;

    # blacklist nouveau
    cat <<EOF > /etc/modprobe.d/nvidia-dkms.conf;
#
# Automatically generated file; DO NOT EDIT.
# ${0##*/} modprobe.conf
#
blacklist nouveau
EOF

    # the modules lives there
    modpath="/var/lib/dkms/$mod/$ver/$krn/$arc/module";
    get_module_list "$modpath";

    # regenerate the dracut configuration
    steamos-info "Adding modules ${module_list[*]} to dracut";
    cat <<EOF > /etc/dracut.conf.d/99-nvidia-dkms.conf;
#
# Automatically generated file; DO NOT EDIT.
# ${0##*/} dracut.conf
#
add_drivers+=" ${module_list[*]%.ko*} "
install_items+=" /etc/modprobe.d/nvidia-dkms.conf "
EOF

    # the initrd is outdated, rebuild it!
    get_initrd "$ver";
    if ! initrd_is_up_to_date "$initrd" "$modpath" "${module_list[@]}";
    then
        steamos-notice "Rebuilding the Initial RAM-disk...";
        update-dracut;
    fi;
}

for mod in "$@";
do
    if [[ "$mod" =~ ^nvidia ]];
    then
        if ! check_nvidia_dkms_modules;
        then
            # remove the dracut configuration
            rm -f /etc/dracut.conf.d/99-nvidia-dkms.conf;

            # do not blacklist nouveau
            rm -f /etc/modprobe.d/nvidia-dkms.conf;
        fi;

        # the module nvidia is not loaded...
        if ! grep -q '^nvidia' /proc/modules;
        then
            # the system has rebooted already, load nouveau instead!
            if [[ -e /var/lib/.steamos-glx-driver-has-rebooted ]];
            then
                rm -f /var/lib/.steamos-glx-driver-has-rebooted;
                steamos-warning "The system falls back to the driver nouveau!";
                if ! grep -q "^nouveau" /proc/modules
                then
                    modprobe nouveau;
                fi
            # the system has not rebooted yet, reboot!
            else
                touch /var/lib/.steamos-glx-driver-has-rebooted;
                steamos-notice "Rebooting in 5 seconds...";
                reboot;
            fi;
            sleep 5s;
        else
            rm -f /var/lib/.steamos-glx-driver-has-rebooted;
        fi;

        # clear the splash screen
        steamos-notice;
    fi;
done;
