#!/usr/bin/bash
set -Eeuo pipefail
PATH=${NABU_MAINT_PATH:-/usr/sbin:/usr/bin}
export PATH

readonly state_dir=${NABU_MAINT_STATE_DIR:-/var/lib/nabu-kernel-maintenance}
readonly pending_dir=$state_dir/pending.d
readonly prepared_dir=$state_dir/prepared.d
readonly identity_tool=${NABU_MAINT_IDENTITY_TOOL:-/usr/libexec/senemos-nabu/kernel-build-identity}
readonly boot_dir=${NABU_MAINT_BOOT_DIR:-/boot}
readonly esp=${NABU_MAINT_ESP:-/boot/efi}
readonly regenerate_command=${NABU_MAINT_REGENERATE_COMMAND:-/usr/bin/nabu-regenerate-uki}
readonly configure_command=${NABU_MAINT_CONFIGURE_COMMAND:-/usr/bin/nabu-configure-boot-manager}
readonly compatible_path=${NABU_MAINT_COMPATIBLE_PATH:-/sys/firmware/devicetree/base/compatible}
readonly config_file=${NABU_MAINT_CONFIG:-/etc/nabu/kernel.conf}

lock_file=${NABU_MAINT_LOCK_FILE:-/run/nabu-kernel-maintenance.lock}
install -d -m0755 "${lock_file%/*}"
exec 9>"$lock_file"
flock -n 9 || exit 0

if [[ -r $compatible_path ]] && ! grep -aq 'xiaomi,nabu' "$compatible_path"; then
    exit 0
fi

declare -Ar selector=(
    [alpha]=senemos-nabu-kernel-alpha
    [stable]=senemos-nabu-kernel
    [mainline]=senemos-nabu-kernel-mainline-alpha
    [mainline-unstable]=senemos-nabu-kernel-mainline-unstable
    [lts]=senemos-nabu-kernel-lts
)
declare -Ar efi_family=(
    [alpha]=SENEMOS6
    [stable]=SENEMOS616
    [mainline]=SENEMOS7
    [mainline-unstable]=SENEMOS7U
    [lts]=SENEMOS6LTS
)

identity_is_newer() {
    local candidate=$1 current=$2 newest
    if [[ $candidate =~ ^[0-9]{10}$ && ! $current =~ ^[0-9]{10}$ ]]; then
        return 0
    fi
    if [[ ! $candidate =~ ^[0-9]{10}$ && $current =~ ^[0-9]{10}$ ]]; then
        return 1
    fi
    newest=$(printf '%s\n%s\n' "$current" "$candidate" | sort -V | tail -n1)
    [[ $newest == "$candidate" && $candidate != "$current" ]]
}

install -d -m0755 "$state_dir" "$pending_dir" "$prepared_dir"
declare -A selected_kernel=() selected_identity=() selected_channel=()
installed_count=0
for channel in alpha stable mainline mainline-unstable lts; do
    rpm -q "${selector[$channel]}" >/dev/null 2>&1 || continue
    installed_count=$((installed_count + 1))
    package_nevra=$(rpm -q --qf '%{NEVRA}\n' "${selector[$channel]}" | tail -n1)
    mapfile -t kernels < <(rpm -ql "$package_nevra" | sed -n "s@^${boot_dir}/vmlinuz-@@p")
    [[ ${#kernels[@]} -eq 1 ]] || {
        printf 'Expected one kernel image in %s, found %s.\n' "$package_nevra" "${#kernels[@]}" >&2
        exit 3
    }
    kernel=${kernels[0]}
    owner=$(rpm -qf --qf '%{NAME}\n' "$boot_dir/vmlinuz-$kernel" 2>/dev/null | head -n1 || :)
    [[ $owner == "${selector[$channel]}" ]] || {
        printf 'Kernel %s is not owned by %s.\n' "$kernel" "${selector[$channel]}" >&2
        exit 4
    }
    [[ -s $boot_dir/vmlinuz-$kernel && -d ${NABU_MAINT_MODULES_DIR:-/usr/lib/modules}/$kernel ]] || {
        printf 'Incomplete kernel payload: %s\n' "$kernel" >&2
        exit 5
    }
    identity=$("$identity_tool" --field=stamp "$kernel" 2>/dev/null || :)
    if [[ ! $identity =~ ^[0-9]{10}$ ]]; then
        package_release=$(rpm -q --qf '%{RELEASE}\n' "$package_nevra" | tail -n1)
        identity=$(sed -nE 's/^([0-9]{10})([.]|$).*/\1/p' <<<"$package_release")
    fi
    [[ $identity =~ ^[0-9]{10}$ ]] || {
        printf 'Cannot derive the SENEMOS identity for %s from the kernel or RPM release.\n' \
            "$kernel" >&2
        exit 6
    }
    family=${efi_family[$channel]}
    if [[ -z ${selected_kernel[$family]:-} ]] ||
       identity_is_newer "$identity" "${selected_identity[$family]}"; then
        selected_kernel[$family]=$kernel
        selected_identity[$family]=$identity
        selected_channel[$family]=$channel
    fi
done
(( installed_count > 0 )) || { printf 'No managed Nabu kernel family is installed.\n' >&2; exit 2; }

preferred=alpha
if [[ -r $config_file ]]; then
    configured=$(sed -nE 's/^preferred=(alpha|stable|mainline|mainline-unstable|lts)$/\1/p' "$config_file" | tail -n1)
    [[ -z $configured ]] || preferred=$configured
fi
preferred_family=${efi_family[$preferred]}
if [[ -z ${selected_kernel[$preferred_family]:-} ]]; then
    mapfile -t available_families < <(printf '%s\n' "${!selected_kernel[@]}" | sort -V)
    preferred_family=${available_families[0]}
fi

android_entry=$esp/loader/entries/android.conf
android_efi=
android_entry_digest=
android_efi_digest=
if [[ -r $android_entry ]]; then
    android_efi=$(sed -nE 's@^efi[[:space:]]+(/.*)@\1@p' "$android_entry" | head -n1)
    [[ -n $android_efi && -s $esp$android_efi ]] || {
        printf 'Android return entry is incomplete.\n' >&2
        exit 7
    }
    android_entry_digest=$(sha256sum "$android_entry" | cut -d' ' -f1)
    android_efi_digest=$(sha256sum "$esp$android_efi" | cut -d' ' -f1)
fi

# Prepare non-default families first. The final sync makes the preferred family
# the default while retaining exactly one latest UKI for every known family.
mapfile -t ordered_families < <(printf '%s\n' "${!selected_kernel[@]}" | sort -V)
for family in "${ordered_families[@]}" "$preferred_family"; do
    [[ $family != "$preferred_family" || ${family_seen:-0} -eq 0 ]] || continue
    [[ $family == "$preferred_family" ]] && family_seen=1
    kernel=${selected_kernel[$family]}
    identity=${selected_identity[$family]}
    uki_name="${family}-${identity}.efi"
    prepared_file=$prepared_dir/$family
    regenerate=1
    if [[ -r $prepared_file ]] && read -r prepared_kernel prepared_digest <"$prepared_file" &&
       [[ $prepared_kernel == "$kernel" && $prepared_digest =~ ^[0-9a-f]{64}$ ]] &&
       [[ -s $esp/EFI/SENEMOS/$uki_name ]] &&
       [[ $(sha256sum "$esp/EFI/SENEMOS/$uki_name" | cut -d' ' -f1) == "$prepared_digest" ]]; then
        regenerate=0
    fi
    if (( regenerate )); then
        "$regenerate_command" --family "$family" "$kernel"
        [[ -s $esp/EFI/SENEMOS/$uki_name ]] || {
            printf 'Generated family UKI is incomplete: %s\n' "$uki_name" >&2
            exit 8
        }
        temporary=$(mktemp "$prepared_dir/.${family}.XXXXXX")
        uki_digest=$(sha256sum "$esp/EFI/SENEMOS/$uki_name" | cut -d' ' -f1)
        printf '%s %s\n' "$kernel" "$uki_digest" >"$temporary"
        chmod 0644 "$temporary"
        mv -f "$temporary" "$prepared_file"
    fi
    for channel in alpha stable mainline mainline-unstable lts; do
        [[ ${efi_family[$channel]} == "$family" ]] && rm -f -- "$pending_dir/$channel"
    done
done

preferred_uki="${preferred_family}-${selected_identity[$preferred_family]}.efi"
"$configure_command" --sync-only --uki "$preferred_uki"

if [[ -n $android_efi ]]; then
    [[ $(sha256sum "$android_entry" | cut -d' ' -f1) == "$android_entry_digest" &&
       $(sha256sum "$esp$android_efi" | cut -d' ' -f1) == "$android_efi_digest" ]] || {
        printf 'Android return artifact changed unexpectedly.\n' >&2
        exit 9
    }
fi
printf 'Prepared %s managed kernel families; default family: %s.\n' \
    "${#selected_kernel[@]}" "$preferred_family"
