#!/usr/bin/env bash
# Copyright (C) 2024-2025 Vasiliy Stelmachenok for CachyOS
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

set -e

die() {
    echo "$@" >&2
    exit 1
}

if [ "$EUID" -gt 0 ]; then
    die "This script can be run only from root"
fi

grub_setup_params() {
    if [ ! -e /boot/grub/grub.cfg ]; then
        die "GRUB doesn't seem to be installed? Please set crashkernel=128M manually."
    fi

    echo "Setting up kernel parameters for GRUB..."

    grub_cmdline="GRUB_CMDLINE_LINUX_DEFAULT"
    echo "$grub_cmdline=\"\$$grub_cmdline crashkernel=256M\"" >/etc/grub.d/00-kdump
    chmod +x /etc/grub.d/00-kdump

    if grub-mkconfig -o /boot/grub/grub.cfg 1>/dev/null; then
        die "Failed to update grub config"
    fi
}

rebind_setup_params() {
    if [ ! -e /boot/refind_linux.conf ]; then
        die "rEFInd doesn't seem to be installed? Please set crashkernel=256M manually."
    fi

    echo "rEFInd bootloader detected"

    # Ugly way to do it, but should work...
    sed -i "0,/\"$/{s// crashkernel=128M\"/}" /boot/refind_linux.conf
}

systemd_boot_setup_params() {
    if ! command -v sdboot-manage 2>/dev/null; then
        die "systemd-boot-manager is not installed. Please set crashkernel=256M manually."
    fi

    echo "systemd-boot bootloader detected"
    echo 'LINUX_OPTIONS+=" crashkernel=256M"' > /etc/sdboot-manage.conf.d/10-kdump.conf

    sdboot-manage update
}

limine_setup_params() {
    if ! command -v limine-update 2>/dev/null; then
        die "limine-entry-tool is not installed. Please set crashkernel=256M manually."
    fi

    echo "Limie bootloader detected"
    echo 'KERNEL_CMDLINE[default]+="crashkernel=256M"' > /etc/limine-entry-tool.d/10-kdump.conf
    limine-update
}

try_to_detect_bootloader() {
    bootloader_type="$(($(</proc/sys/kernel/bootloader_type) >> 4))"

    case "$bootloader_type" in
        $((0x07))) # GRUB
            grub_setup_params
            ;;
        $((0x02))) # EFIstub based bootloader
            if bootctl is-installed &>/dev/null; then
                systemd_boot_setup_params
            else
                # Assume that refind is used if systemd-boot is not installed
                rebind_setup_params
            fi
            ;;
        $((0x0F))) # Limine bootloader
            limine_setup_params
            ;;
        *)
            die "Unknown type of bootloader. Please set crashkernel=256M manually."
            ;;
    esac

    echo "The parameter setup is complete, please restart the system."
}

cleanup_files() {
    if [ -e /etc/grub.d/00-kdump ]; then
        rm /etc/grub.d/00-kdump && echo "Config for GRUB has been removed"
    fi

    if [ -e /etc/sdboot-manage.conf.d/10-kdump.conf ]; then
        rm /etc/sdboot-manage.conf.d/10-kdump.conf && echo "Config for systemd-boot has been removed"
    fi

    if [ -e /boot/refind_linux.conf ]; then
        sed -i "s/crashkernel=256M//" /boot/refind_linux.conf && echo "crashkernel paramater in rEFInd config has been removed"
    fi

    if [ -e /etc/limine-entry-tool.d/10-kdump.conf ]; then
        rm /etc/limine-entry-tool.d/10-kdump.conf && echo "Config for Limine has been removed"
    fi
}

load_kexec() {
    local crashkernel="$(cat /sys/kernel/kexec_crash_size)"

    if [ ! "$crashkernel" -gt 0 ]; then
        echo "The crashkernel parameter is not set." >&2
        echo "Please run kdump setup or add crashkernel=128M parameter manually." >&2
        exit 1
    fi

    local kernelname="$(cat /lib/modules/"$(uname -r)"/pkgbase)"
    local vmlinux="/boot/vmlinuz-${kernelname}"
    local initramfs="/boot/initramfs-${kernelname}.img"

    if command -v limine-mkinitcpio &>/dev/null; then
        local machine_id="$(cat /etc/machine-id)"
        vmlinux="/boot/${machine_id}/${kernelname}/vmlinuz-${kernelname}"
        initramfs="/boot/${machine_id}/${kernelname}/initramfs-${kernelname}"
    fi

    if [ ! -e "$vmlinux" ]; then
        die "vmlinuz for current kernel not found. Do you have /boot unmounted?"
    fi

    if [ ! -e "$initramfs" ]; then
        die "initramfs for current kernel not found. Do you have /boot unmounted?"
    fi

    if [ ! -e /var/crash ]; then
        mkdir -p /var/crash
    fi

    # We assume that if the current kernel has reached the stage of booting
    # systemd and the kdump service, it should be suitable for creating dumps
    # inside initramfs and will not cause panic at boot time.
    local cmdline="$(cat /proc/cmdline)"
    cmdline="${cmdline//quiet/}"
    cmdline="${cmdline//crashkernel=256M/}"
    if kexec -p "$vmlinux" --initrd="$initramfs" \
        --command-line="$cmdline loglevel=5 fsck.mode=force fsck.repair=yes nr_cpus=1 irqpoll reset_devices systemd.unit=emergency.target nomodeset"
    then
        echo "The fallback kernel has been successfully loaded"
    else
        die "Failed to load a fallback kernel using kexec"
    fi
}

unload_kexec() {
    current_state="$(cat /sys/kernel/kexec_crash_loaded)"
    if [ "$current_state" -gt 0 ]; then
        kexec -p -u
        echo "The fallback kernel was successfully unloaded from memory."
    else
        echo "The fallback kernel was not loaded, nothing to do"
    fi
}

collect() {
    if [ ! -e /proc/vmcore ]; then
        die "Failed to collect kernel dump, /proc/vmcore is not exist"
    fi

    umask 0077
    makedumpfile --dump-dmesg /proc/vmcore "/var/crash/crashdump-$(date +'%F-%T').log"
    makedumpfile -l -d 31 /proc/vmcore "/var/crash/crashdump-$(date +'%F-%T')"
}

usage() {
    cat <<EOF
${0##*/} <ACTION>
Utility for quick kdump configuration

Options:
    setup     Attempt to automatically set the crashkernel parameter
    load      Load a fallback kernel for panic via kexec
    unload    Unload fallback kernel from reserved area in memory
    clean     Clear config files created by setup
EOF
}

case "$1" in
    setup)
        try_to_detect_bootloader
        ;;
    load)
        load_kexec
        ;;
    unload)
        unload_kexec
        ;;
    collect)
        collect
        ;;
    cleanup)
        cleanup_files
        ;;
    *)
        usage
        ;;
esac

exit 0
