#!/bin/bash

[ -e /etc/default/grub ] && source /etc/default/grub
[ -e /etc/default/grub-aosc ] && source /etc/default/grub-aosc
. /usr/share/grub/grub-mkconfig_lib

indent=""
tab=$'\t'
arch=$(uname -m)
KERNELS=()
# Using an associative array to keep track of unprocessed kernels.
# They can be easily modified, unlike ordinary arrays.
declare -A UNPROCESSED_VERSIONS
KERNEL_VERSIONS=()
OS=$(gettext_printf "AOSC OS")
OS_ID="aosc"
KERNEL_DIR=${KERNEL_DIR:-/boot}

shopt -s extglob

bool() {
	[[ "$1" == "true" || "$1" == "1" || "$1" =~ [yY] ]]
}

info() {
	echo "-- $@" >&2
}

if bool "$AOSC_GRUB_DEBUG" ; then
	#set -x
	debug() {
		echo -e ">> $@" >&2
	}
else
	debug() {
		: noop
	}
fi

if [ -z "$GRUB_DEVICE" ] ; then
	GRUB_DEVICE="`${grub_probe} --target=device /`"
	GRUB_DEVICE_UUID="`${grub_probe} --device ${GRUB_DEVICE} --target=fs_uuid 2> /dev/null`" || true
	GRUB_DEVICE_PARTUUID="`${grub_probe} --device ${GRUB_DEVICE} --target=partuuid 2> /dev/null`" || true
fi

if [ -z "$GRUB_DEVICE_BOOT" ] ; then
	# Device containing our /boot partition.  Usually the same as GRUB_DEVICE.
	GRUB_DEVICE_BOOT="`${grub_probe} --target=device /boot`"
	GRUB_DEVICE_BOOT_UUID="`${grub_probe} --device ${GRUB_DEVICE_BOOT} --target=fs_uuid 2> /dev/null`" || true
fi

if [[ -z "$GRUB_DEVICE" || -z "$GRUB_DEVICE_BOOT" ]] ; then
	gettext_printf "Internal error - Failed to gather information of the root partition\n"
	exit 1
fi

# From GRUB: specify subvolume ID for btrfs, if the filesystem has subvolumes.
if [ "$GRUB_FS" = "btrfs" ] ; then
	rootsubvol="`make_system_path_relative_to_its_root /`"
	rootsubvol="${rootsubvol#/}"
	if [ "x${rootsubvol}" != x ]; then
		GRUB_CMDLINE_LINUX="rootflags=subvol=${rootsubvol} ${GRUB_CMDLINE_LINUX}"
	fi
fi
# Depends on whether /boot is a mounted filesystem, the path GRUB loads the
# kernel changes.
if [ "$GRUB_DEVICE" == "$GRUB_DEVICE_BOOT" ] ; then
	GRUB_KERNEL_DIR="${KERNEL_DIR%%/}/"
else
	GRUB_KERNEL_DIR="/"
fi

MENUENTRY_PREP="$(prepare_grub_to_access_device "$GRUB_DEVICE_BOOT")"

if [ -z "$GRUB_GFXPAYLOAD_LINUX" ] ; then
	_GFXPAYLOAD_LINUX="set gfxpayload=keep"
else
	_GFXPAYLOAD_LINUX="set gfxpayload=$GRUB_GFXPAYLOAD_LINUX"
fi

# Returns the default page size suffix.
# WARN $arch is the output of `uname -m`, which is not exactly the same as
# `dpkg --print-architecture`.
default_pagesize() {
	case "$arch" in
		loongarch64)
			printf "16k"
			;;
		*)
			printf "4k"
			;;
	esac
}

pagesize_str() {
	case "$1" in
		4k)
			if [ "$(default_pagesize)" = "$1" ] ; then
				printf ""
			else
				gettext_printf "4K Page Size"
			fi
			;;
		16k)
			if [ "$(default_pagesize)" = "$1" ] ; then
				printf ""
			else
				gettext_printf "16K Page Size"
			fi
			;;
		64k)
			if [ "$(default_pagesize)" = "$1" ] ; then
				printf ""
			else
				gettext_printf "64K Page Size"
			fi
			;;
		*)
			# Unknown page size.
			echo ""
			;;
	esac
}

# Localize the given variant.
variant_str() {
	case "$1" in
		rc)
			gettext_printf "RC Kernel"
			;;
		lts)
			gettext_printf "LTS Kernel"
			;;
		main)
			: noop
			;;
		*)
			printf "Kernel variant '%s'" "$1"
			;;
	esac
}

# Break localversion apart, return the processed comments.
# E.g.
# - "-aosc-main" -> ""
# - "-aosc-lts" -> " (LTS Kernel)"
# - "-aosc-main-4k" -> " (4K Page Size)"
# - "-aosc-rc-4k" -> " (RC Kernel, 4K Page Size)"
get_comments() {
	local _str="$1" _fullver="$1" segs=() comments=() sep=", " variant pagesize
	# Not an AOSC OS kernel
	if ! [[ "$_str" =~ \.[0-9]+-aosc- ]] ; then
		return
	fi

	_str="${_str##*-aosc-}"
	IFS=-
	segs=($_str)
	unset IFS
	variant="$(variant_str "${segs[0]}")"
	pagesize="$(pagesize_str "${segs[1]}")"
	comment="${variant}${sep}${pagesize}"
	if [ "x$comment" = "x$sep" ] ; then
		return
	fi
	comment="${comment##$sep}"
	comment="${comment%%$sep}"
	printf "%s" " ($comment)"
}

linux_entry() {
	title="$1"
	kernel="$2"
	initrd="$3"
	args="$4"
	_id="$5"
	
	printf "${indent}menuentry '%s' --class linux \$menuentry_id_option '%s' {\n" "$title" "$_id"
	printf "$(echo "$MENUENTRY_PREP" | sed "s|^|\t${indent}|g")\n"
	printf "${indent}\t$_GFXPAYLOAD_LINUX\n"
	printf "${indent}\tinsmod	gzio\n"
	printf "${indent}\techo	'%s'\n" "$(gettext_printf "Loading kernel $kernel ...")"
	printf "${indent}\tlinux	%s %s\n" "$kernel" "$args"
	if [ -n "$initrd" ] ; then
		printf "${indent}\techo	'%s'\n" "$(gettext_printf "Loading initrd image ...")"
		printf "${indent}\tinitrd	%s\n" "$initrd"
	fi
	printf "${indent}\techo	'%s'\n" "$(gettext_printf "Booting up $OS ...")"
	printf "${indent}}\n\n"
}

find_rootdev() {
	local ident_type=uuid
	if ! [ -e "$GRUB_DEVICE" ] ; then
		ident_type=error
	elif bool "$GRUB_DISABLE_LINUX_UUID" ; then
		if bool "$GRUB_DISABLE_LINUX_PARTUUID" ; then
			ident_type=path
		else
			ident_type=partuuid
		fi
	else
		if [ ! -e "/dev/disk/by-uuid/${GRUB_DEVICE_UUID,,}" ] ; then
			ident_type=partuuid
		fi
		if [ ! -e "/dev/disk/by-partuuid/${GRUB_DEVICE_PARTUUID,,}" ] || \
			bool "${GRUB_DISABLE_LINUX_PARTUUID}" || \
			uses_abstraction "${GRUB_DEVICE}" lvm ; then
			ident_type=path
		fi

	fi
	case "$ident_type" in
		path)
			printf "%s" "$GRUB_DEVICE"
			;;
		uuid)
			printf "UUID=%s" "$GRUB_DEVICE_UUID"
			return
			;;
		partuuid)
			printf "PARTUUID=%s" "$GRUB_DEVICE_PARTUUID"
			;;
		*)
			# Propogate the error
			;;
	esac
}

# Prioritize kernel names in "$@", based on the default size.
# The kernel with the default size will be always on top. If it does
# not exist, then the kernel without page size suffix will be on top.
prioritize() {
	local kernels newkernels default_pgsize flag
	kernels=("$@")
	newkernels=()
	default_pgsize="$(default_pagesize)"
	for k in "${kernels[@]}" ; do
		# This system probably don't have other page sizes.
		if [ "${k%%-[0-9]*k}" = "$k" ] ; then
			newkernels=("$k" "${newkernels[@]}")
			continue
		fi
		if [ "${k%%-$default_pgsize}" != "$k" ] ; then
			newkernels=("$k" "${newkernels[@]}")
			continue
		fi
		newkernels+=("$k")
	done
	OLDIFS="$IFS"
	IFS=$'\n'
	echo -n "${newkernels[*]}"
	IFS="$OLDIFS"
}

scan_for_kernels() {
	IFS=$'\n'
	KERNELS=($(find $KERNEL_DIR -maxdepth 1 -type f -a \( -name 'vmlinux-*' -o -name 'vmlinuz-*' -o -name 'kernel-*' \) -printf '%P\n'))
	if [ "${#KERNELS[@]}" -lt 1 ] ; then
		unset IFS
		return
	fi
	KERNEL_VERSIONS+=($(echo -e "${KERNELS[*]}" | grep -o '[0-9]*\.[0-9]*\.[0-9]*' | sort -rVu))
	for ver in "${KERNEL_VERSIONS[@]}" ; do
		UNPROCESSED_VERSIONS["$ver"]="$ver"
	done
	debug "Kernel versions:\n${KERNEL_VERSIONS[*]}"
	unset IFS k
}

# Scan for a corresponding initrd image for the kernel.
# $1: Full version of the kernel.
get_initrd() {
	local rd
	# Stub
	rd="initramfs-$1.img"
	if [ -e "$KERNEL_DIR/$rd" ] ; then
		echo "$GRUB_KERNEL_DIR/$rd"
	else
		return
	fi
}

process_kernels() {
	local kernel uname_r latest latest_kernels localpart
	# Use the latest non-RC kernel.
	# Whatever. I don't care.
	IFS=$'\n'
	latest="${KERNEL_VERSIONS[0]}"
	latest_kernels=($(echo "${KERNELS[*]}" | grep "$latest"))
	latest_kernels=($(prioritize ${latest_kernels[@]}))
	unset IFS
	ROOT_PARAM="root=$(find_rootdev)"
	if [ -z "${ROOT_PARAM##root=}" ] ; then
		gettext_printf "Internal error - Failed to find the root device\n" >&2
		exit 1
	fi
	ARGS="${ROOT_PARAM} ${GRUB_CMDLINE_LINUX} ${GRUB_CMDLINE_LINUX_DEFAULT}"
	for kernel in "${latest_kernels[@]}" ; do
		uname_r=${kernel/#@(vmlinu?|kernel)-/}
		info $(gettext_printf "Processing kernel '%s'\n" "$uname_r")
		localpart="${uname_r//$latest/}"
		comments="$(get_comments "$uname_r")"
		title="${OS}${comments}"
		linux_entry "$title" "$GRUB_KERNEL_DIR$kernel" "$(get_initrd "$kernel")" "$ARGS" "${OS_ID}-${kernel}"
	done
	# Now we processed the latest versions of the kernel, time to put
	# other kernel versions into the submenu.
	unset UNPROCESSED_VERSIONS["$latest"]
	if [ "${#UNPROCESSED_VERSIONS[@]}" = "0" ] ; then
		return
	fi
	submenu_title="$(gettext_printf "More kernel options for %s" "$OS")"
	printf "submenu '%s' {\n" "$submenu_title"
	indent="$indent$tab"
	for version in "${UNPROCESSED_VERSIONS[@]}" ; do
		IFS=$'\n'
		kernels=($(echo "${KERNELS[*]}" | grep "$version"))
		unset IFS
		for kernel in "${kernels[@]}" ; do
			uname_r=${kernel/#@(vmlinu?|kernel)-/}
			info $(gettext_printf "Processing kernel '%s'\n" "$uname_r")
			localpart="${uname_r//$latest/}"
			comments="$(gettext_printf "With kernel %s" "$uname_r")"
			comments=" ($(echo "$comments" | grub_quote))"
			title="${OS}${comments}"
			linux_entry "$title" "$GRUB_KERNEL_DIR/$kernel" "$(get_initrd "$kernel")" "$ARGS" "${OS_ID}-${kernel}"
		done
	done
	printf "}\n"
}

scan_for_kernels
process_kernels
info "$(gettext_printf "Done generating entries for %s." "$OS")"