#!/bin/bash
set -euo pipefail

/usr/sbin/grub-install.real "$@"

# Check if /boot is a separate partition
if findmnt -M "/boot" >/dev/null 2>&1; then
    echo "==> Separate /boot partition detected. Localize not needed"
    exit 0
fi

if [ ! -d /sys/firmware/efi ]; then
    echo "==> Legacy BIOS / CSM System no efi patching"
    exit 0
fi

ROOT_SOURCE=$(findmnt -n -o SOURCE /)
ROOT_DEV=$(printf '%s\n' "$ROOT_SOURCE" | sed 's/\[.*$//')
ROOT_DEV=$(readlink -f "$ROOT_DEV")
LUKS_RAW_UUID=$(lsblk -s -n -o FSTYPE,UUID "$ROOT_DEV" | awk '$1=="crypto_LUKS" {print $2; exit}')
if [ -z "$LUKS_RAW_UUID" ]; then
    echo "==> Unencrypted system detected (no LUKS2 volume found)."
    exit 0
fi

### 0. Configuration ###
# TEST_MODE=1: Creates a SEPARATE efibootmgr entry, leaves /EFI/TUXEDO untouched.
# TEST_MODE=0: Directly replaces the production grubx64.efi (with backup).
TEST_MODE=0

### 1. Prerequisites ###
for cmd in ckbcomp grub-mklayout grub-mkimage findmnt blkid; do
    command -v "$cmd" >/dev/null || { echo "ERROR: $cmd missing"; exit 1; }
done


ESP_MOUNT=$(findmnt -rn -t vfat -M /boot/efi -o TARGET)
ESP_UUID=$(findmnt -rn -t vfat -M /boot/efi -o UUID)
WORKDIR=$(mktemp -d)
trap 'rm -rf "$WORKDIR"' EXIT

if [ "$TEST_MODE" = "1" ]; then
    TARGET_DIR="$ESP_MOUNT/EFI/keymaptest"
else
    TARGET_DIR="$ESP_MOUNT/EFI/TUXEDO"
    [ -d "$TARGET_DIR" ] || { echo "ERROR: $TARGET_DIR does not exist"; exit 1; }
fi
mkdir -p "$TARGET_DIR"
GRUB_TARGET="$TARGET_DIR/grubx64.efi"

### 2. Dynamically read keymap ###
RAW_KEYMAP=""
KEYMAP=""

if [ -f /etc/default/keyboard ]; then
    XKBLAYOUT=""
    . /etc/default/keyboard
    RAW_KEYMAP="${XKBLAYOUT:-}"
fi

if [ -z "$RAW_KEYMAP" ] && [ -f /etc/vconsole.conf ]; then
    KEYMAP=""
    XKBLAYOUT=""
    . /etc/vconsole.conf 2>/dev/null || true
    RAW_KEYMAP="${KEYMAP:-${XKBLAYOUT:-}}"
fi

if [ -z "$RAW_KEYMAP" ]; then
    RAW_KEYMAP=$(
        localectl status 2>/dev/null |
        awk -F: '/VC Keymap:/ {
            gsub(/^[ \t]+|[ \t]+$/, "", $2)
            if ($2 != "(unset)") print $2
            exit
        }'
    )
fi

if [ -z "$RAW_KEYMAP" ]; then
    RAW_KEYMAP=$(
        localectl status 2>/dev/null |
        awk -F: '/X11 Layout:/ {
            gsub(/^[ \t]+|[ \t]+$/, "", $2)
            print $2
            exit
        }'
    )
fi

KEYMAP="${RAW_KEYMAP%%,*}"
KEYMAP=$(printf '%s' "$KEYMAP" | sed 's/[()[:space:]]//g')
KEYMAP="${KEYMAP:-de}"

echo "==> Layout: $KEYMAP"

### 3. Build .gkb layout file ###
PROD_LAYOUT_DIR="$ESP_MOUNT/EFI/TUXEDO/layouts"
mkdir -p "$PROD_LAYOUT_DIR"
ckbcomp "$KEYMAP" | grub-mklayout -o "$PROD_LAYOUT_DIR/${KEYMAP}.gkb"

### 4. Early-Config ###
cat > "$WORKDIR/early.cfg" << EOF
insmod part_gpt
insmod fat
search --no-floppy --fs-uuid --set=esp $ESP_UUID

insmod terminal
insmod at_keyboard
insmod keylayouts
terminal_input at_keyboard
terminal_output console

keymap (\$esp)/EFI/TUXEDO/layouts/${KEYMAP}.gkb

insmod cryptodisk
insmod luks2
insmod btrfs
insmod lzopio
insmod zfs
insmod xzio
insmod gzio
insmod zstd
insmod zstdio
insmod ext2
insmod xfs
insmod f2fs
insmod lvm
insmod search_fs_uuid

configfile (\$esp)/EFI/TUXEDO/grub.cfg
EOF

### 5. Build core image ###
grub-mkimage \
    -O x86_64-efi \
    -o "$WORKDIR/grubx64_new.efi" \
    -c "$WORKDIR/early.cfg" \
    -p "" \
    part_gpt part_msdos fat btrfs zstd zstdio ext2 xfs f2fs lvm lzopio xzio gzio \
    crypto cryptodisk luks luks2 gcry_rijndael gcry_sha256 gcry_sha512 pbkdf2 argon2 \
    search search_fs_uuid keylayouts \
    at_keyboard terminal \
    configfile normal boot linux echo test true help ls cat halt reboot \
    efi_gop video video_fb gfxterm loadenv

### 6. Install image (unsigned) ###
if [ "$TEST_MODE" = "1" ]; then
    install -Dm644 "$WORKDIR/grubx64_new.efi" "$GRUB_TARGET"

    ORIG_SHIM=$(find "$ESP_MOUNT/EFI" -iname "shimx64.efi" ! -path "*keymaptest*" | head -n1 || true)
    if [ -n "$ORIG_SHIM" ]; then
        install -Dm644 "$ORIG_SHIM" "$TARGET_DIR/shimx64.efi"
    fi

    DISK=$(findmnt /boot/efi -no SOURCE | sed -E 's/p?[0-9]+$//')
    PARTNUM=$(findmnt /boot/efi -no SOURCE | grep -oE '[0-9]+$')
    efibootmgr --create --disk "$DISK" --part "$PARTNUM" \
        --loader '\EFI\keymaptest\grubx64.efi' \
        --label "GRUB Keymap Test ($KEYMAP)"

    echo "==> TEST_MODE: New separate boot entry created."
    echo "==> Existing production entry remains UNCHANGED."
    echo "==> Select this entry in firmware boot menu (F11/F12/Esc) for testing."
else
    if [ -f "$GRUB_TARGET" ]; then
        BACKUP="$GRUB_TARGET.bak-$(date +%Y%m%d-%H%M%S)"
        cp -a "$GRUB_TARGET" "$BACKUP"
        echo "==> Backup of previous grubx64.efi: $BACKUP"
    fi
    install -Dm644 "$WORKDIR/grubx64_new.efi" "$GRUB_TARGET"

    echo "==> PRODUCTION: $GRUB_TARGET successfully replaced."
    echo "==> Existing efibootmgr entry already points to this path."
fi

