#!/bin/bash

PARAM=$#
FBARG=""
FBRESOLUTION=""
GRUBDEFAULT="/etc/default/grub"
UPDATEGRUBCMD="/usr/sbin/update-grub"
MKINITCPIOCONF="/etc/mkinitcpio.conf"



check_system() {
    local missing_file=""

    [ ! -e "${GRUBDEFAULT}" ] && missing_file="${GRUBDEFAULT}"
    [ ! -e "${UPDATEGRUBCMD}" ] && missing_file="${UPDATEGRUBCMD}"
    [ ! -e "${MKINITCPIOCONF}" ] && missing_file="${MKINITCPIOCONF}"

    if [ "${missing_file}" != "" ]; then
        echo "warning: '${missing_file}' does not exists! Abording..."
        echo ""
        echo "note: mhwd-fb couldn't find grub installed on your system. mhwd-fb sets kernel arguments to tell the mhwd-fb initramfs hook which framebuffer should be loaded. Check the wiki how to add this kernel arguments manually."
        exit 1
    fi

    # Check if mhwd-fb hook is enabled
    . "${MKINITCPIOCONF}"
    local found="";

    for arg in ${HOOKS[@]}
    do
        if [ "${arg}" == "mhwd-fb" ]; then
            found="true"
            break
        fi 
    done

    if [ "${found}" != "true" ]; then
        echo "error: mhwd-fb hook isn't enabled! Abording..."
        exit 1
    fi
}



get_framebuffer_resolution() {
    # Find fitting resolution if not set
    [ "${FBRESOLUTION}" == "" ] && FBRESOLUTION="$(LANG=C LC_MESSAGES=C xrandr | grep \* | cut -d' ' -f4 | head -n 1)-32"
    [ "${FBRESOLUTION}" == "" ] && FBRESOLUTION="$(LANG=C LC_MESSAGES=C hwinfo --framebuffer | grep -w 'Mode' | cut -d: -f2 | sed 's/^ *//g' | sed 's/ .*//g' | sort -nr | head -n 1)-32"

    # Check if resolution is valid
    if ! [[ "$(echo "${FBRESOLUTION}" | sed 's/-//' | sed 's/x//')" =~ ^[0-9]+$ ]] || [[ "${FBRESOLUTION}" != *[0-9]"x"[0-9]* ]]; then
        echo "error: '${FBRESOLUTION}' isn't a valid framebuffer resolution!"
        exit 1
    fi
}



# param 1: "kms", "uvesafb" or "none"
set_framebuffer() {
    echo ":: setup mhwd-fb"

    . "${GRUBDEFAULT}"
    local NEW_GRUB_CMDLINE_LINUX=""

    for arg in ${GRUB_CMDLINE_LINUX}
    do
        if [[ "${arg}" != "mhwdfb="* ]]; then
            NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}${arg} "
        fi
    done

    if [ "$1" == "kms" ]; then
        NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}mhwdfb=kms"
        echo "  using kms framebuffer"

        # Enable GRUB_GFXPAYLOAD_LINUX
        sed -i '/#[]*GRUB_GFXPAYLOAD_LINUX[ ]*=.*/ s/^.//' "${GRUBDEFAULT}"
    elif [ "$1" == "uvesafb" ]; then
        get_framebuffer_resolution

        if [ ${FBRESOLUTION} == "" ]; then
            NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}mhwdfb=uvesafb"
            echo "  using uvesafb framebuffer"
        else
            NEW_GRUB_CMDLINE_LINUX="${NEW_GRUB_CMDLINE_LINUX}mhwdfb=uvesafb:${FBRESOLUTION}"
            echo "  using uvesafb framebuffer with resolution: ${FBRESOLUTION}"
        fi

        # Disable GRUB_GFXPAYLOAD_LINUX
        sed -i "/#/!s|GRUB_GFXPAYLOAD_LINUX[ ]*=.*|#&|g" "${GRUBDEFAULT}"
    elif [ "$1" == "none" ]; then
        echo "  disabling mhwd-fb framebuffer option"

        # Enable GRUB_GFXPAYLOAD_LINUX
        sed -i '/#[]*GRUB_GFXPAYLOAD_LINUX[ ]*=.*/ s/^.//' "${GRUBDEFAULT}"
    else
        echo "error: invalid function argument!"
        exit 1
    fi
    
    # Adjust grub default file
    sed -i "/#/!s|GRUB_CMDLINE_LINUX[ ]*=.*|GRUB_CMDLINE_LINUX=\"${NEW_GRUB_CMDLINE_LINUX}\"|g" "${GRUBDEFAULT}"

    # Update grub configuration file
    echo ""
    echo ":: update-grub output:"
    ${UPDATEGRUBCMD}
    if [ $? != 0 ]; then
        echo "error: grub-update failed!"
        exit 1
    fi
}



print_status() {
    . "${GRUBDEFAULT}"
    local FB_DRIVER=""
    local FB_RESOLUTION=""

    for arg in ${GRUB_CMDLINE_LINUX}
    do
        if [[ "${arg}" == "mhwdfb="* ]]; then
            arg="$(echo "${arg}" | cut -d= -f2)"
            FB_DRIVER="$(echo "${arg}" | cut -d: -f1)"
            FB_RESOLUTION="$(echo "${arg}:" | cut -d: -f2)" # additional ':' if f2 isn't set -> empty string
            break;
        fi
    done

    if [ "${FB_DRIVER}" != "kms" ] && [ "${FB_DRIVER}" != "uvesafb" ] && [ "${FB_DRIVER}" != "" ]; then
        echo "warning: set mhwd-fb driver is invalid: '${FB_DRIVER}'"
        exit 1
    elif [ "${FB_DRIVER}" == "" ]; then
        FB_DRIVER="disabled"
        FB_RESOLUTION="-"
    else
        [ "${FB_RESOLUTION}" == "" ] && FB_RESOLUTION="auto"
    fi

    echo ":: current mhwd-fb state"
    echo -e "  framebuffer:\t${FB_DRIVER}"
    echo -e "  resolution:\t${FB_RESOLUTION}"
}



print_help() {
    echo "mhwd-fb [OPTION] [...]"
    echo ""
    echo "   --help                          show help"
    echo "   --status                        show current status"
    echo "   --setfb [kms/uvesafb/none]      set framebuffer"
    echo "   --setres [resolution]           set uvesafb resolution. ex.: 1440x900 or 1440x900-32"
    echo ""
    echo "note: mhwd-fb sets kernel arguments to tell the mhwd-fb initramfs hook which framebuffer should be loaded. mhwd-fb requires an installed grub which is configured to this system to work. Otherwise check the wiki how to add mhwd-fb kernel arguments manually."
    echo ""
}

if [ "${PARAM}" -lt 1 ]; then
    print_status
    exit 0
fi

for (( I=1; $I <= $PARAM; I++ ))
do
    case "$1" in
        --help)
            print_help
            exit 0
            ;;
        --setfb)
            shift
            FBARG="$1"
            ;;
        --setres)
            shift
            FBRESOLUTION="$1"
            ;;
        --status)
            print_status
            exit 0
            ;;
        "") ;;
        *)
            echo "error: invalid argument: $1"
            echo ""
            print_help
            exit 1
            ;;
    esac

    shift
done


# Check root
if [[ $EUID -ne 0 ]]; then
   echo "error: you cannot perform this operation unless you are root." 1>&2
   exit 1
fi

# Check system
check_system

# Set framebuffer
if [ "${FBARG}" != "" ]; then
    if [ "${FBARG}" != "kms" ] && [ "${FBARG}" != "uvesafb" ] && [ "${FBARG}" != "none" ]; then
        echo "error: invalid argument '${FBARG}'"
        exit 1
    fi

    set_framebuffer "${FBARG}"
fi

exit 0
