#!/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)
# Using an associative array to keep track of unprocessed kernels.
# They can be easily modified, unlike ordinary arrays.
declare -A UNPROCESSED_VERSIONS KERNELS TOPLEVEL_KERNELS
KERNEL_VERSIONS=()
OS=$(gettext_printf "AOSC OS")
OS_ID="aosc"
KERNEL_DIR=${KERNEL_DIR:-/boot}
SAVEDEFAULT=

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 on 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%%/}/"
	# Convert them into paths relative to the "actual" root, say, namespaces like
	# BtrFS subvolumes.
	GRUB_KERNEL_DIR="$(make_system_path_relative_to_its_root $GRUB_KERNEL_DIR)/"
else
	# If /boot is a separate partition, the path GRUB loads the kernel should
	# be /KERNEL_FILE instead of /boot/KERNEL_FILE.
	GRUB_KERNEL_DIR="/"
fi

# search --fs-uuid commands.
MENUENTRY_PREP="$(prepare_grub_to_access_device "$GRUB_DEVICE_BOOT")"

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

# `savedefault' for saving current selection upon boot.
if [ "$GRUB_SAVEDEFAULT" = "true" ] ; then
	SAVEDEFAULT="savedefault"
fi

# Handle microcodes.
for rd in ${GRUB_EARLY_INITRD_LINUX_STOCK[@]} \
	${GRUB_EARLY_INITRD_LINUX_CUSTOM[@]} ; do
	if [ -e "$KERNEL_DIR"/"$rd" ] ; then
		INITRD_EARLY="${INITRD_EARLY} ${GRUB_KERNEL_DIR}${rd}"
	fi
	INITRD_EARLY="${INITRD_EARLY##[ 	]}"
done

# 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_configuration() {
	case "$arch" in
		loongarch64)
			printf "16k"
			;;
		*)
			printf "4k"
			;;
	esac
}

# Localize page size comment string.
configuration_str() {
	case "$1" in
		4k)
			if [ "$(default_configuration)" = "$1" ] ; then
				printf ""
			else
				gettext_printf "4KiB Kernel Page Size"
			fi
			;;
		16k)
			if [ "$(default_configuration)" = "$1" ] ; then
				printf ""
			else
				gettext_printf "16KiB Kernel Page Size"
			fi
			;;
		64k)
			if [ "$(default_configuration)" = "$1" ] ; then
				printf ""
			else
				gettext_printf "64KiB Kernel 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"
			;;
		vanillarc)
			gettext_printf "Vanilla RC Kernel"
			;;
		main)
			: noop
			;;
		asahi)
			gettext_printf "Apple Silicon"
			;;
		*)
			gettext_printf "Kernel variant '%s'" "$1"
			;;
	esac
}

# Break localversion apart, return the processed comments.
# The LOCALVERSION comes with several parts:
#   -VENDOR-VARIANT-CONFIGURATOIN
# Where:
#   - VENDOR: Undoubtably, "aosc".
#   - VARIANT: The variant of the kernel, built for specific purposes or
#     devices, e.g. "main" for mainline kernel, "lts" for LTS kernel, and
#     "asahi" for Apple Silicon devices.
#   - CONFIGURATION: Variant specific configurations, e.g. page sizes.
# 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 configuration
	# 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]}")"
	configuration="$(configuration_str "${segs[1]}")"
	comment="${variant}${sep}${configuration}"
	if [ "x$comment" = "x$sep" ] ; then
		return
	fi
	comment="${comment##$sep}"
	comment="${comment%%$sep}"
	printf "%s" " ($comment)"
}

# Generates a menuentry for linux.
linux_entry() {
	local title="$1"
	local kernel="$2"
	local initrd="$3"
	local args="$4"
	local _id="$5"
	local os_class="$(get_os_id)"

	case "${GRUB_FB_ROTATION}" in
		inverted | 180)
			args="$args fbcon=rotate:2"
			;;
		left | 90)
			args="$args fbcon=rotate:3"
			;;
		right | 270)
			args="$args fbcon=rotate:1"
			;;
	esac

	if [ ! -z "$os_class" ];  then
		os_class="--class $(echo $os_class | grub_quote)"
	fi

	printf "${indent}menuentry '%s' $os_class --class gnu --class gnu-linux --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}\t$SAVEDEFAULT\n"
	printf "${indent}\tinsmod	gzio\n"
	printf "${indent}\techo	'%s'\n" "$(gettext_printf "Loading kernel %s ..." "$kernel")"
	printf "${indent}\tlinux	%s %s\n" "${kernel##[	 ]}" "$args"
	if [ -n "$initrd" ] ; then
		printf "${indent}\techo	'%s'\n" "$(gettext_printf "Loading initial ramdisk ...")"
		printf "${indent}\tinitrd	%s\n" "${initrd##[	 ]}"
	fi
	printf "${indent}\techo	'%s'\n" "$(gettext_printf "Starting %s ..." "$OS")"
	printf "${indent}}\n\n"
}

# Get the root= parameter depending on various conditions.
find_rootdev() {
	local ident_type=uuid
	# If UUID is disabled, use PARTUUID if PARTUUID is available.
	# Otherwise use plain paths (which is dangerous, but there's
	# nothing we can do anyways.
	if ! [ -e "$GRUB_DEVICE" ] ; then
		ident_type=error
	elif bool "$GRUB_DISABLE_LINUX_UUID" || [ "$1" = "fallback" ] ; 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_configuration flag \
		variant cfg k k1
	kernels=("$@")
	newkernels=()
	default_cfg="$(default_configuration)"
	for k in "${kernels[@]}" ; do
		k1="${k#@(kernel-|vmlinu?-)}"
		IFS='-'
		components=($k1)
		unset IFS
		variant="${components[2]}"
		cfg="${components[3]}"
		# NOTE: This code can NOT insert an element in the middle
		# of the array, because it does not know where to put it.
		# If this ever results in the default configuration entry
		# showing up later than other configurations, please do
		# not attempt to fix this.
		case "${variant,,}" in
			main)
				# This is the default configration.
				# Put it on top.
				if [ -z "$cfg" ] || [ "$cfg" = "$default_cfg" ] ; then
					newkernels=("$k" "${newkernels[@]}")
				else
					newkernels+=("$k")
				fi
				;;
			*)
				newkernels+=("$k")
				;;
		esac
	done
	OLDIFS="$IFS"
	IFS=$'\n'
	echo -n "${newkernels[*]}"
	IFS="$OLDIFS"
}

get_os_id() {
(
	if [ ! -e /etc/os-release ] ; then
		return
	fi
	source "/etc/os-release"
	echo "$ID"
)
}

scan_for_kernels() {
	local scanned
	IFS=$'\n'
	scanned=($(find $KERNEL_DIR -maxdepth 1 -type f -a \( -name 'vmlinux-*' -o -name 'vmlinuz-*' -o -name 'kernel-*' \) -printf '%P\n'))
	if [ "${#scanned[@]}" -lt 1 ] ; then
		unset IFS
		return
	fi
	for kernel in "${scanned[@]}" ; do
		# Keep track of unprocessed kernels
		KERNELS["$kernel"]="$kernel"
	done
	KERNEL_VERSIONS+=($(echo -e "${scanned[*]}" | 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 ret="${INITRD_EARLY}" initrd_found="0"
	for rd in \
		initramfs-"$1".img \
		initramfs.img-"$1" \
		initrd-"$1".img \
		initrd.img-"$1" \
		initrd-"$1".gz ; do
		if [ -e "$KERNEL_DIR/$rd" ] ; then
			if [ -z "${TOPLEVEL_KERNELS["$kernel"]}" ] ; then
				info "$(gettext_printf "Initial ramdisk image for kernel '%s': '%s'" "$1" "${GRUB_KERNEL_DIR}${rd}")"
			fi
			ret="${ret} ${GRUB_KERNEL_DIR}${rd}"
			initrd_found="1"
			break
		else
			continue
		fi
	done
	if [ -n "$ret" ] ; then
		printf "%s" "$ret"
	else
		if [ -z "${TOPLEVEL_KERNELS["$kernel"]}" ] ; then
			info "$(gettext_printf "No initial ramdisk image found for kernel '%s'" "$1")"
		fi
	fi
}

# Make all variants of the latest version appear on the top level menu.
# Everything else go into the submenu.
# NOTE: Need to show one entry for each known variants. This implementation
#       may look dirty, but it is pretty straightforward, and it may work.
#       It is too difficult to map every kernel versions, variants and
#       configuration, and sort them based on keys, using Bash.
process_kernels() {
	local kernel uname_r latest latest_kernels localpart
	IFS=$'\n'
	# Process the latest kernel(s).
	# Never asusme there's only one variant here!
	latest_ver="${KERNEL_VERSIONS[0]}"
	declare -A latest_variants
	toplevel_kernels=()
	# Find the latest kernel for each variant
	for ver in "${KERNEL_VERSIONS[@]}" ; do
		kernels=($(echo "${KERNELS[*]}" | grep "$ver"))
		for kernel in "${kernels[@]}" ; do
			for variant in main lts rc vanillarc ; do
				# NOTE: Treat -main$ and -main- separately to
				# handle partial name collisions.
				if [ -z "${latest_variants[$variant]}" ] && (
					[ "${kernel%%-$variant}" != "$kernel" ] || \
					[ "${kernel//-$variant-}" != "$kernel" ] ) ; then
					debug "Latest $variant kernel: $ver"
					latest_variants["$variant"]="$ver"
				fi
			done
		done
	done
	# For each variant, add kernels with that version to the top level
	# menu. This makes sure we see all configurations for this variant
	# in the top level menu.
	for variant in main lts rc vanillarc ; do
		if [ -z "${latest_variants[$variant]}" ] ; then
			continue
		fi
		debug "Processing latest $variant kernel..."
		version="${latest_variants[$variant]}"
		# NOTE: Treat -main$ and -main- separately to handle partial
		# name collisions.
		kernels=($(echo "${KERNELS[*]}" | grep "$version"))
		for kernel in "${kernels[@]}" ; do
			if [ "${kernel%%-$variant}" != "$kernel" ] || \
				[ "${kernel//-$variant-}" != "$kernel" ] ; then
				toplevel_kernels+=("$kernel")
			fi
		done
	done
	if [ "${#toplevel_kernels[@]}" -lt 1 ] ; then
		# Perhaps they are using other variants.
		latest_kernel=($(echo "${KERNELS[*]}" | grep "$latest_ver"))
		toplevel_kernels+=("${latest_kernel[@]}")
	fi
	toplevel_kernels=($(prioritize "${toplevel_kernels[@]}"))
	unset IFS
	ROOT_PARAM="root=$(find_rootdev)"
	ROOT_PARAM_FALLBACK="root=$(find_rootdev fallback)"
	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 "${toplevel_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}"
		initrd="$(get_initrd "$uname_r")"
		if [ -z ${initrd/${INITRD_EARLY}/} ] ; then
			args="${ARGS/$ROOT_PARAM/$ROOT_PARAM_FALLBACK}"
		else
			args="$ARGS"
		fi
		linux_entry "$title" "${GRUB_KERNEL_DIR}${kernel}" "$initrd" "$args" "${OS_ID}-${kernel}"
		TOPLEVEL_KERNELS["$kernel"]="$kernel"
	done
	# Now we processed the latest versions of the kernel, time to put
	# other kernel versions into the submenu.
	if [ "${#KERNELS[@]}" = "0" ] ; then
		return
	fi
	submenu_title="$(gettext_printf "All kernel options for %s" "$OS")"
	# Sort other versions. However we can't prioritize them, it will
	# mess up the order.
	printf "submenu '%s >>' {\n" "$submenu_title"
	indent="$indent$tab"
	for version in "${KERNEL_VERSIONS[@]}" ; do
		IFS=$'\n'
		# Entries in submenus are sorted by versions.
		cur_version_kernels=($(echo "${KERNELS[*]}" | grep "$version"))
		unset IFS
		for kernel in "${cur_version_kernels[@]}" ; do
			uname_r=${kernel/#@(vmlinu?|kernel)-/}
			if [ -z "${TOPLEVEL_KERNELS["$kernel"]}" ] ; then
				info $(gettext_printf "Processing kernel '%s'\n" "$uname_r")
			fi
			localpart="${uname_r//$latest/}"
			comments="$(gettext_printf "With kernel %s" "$uname_r")"
			comments=" ($(echo "$comments" | grub_quote))"
			title="${OS}${comments}"
			initrd="$(get_initrd "$uname_r")"
			if [ -z ${initrd/${INITRD_EARLY}/} ] ; then
				args="${ARGS/$ROOT_PARAM/$ROOT_PARAM_FALLBACK}"
			else
				args="$ARGS"
			fi
			linux_entry "$title" "${GRUB_KERNEL_DIR}${kernel}" "$initrd" "$args" "${OS_ID}-${kernel}"
		done
	done
	printf "}\n"
}

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