#!/bin/bash
set -e
[ -n "$DEBUG" ] && set -x

. /lib/gentoo/functions.sh

EFI_PARTITION=/boot/efi
EFI_DIR=/efi/gentoo

_() {
	if [ "$NOOP" ]; then
		echo "+" "$*"
	else
		echo "+" "$*"
		"$@"
	fi
}

mount_efi() {
	mount "${EFI_PARTITION}" || true
	mount -o remount,rw /sys/firmware/efi/efivars || true
	trap '{ umount ${EFI_PARTITION}; mount -o remount,ro /sys/firmware/efi/efivars; }' EXIT
}
path_from_efi() {
	sed -e 's,\\,/,g' -e "s,^,${EFI_PARTITION},"
}
path_to_efi() {
	sed -e "s,${EFI_PARTITION},," -e 's,/,\\,g'
}
hashsum() {
	sha1sum | cut -d' ' -f1
}

parse_efi() (
	local bootorder bootcurrent
	bootorder=$(efibootmgr -u | grep '^BootOrder' | sed -e 's,BootOrder: ,,' -e 's/,/\n/g' | awk '{print $0 "\t" NR}')
	bootcurrent=$(efibootmgr -u | grep '^BootCurrent' | sed -e 's,BootCurrent: ,,')

	efibootmgr -u |
		grep -E '^Boot[0-9]' |
		while read -r ln; do
			local hash entry loader id

			hash=$(hashsum <<<"$ln")
			ln=$(sed -e 's|Boot\([^\*]*\)\*\s\(.*\)\sHD(\([0-9]*\),[^,]*,\([^,]*\),.*)/File(\([^)]*\))\(.*\)$|\1\t\2\t\3\t\4\t\5\t\6|' <<<"$ln")
			entry=$(cut -d $'\t' -f1 <<<"$ln")
			loader=$(cut -d $'\t' -f5 <<<"$ln")
			version=$(sed -e 's,^.*vmlinuz-\(.*\)\.efi$,\1,' <<<"$loader")
			if [ -f "$(path_from_efi <<<"$loader")" ]; then
				jq -cn -S \
					--arg hash "$hash" \
					--arg entry "$entry" \
					--arg loader "$loader" \
					--arg version "$version" \
					--argjson bootorder "$(grep "$entry" <<<"$bootorder" | cut -d $'\t' -f2)" \
					--argjson current "$(test "$entry" = "$bootcurrent" && echo true || echo false)" \
					--arg disk "$(cut -d $'\t' -f4 <<<"$ln")" \
					--arg label "$(cut -d $'\t' -f2 <<<"$ln")" \
					--argjson part "$(cut -d $'\t' -f3 <<<"$ln")" \
					--arg args "$(cut -d $'\t' -f6 <<<"$ln")" \
					'{version:$version,entry:$entry,label:$label,loader:$loader,args:$args,bootorder:$bootorder,current:$current,disk:$disk,part:$part,hash:$hash}'
			fi
		done |
		jq -s 'sort_by(.bootorder)'
)

create_efi() (
	jq -rc '.[]' |
		while read -r ln; do
			local entry hash

			entry=$(jq -r '.entry' <<<"$ln" || echo none)
			hash=$(jq -r '.hash' <<<"$ln" || echo none)
			loader=$(jq -r '.loader' <<<"$ln")
			bootorder=$(jq -r '.bootorder' <<<"$ln")
			label=$(jq -r '.label' <<<"$ln")
			args=$(jq -r '.args' <<<"$ln")
			disk=$(findmnt -no SOURCE --target "$(path_from_efi <<<"$loader")" || echo none)
			part=$(sed -e 's,^.*p\([0-9]\)$,\1,' <<<"$disk")
			disk=${disk//p${part}/}

			_ efibootmgr -u --remove-dups --create --disk "$disk" --part "$part" --label "$label" --loader "$loader" -u "$args" --index "$bootorder"
		done
)

delete_efi() (
	local entry=$1
	_ efibootmgr -u --delete-bootnum --bootnum "$entry"
)
bootnext_efi() (
	local entry=$1
	_ efibootmgr -u --bootnext "$entry"
)
activate_efi() (
	local entry=$1
	_ efibootmgr -u --active --bootnum "$entry"
)
deactivate_efi() (
	local entry=$1
	_ efibootmgr -u --inactive --bootnum "$entry"
)

parse_kernels() (
	local kernel_dir=${1:-/boot}
	test -d "$kernel_dir" || eerror "kernel_dir should be a directory"

	local current_kernel current_hash
	current_kernel=$(parse_efi | jq -rc '.[] | select(.current==true) | .loader' | path_from_efi)
	current_hash=$(hashsum <"$current_kernel")

	find "$kernel_dir" -maxdepth 1 -type f -name 'vmlinuz*' -print |
		while read -r ln; do
			local file mtime ver ext hash aux="[]"

			file=$ln
			aux=$(jq -nc \
				--arg img "$(kernel_to_initramfs "$file")" \
				--arg sys "$(kernel_to_aux "$file" System.map)" \
				--arg con "$(kernel_to_aux "$file" config)" \
				'[$img,$sys]')

			hash=$(hashsum <"$file")
			mtime=$(stat -c '%Y' "$file")
			ln=$(sed -e "s,^${kernel_dir}/\([^-]*\)-\([0-9]*\.[0-9]*\.[0-9]*[^\.]*\)\(.*\)\$,\1\t\2\t\3," <<<"$ln")
			ver=$(cut -d $'\t' -f2 <<<"$ln")
			ext=$(cut -d $'\t' -f3 <<<"$ln")
			jq -cn -S \
				--arg file "$file" \
				--argjson old "$(grep -q old <<<"$ext" && echo true || echo false)" \
				--argjson current "$(test "$hash" = "$current_hash" && echo true || echo false)" \
				--arg version "$ver" \
				--argjson mtime "$mtime" \
				--arg hash "$hash" \
				--argjson aux "$aux" \
				'{file:$file,old:$old,version:$version,mtime:$mtime,hash:$hash,current:$current,aux:$aux}'
		done |
		jq -s 'sort_by(.mtime)|reverse'
)

kernel_to_aux() {
	local kernel=$1 aux_name=$2 aux_ext=$3

	aux=$(dirname "$kernel")/$(basename "$kernel" | sed -e "s,^vmlinuz,${aux_name},")
	if [ -n "$aux_ext" ]; then
		if grep -Eq '\.efi$' <<<"$aux"; then
			aux=$(sed -e "s,\.efi$,.${aux_ext}," <<<"$aux")
		elif grep -Eq '\.old$' <<<"$aux"; then
			aux=$(sed -e "s,\.old$,.${aux_ext}.old," <<<"$aux")
		else
			aux=$(sed -e "s,$,.${aux_ext}," <<<"$aux")
		fi
	fi
	# echo "+" "$kernel" "->" "$aux" >&2
	echo "$aux"
}

kernel_to_initramfs() {
	kernel_to_aux "$1" initramfs img
}

install_kernel() {
	local file=$1 target=${EFI_PARTITION}${EFI_DIR} newfile
	newfile=$target/$(basename "$file").efi

	_ cp -vf "$file" "$newfile" >&2
	_ cp -vf "$(kernel_to_initramfs "$file")" "$target/" >&2

	echo "$newfile"
}

remove_kernel() {
	local kernel_version=$1 file

	kernel_entries=$(parse_kernels | jq -cr --arg version "$kernel_version" '.[] | select(.version==$version)')
	for kernel in "$kernel_entries"; do
		file="$(jq -r '.file' <<<"$kernel")"
		test -f "$file" && _ rm -vf "$file" >&2

		efi_kernel=${EFI_PARTITION}${EFI_DIR}/$(basename "$file").efi
		test -f "$efi_kernel" && _ rm -vf "$efi_kernel" >&2

		for file in $(jq -cr '.aux[]' <<<"$kernel"); do
			test -f "$file" && _ rm -vf "$file" >&2

			if [[ "$file" = *initramfs* ]]; then
				efi_initramfs=${EFI_PARTITION}${EFI_DIR}/$(basename "$file")
				test -f "$efi_initramfs" && _ rm -vf "$efi_initramfs" >&2
			fi
		done
	done
}

# $1 : new kernel version, e.g. 6.1.27-gentoo-dist
# $2 : new kernel path, e.g. /boot/vmlinuz-6.1.27-gentoo-dist
new_kernel() {
	local version=$1 kernel=$2
	local kernel_dir efi_kernel efi_entries

	test -f "$kernel" || eerror "kernel needs to be a file"

	einfo "New kernel version ${version}"

	kernel_dir=$(dirname "$kernel")
	efi_kernel=${EFI_PARTITION}${EFI_DIR}/$(basename "$kernel").efi
	efi_entries=$(parse_efi | jq -rc)

	ebegin "Remove old and duplicate boot entries"
	# efi entries that are NOT the current version
	efi_entries_old=$(jq -rc '. as $in | .[] | select(.current) | .version as $current | $in | map(select(.version!=$current))[] | {entry,version}' <<<"$efi_entries")
	for efi_entry in $efi_entries_old; do
		delete_efi "$(jq -rc '.entry' <<<"$efi_entry")"
		remove_kernel "$(jq -rc '.version' <<<"$efi_entry")"
	done
	eend $?

	ebegin "Install new kernel"
	efi_kernel=$(install_kernel "$kernel")
	eend $?

	ebegin "Create new boot entry (debug)"
	jq -cnS \
		--arg loader "$(path_to_efi <<<"$efi_kernel")" \
		--arg label "Gentoo Linux $version (debug)" \
		--arg args "initrd=$(kernel_to_initramfs "$efi_kernel" | path_to_efi) rd.cmdline=ask" \
		'[{label:$label,loader:$loader,args:$args}]' |
		create_efi
	eend $?

	ebegin "Create new boot entry (default)"
	jq -cnS \
		--arg loader "$(path_to_efi <<<"$efi_kernel")" \
		--arg label "Gentoo Linux $version" \
		--arg args "initrd=$(kernel_to_initramfs "$efi_kernel" | path_to_efi)" \
		'[{label:$label,loader:$loader,args:$args}]' |
		create_efi
	eend $?

	# refresh efi+kernel entries
	efi_entries=$(parse_efi | jq -rc)
	kernel_entries=$(parse_kernels | jq -rc)

	ebegin "Remove unreferenced kernels"
	kernel_entries_unref=$(jq -rc --argjson efi "$efi_entries" '. as $kernels | $efi | map(.version) as $refs | $kernels | map(select(. as $this | ($refs|contains([($this|.version)])|not) )) | .[].version' <<<"$kernel_entries")
	for kernel_version in $kernel_entries_unref; do
		remove_kernel "$kernel_version"
	done
	eend $?
}

# check for cli usage
efi=1
kernel=
interactive=
list=
json=
while getopts KEhlij opt; do
	case ${opt} in
	E) efi=1 kernel= ;;
	K) kernel=1 efi= ;;
	i) interactive=1 ;;
	l) list=1 ;;
	j) json=1 ;;
	h | ?)
		printf "usage: %s [-hilj]\n" "$(basename "$0")"
		printf "\n"
		exit 2
		;;
	esac
done
shift $((OPTIND - 1))

if [ "$(id -u)" -ne 0 ]; then
	eerror "need to be root"
fi

# always mount efivarfs first
mount_efi

list_efi() {
	parse_efi |
		jq 'map({A:(if .current then "*" else " " end),boot:.bootorder,entry,label})' |
		/home/robert/bin/jq-table
}
list_kernels() {
	parse_kernels |
		jq 'map({A:(if .current then "*" else " " end),hash:(.hash[:6]),version,O:(if .old then "*" else " " end)})|sort_by(.version)' |
		/home/robert/bin/jq-table
}

# run list
if [ -n "$list" ]; then
	[ -n "$efi" ] && list_efi
	[ -n "$kernel" ] && list_kernels
	exit $?
fi

if [ -n "$json" ]; then
	[ -n "$efi" ] && parse_efi
	[ -n "$kernel" ] && parse_kernels
	exit $?
fi

# run interactive
if [ -n "$interactive" ]; then
	export BEMENU_BACKEND=curses

	entry=$(list_efi | tail -n+2 |
		bemenu -p ' efi: select boot entry ' |
		cut -c5- |
		sed -e 's,^\s*,,')
	[ -z "$entry" ] && exit

	bootnum=$(cut -d' ' -f1 <<<"$entry")
	[ -z "$bootnum" ] && exit

	action=$(cat <<-EOF | bemenu -p " efi: $(cut -d' ' -f2 <<<"$entry") "
	delete
	activate
	deactivate
	bootnext
EOF
	)
	[ -z "$action" ] && exit

	echo ">>>" "$action"

	case "$action" in
	delete) delete_efi "$bootnum" ;;
	activate) activate_efi "$bootnum" ;;
	deactivate) deactivate_efi "$bootnum" ;;
	bootnext) bootnext_efi "$bootnum" ;;
	esac
	exit $?
fi


# run install routine, used by installkernel
test -n "$1" || eerror "need version argument"
test -n "$2" || eerror "need kernel argument"
new_kernel "$1" "$2"
exit $?