#!/bin/sh

# u-boot-update - Update U-Boot configuration
# Copyright (C) 2025 BredOS <https://www.bredos.org>

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

set -e

_U_BOOT_DIRECTORY="/boot/extlinux"

# Debug function to print variables
debug_print() {
    echo "DEBUG: $1 = $2"
}

Update ()
{
    # Update target file using source content
    _TARGET="${1}"
    _SOURCE="${2}"

    _TMPFILE="${_TARGET}.tmp"
    rm -f "${_TMPFILE}"

    echo "${_SOURCE}" > "${_TMPFILE}"

    if [ -e "${_TARGET}" ] && cmp -s "${_TARGET}" "${_TMPFILE}"
    then
        rm -f "${_TMPFILE}"
    else
        # FIXME: should use fsync here
        echo "P: Updating ${_TARGET}..."
        mv -f "${_TMPFILE}" "${_TARGET}"
    fi
}

# FIXME: switch to check extlinux file can be written to
# User is unprivileged
if [ "$(id -u)" -ne 0 ]
then
    echo "E: need root privileges"
    exit 1
fi

# Redirect stdout to stderr due Debconf usage
exec 1>&2

. /usr/share/u-boot-menu/read-config

# Checking extlinux directory
printf '%s' "P: Checking for EXTLINUX directory..."

# Creating extlinux directory
if [ ! -e "${_U_BOOT_DIRECTORY}" ]
then
    echo " not found."

    printf '%s' "P: Creating EXTLINUX directory..."
    mkdir -p "${_U_BOOT_DIRECTORY}"
    echo " done: ${_U_BOOT_DIRECTORY}"
else
    echo " found."
fi

# Create extlinux.conf
_CONFIG="\
## ${_U_BOOT_DIRECTORY}/extlinux.conf
##
## IMPORTANT WARNING
##
## The configuration of this file is generated automatically.
## Do not edit this file manually, use: u-boot-update

default ${U_BOOT_DEFAULT}
menu title BredOS U-Boot menu
prompt ${U_BOOT_PROMPT}
timeout ${U_BOOT_TIMEOUT}
"

# Find parameter for root from fstab
if [ -z "${U_BOOT_ROOT}" ]
then
    # Find root partition
    while read _LINE
    do
        # Skip empty lines and comments
        case "${_LINE}" in
            ""|"#"*) continue ;;
        esac

        read _FS_SPEC _FS_FILE _FS_VFSTYPE _FS_MNTOPS _FS_FREQ _FS_PASSNO << EOF
${_LINE}
EOF

        if [ "${_FS_FILE}" = "/" ]
        then
            # Check if root device is encrypted (LUKS)
            _IS_LUKS=0
            _ROOT_DEV="${_FS_SPEC}"
            
            # If root is specified by UUID
            if echo "${_ROOT_DEV}" | grep -q "^UUID="; then
                _ROOT_UUID=$(echo "${_ROOT_DEV}" | cut -d= -f2)
                _ROOT_DEV=$(findfs UUID="${_ROOT_UUID}" 2>/dev/null || echo "${_ROOT_DEV}")
            fi
            
            # If root is specified by LABEL
            if echo "${_ROOT_DEV}" | grep -q "^LABEL="; then
                _ROOT_LABEL=$(echo "${_ROOT_DEV}" | cut -d= -f2)
                _ROOT_DEV=$(findfs LABEL="${_ROOT_LABEL}" 2>/dev/null || echo "${_ROOT_DEV}")
            fi
            
            # Check if root is on LUKS
            if echo "${_ROOT_DEV}" | grep -q "^/dev/mapper/luks-"; then
                _IS_LUKS=1
                echo "P: Detected encrypted root filesystem at ${_ROOT_DEV}"
                
                # Find the real device that contains the LUKS container
                if command -v lsblk >/dev/null 2>&1 && command -v jq >/dev/null 2>&1; then
                    _REAL_DEV=$(lsblk -sJp | jq -r --arg dsk "${_ROOT_DEV}" '.blockdevices | .[] | select(.name == $dsk) | .children | .[0] | .name' 2>/dev/null)
                    
                    if [ -n "${_REAL_DEV}" ]; then
                        echo "P: Found LUKS physical device: ${_REAL_DEV}"
                        _REAL_UUID=$(blkid -s UUID -o value "${_REAL_DEV}" 2>/dev/null)
                        
                        if [ -n "${_REAL_UUID}" ]; then
                            echo "P: Real device UUID: ${_REAL_UUID}"
                            # Add cryptdevice parameter
                            U_BOOT_CRYPTDEVICE="cryptdevice=UUID=${_REAL_UUID}:luks-${_REAL_UUID}"
                        fi
                    fi
                else
                    echo "W: lsblk or jq not available, unable to determine physical device for LUKS"
                fi
            fi
            
            # Get the UUID of the root filesystem
            _ROOT_UUID=$(blkid -s UUID -o value "${_ROOT_DEV}" 2>/dev/null)
            if [ -n "${_ROOT_UUID}" ]; then
                U_BOOT_ROOT="root=UUID=${_ROOT_UUID}"
            else
                U_BOOT_ROOT="root=${_FS_SPEC}"
            fi
            
            # Auto-detect filesystem type and add appropriate parameters
            if [ "${_FS_VFSTYPE}" = "btrfs" ]; then
                echo "P: Detected btrfs filesystem for root"
                if ! echo "${U_BOOT_PARAMETERS}" | grep -q "rootfstype="; then
                    U_BOOT_PARAMETERS="${U_BOOT_PARAMETERS} rootfstype=btrfs"
                fi
                # Add subvolume parameter if not already present
                if ! echo "${U_BOOT_PARAMETERS}" | grep -q "rootflags="; then
                    # Check for subvolume in mount options
                    if echo "${_FS_MNTOPS}" | grep -q "subvol="; then
                        _SUBVOL=$(echo "${_FS_MNTOPS}" | sed -n 's/.*subvol=\([^,]*\).*/\1/p')
                        # Strip leading slash if present
                        _SUBVOL=$(echo "${_SUBVOL}" | sed 's|^/||')
                        echo "P: Found subvolume: ${_SUBVOL}"
                        U_BOOT_PARAMETERS="${U_BOOT_PARAMETERS} rootflags=subvol=${_SUBVOL}"
                    else
                        # Default to @ subvolume common in Arch Linux
                        echo "P: No subvolume specified, using default @"
                        U_BOOT_PARAMETERS="${U_BOOT_PARAMETERS} rootflags=subvol=@"
                    fi
                fi
            elif [ "${_FS_VFSTYPE}" = "ext4" ]; then
                echo "P: Detected ext4 filesystem for root"
                if ! echo "${U_BOOT_PARAMETERS}" | grep -q "rootfstype="; then
                    U_BOOT_PARAMETERS="${U_BOOT_PARAMETERS} rootfstype=ext4"
                fi
            fi
            
            # Add rootwait parameter
            if ! echo "${U_BOOT_PARAMETERS}" | grep -q "rootwait"; then
                U_BOOT_PARAMETERS="${U_BOOT_PARAMETERS} rootwait"
            fi
            
            # For LUKS, add the cryptdevice parameter
            if [ "${_IS_LUKS}" = "1" ] && [ -n "${U_BOOT_CRYPTDEVICE}" ]; then
                U_BOOT_PARAMETERS="${U_BOOT_PARAMETERS} ${U_BOOT_CRYPTDEVICE}"
            fi
            
            echo "P: Root parameters: ${U_BOOT_ROOT}"
            echo "P: Kernel parameters: ${U_BOOT_PARAMETERS}"
            break
        fi
    done < /etc/fstab
fi

# if not in fstab, try from current kernel arguments
if [ -z "${U_BOOT_ROOT}" ]
then
    # Check if we're in a chroot environment
    if command -v ischroot >/dev/null 2>&1 && ischroot
    then
        echo "W: Running in chroot and no root found in fstab. Bailing out."
        exit 0
    fi

    echo "P: No root found in fstab, checking kernel command line"
    if [ -f "/proc/cmdline" ]; then
        for param in $(cat /proc/cmdline)
        do
            case $param in
              root=*)
                U_BOOT_ROOT="$param"
                echo "P: Found root in kernel cmdline: ${U_BOOT_ROOT}"
                break
                ;;
            esac
        done
    else
        echo "W: /proc/cmdline not available (chroot?)"
    fi
fi

# If U_BOOT_FDT is specified, search for and copy that file to boot
if [ -n "${U_BOOT_FDT}" ]; then
    echo "P: Searching for specified DTB file: ${U_BOOT_FDT}"
    _DTB_FOUND=0
    
    # Common DTB locations to search - excluding the destination path
    _DTB_SEARCH_PATHS="
        /usr/lib/modules/*/dtbs
        /usr/lib/modules/*/dtb
        /lib/modules/*/dtbs
        /lib/modules/*/dtb
        /boot/dtbs
    "
    
    for _SEARCH_PATH in ${_DTB_SEARCH_PATHS}; do
        # Use find to locate the DTB file (handling wildcards in the path)
        find_cmd="find ${_SEARCH_PATH} -name \"${U_BOOT_FDT}\" -type f 2>/dev/null | head -n 1"
        _FOUND_DTB=$(eval ${find_cmd} || true)
        
        if [ -n "${_FOUND_DTB}" ] && [ -f "${_FOUND_DTB}" ]; then
            echo "P: Found DTB file at: ${_FOUND_DTB}"
            _DTB_FOUND=1
            break
        fi
    done
    
    if [ "${_DTB_FOUND}" -eq 0 ]; then
        echo "W: Specified DTB file '${U_BOOT_FDT}' not found in any standard location."
    fi
fi

# Auto find kernels
_NUMBER=0
_ENTRY=1

# Find all kernel files in /boot
if ls /boot/vmlinuz-* >/dev/null 2>&1; then
    _KERNEL_FILES=$(ls -1 /boot/vmlinuz-* 2>/dev/null | sort -V -r)
fi

# Process each kernel
for _KERNEL_FILE in ${_KERNEL_FILES}; do
    # Get kernel filename without /boot prefix
    _KERNEL=$(basename "${_KERNEL_FILE}")
    
    # Determine version based on kernel name format
    if echo "${_KERNEL}" | grep -q "^vmlinuz-"; then
        _VERSION="${_KERNEL#vmlinuz-}"
    elif echo "${_KERNEL}" | grep -q "^Image.gz-"; then
        _VERSION="${_KERNEL#Image.gz-}"
    elif [ "${_KERNEL}" = "Image.gz" ]; then
        _VERSION="linux"
    else
        # Skip unknown format
        continue
    fi
    
    # Get package base for this kernel
    _PKGBASE=""
    if [ -f "/lib/modules/${_VERSION}/pkgbase" ]; then
        _PKGBASE=$(cat "/lib/modules/${_VERSION}/pkgbase" 2>/dev/null)
        if [ -n "${_PKGBASE}" ]; then
            echo "P: Found package base for kernel ${_VERSION}: ${_PKGBASE}"
        else
            echo "W: Package base file exists but is empty for kernel ${_VERSION}"
            _PKGBASE="${_VERSION}"  # Fallback to version if pkgbase is empty
        fi
    else
        echo "W: No package base file found for kernel ${_VERSION}, using version as package name"
        _PKGBASE="${_VERSION}"
    fi
    
    # Set up package-specific DTB directory path
    _PKG_DTB_DIR="/boot/dtb/${_PKGBASE}"
    
    # Create the directory if it doesn't exist
    if [ ! -d "${_PKG_DTB_DIR}" ]; then
        echo "P: Creating package DTB directory: ${_PKG_DTB_DIR}"
        mkdir -p "${_PKG_DTB_DIR}"
    fi
    
    # If DTB file was found earlier, copy it to the package-specific directory
    if [ -n "${U_BOOT_FDT}" ] && [ -n "${_FOUND_DTB}" ] && [ -f "${_FOUND_DTB}" ]; then
        # Copy DTB to package-specific directory
        echo "P: Copying DTB file to ${_PKG_DTB_DIR}/${U_BOOT_FDT}"
        cp -f "${_FOUND_DTB}" "${_PKG_DTB_DIR}/${U_BOOT_FDT}" 2>/dev/null || true
    fi
    
    # Handle overlays - clean existing overlays first
    if [ -d "${_PKG_DTB_DIR}/overlays" ]; then
        echo "P: Cleaning existing overlays in ${_PKG_DTB_DIR}/overlays"
        rm -rf "${_PKG_DTB_DIR}/overlays"
    fi
    
    # Create overlay directory
    mkdir -p "${_PKG_DTB_DIR}/overlays"
    
    # Process overlays if specified in U_BOOT_FDT_OVERLAYS
    if [ -n "${U_BOOT_FDT_OVERLAYS}" ]; then
        echo "P: Processing specified overlays: ${U_BOOT_FDT_OVERLAYS}"
        
        # Search paths for overlays
        _DTBO_SEARCH_PATHS="
            /boot/dtbs/
            /usr/lib/modules/*/dtb
            ${U_BOOT_FDT_OVERLAYS_DIR}
        "
        
        # Process each overlay specified in U_BOOT_FDT_OVERLAYS
        for _OVERLAY in ${U_BOOT_FDT_OVERLAYS}; do
            _OVERLAY_FOUND=0
            
            for _SEARCH_PATH in ${_DTBO_SEARCH_PATHS}; do
                # Search for the overlay file
                find_cmd="find ${_SEARCH_PATH} -name ${_OVERLAY} -type f 2>/dev/null | head -n 1"
                _FOUND_OVERLAY=$(eval ${find_cmd} || true)
                
                if [ -n "${_FOUND_OVERLAY}" ] && [ -f "${_FOUND_OVERLAY}" ]; then
                    echo "P: Found overlay file at: ${_FOUND_OVERLAY}"
                    echo "P: Copying overlay to ${_PKG_DTB_DIR}/overlays/${_OVERLAY}"
                    cp -f "${_FOUND_OVERLAY}" "${_PKG_DTB_DIR}/overlays/${_OVERLAY}" 2>/dev/null || true
                    _OVERLAY_FOUND=1
                    break
                fi
            done
            
            if [ "${_OVERLAY_FOUND}" -eq 0 ]; then
                echo "W: Specified overlay '${_OVERLAY}' not found in any standard location."
            fi
        done
    fi
    
    # Also check for overlays in the U_BOOT_FDT_OVERLAYS_DIR
    if [ -d "${U_BOOT_FDT_OVERLAYS_DIR}" ]; then
        for _OVERLAY_PATH in "${U_BOOT_FDT_OVERLAYS_DIR}"/*.dtbo; do
            if [ -f "${_OVERLAY_PATH}" ]; then
                _OVERLAY=$(basename "${_OVERLAY_PATH}")
                # Only copy if not already copied from U_BOOT_FDT_OVERLAYS
                if [ ! -f "${_PKG_DTB_DIR}/overlays/${_OVERLAY}" ]; then
                    echo "P: Copying overlay from ${U_BOOT_FDT_OVERLAYS_DIR} to ${_PKG_DTB_DIR}/overlays/${_OVERLAY}"
                    cp -f "${_OVERLAY_PATH}" "${_PKG_DTB_DIR}/overlays/${_OVERLAY}" 2>/dev/null || true
                fi
            fi
        done
    fi
    
    # Auto-detect matching initramfs - try common patterns
    _INITRD=""
    if [ -e "/boot/initramfs-${_VERSION}.img" ]; then
        _INITRD="initrd ${_BOOT_DIRECTORY}/initramfs-${_VERSION}.img"
    elif [ -e "/boot/initrd-${_VERSION}.img" ]; then
        _INITRD="initrd ${_BOOT_DIRECTORY}/initrd-${_VERSION}.img"
    elif [ -e "/boot/initrd.img-${_VERSION}" ]; then
        _INITRD="initrd ${_BOOT_DIRECTORY}/initrd.img-${_VERSION}"
    elif [ -e "/boot/initramfs-linux.img" ] && [ "${_VERSION}" = "linux" ]; then
        # Special case for Image.gz + initramfs-linux.img
        _INITRD="initrd ${_BOOT_DIRECTORY}/initramfs-linux.img"
    else
        echo "W: No matching initramfs found for ${_KERNEL}. Skipping this kernel."
        continue
    fi
    
    # Check for fallback/recovery initramfs
    _HAS_FALLBACK=0
    _FALLBACK_INITRD=""
    if [ -e "/boot/initramfs-${_VERSION}-fallback.img" ]; then
        _FALLBACK_INITRD="initrd ${_BOOT_DIRECTORY}/initramfs-${_VERSION}-fallback.img"
        _HAS_FALLBACK=1
    elif [ -e "/boot/initrd-${_VERSION}-fallback.img" ]; then
        _FALLBACK_INITRD="initrd ${_BOOT_DIRECTORY}/initrd-${_VERSION}-fallback.img"
        _HAS_FALLBACK=1
    elif [ -e "/boot/initramfs-linux-fallback.img" ] && [ "${_VERSION}" = "linux" ]; then
        # Special case for generic kernel + fallback initramfs
        _FALLBACK_INITRD="initrd ${_BOOT_DIRECTORY}/initramfs-linux-fallback.img"
        _HAS_FALLBACK=1
    fi
    
    # Handle device tree files - ensure the FDT entry is added to the configuration
    _FDT=""
    
    # Always check for package-specific DTB first
    if [ -n "${U_BOOT_FDT}" ] && [ -f "${_PKG_DTB_DIR}/${U_BOOT_FDT}" ]; then
        echo "P: Using package-specific DTB for ${_PKGBASE}: ${_PKG_DTB_DIR}/${U_BOOT_FDT}"
        _FDT="fdt ${_BOOT_DIRECTORY}/dtb/${_PKGBASE}/${U_BOOT_FDT}"
    # Then try the common /boot location
    elif [ -n "${U_BOOT_FDT}" ] && [ -f "/boot/${U_BOOT_FDT}" ]; then
        echo "P: Using common DTB file: /boot/${U_BOOT_FDT}"
        _FDT="fdt ${_BOOT_DIRECTORY}/${U_BOOT_FDT}"
    elif [ -n "${U_BOOT_FDT}" ]; then
        echo "W: Specified DTB file '${U_BOOT_FDT}' not found in package-specific or common location."
    fi
    
    # Process the overlays for this kernel
    _FDTOVERLAYS=""
    if [ -d "${_PKG_DTB_DIR}/overlays" ]; then
        _DTBO_LIST=""
        _OVERLAY_COUNT=0
        
        for _OVERLAY_PATH in "${_PKG_DTB_DIR}/overlays/"*.dtbo; do
            if [ -f "${_OVERLAY_PATH}" ]; then
                _OVERLAY=$(basename "${_OVERLAY_PATH}")
                _DTBO_LIST="${_DTBO_LIST} /dtb/${_PKGBASE}/overlays/${_OVERLAY}"
                _OVERLAY_COUNT=$((_OVERLAY_COUNT + 1))
            fi
        done
        
        if [ "${_OVERLAY_COUNT}" -gt 0 ]; then
            echo "P: Found ${_OVERLAY_COUNT} overlay(s) for kernel ${_VERSION}"
            _FDTOVERLAYS="fdtoverlays${_DTBO_LIST}"
        fi
    fi
    
    # Make sure we debug the _FDT value
    if [ -n "${_FDT}" ]; then
        echo "P: Will add FDT configuration: ${_FDT}"
    else
        echo "W: No valid FDT configuration found for kernel ${_VERSION} (${_PKGBASE})"
    fi
    
    if echo "${U_BOOT_ALTERNATIVES}" | grep -q default; then
        # Writing default entry
        echo "P: Creating default entry for kernel ${_VERSION}"
        
        # Build the entry
        _ENTRY="\
label l${_NUMBER}
	menu label ${U_BOOT_MENU_LABEL} ${_VERSION}
	linux ${_BOOT_DIRECTORY}/${_KERNEL}
	${_INITRD}"

        # Only add fdt line if _FDT is not empty
        if [ -n "${_FDT}" ]; then
            _ENTRY="${_ENTRY}
	${_FDT}"
        fi
        
        # Only add fdtoverlays line if _FDTOVERLAYS is not empty
        if [ -n "${_FDTOVERLAYS}" ]; then
            _ENTRY="${_ENTRY}
	${_FDTOVERLAYS}"
        fi
        
        _ENTRY="${_ENTRY}
	append ${U_BOOT_ROOT} ${U_BOOT_PARAMETERS}"
        
        # Add the entry to the configuration
        _CONFIG="${_CONFIG}

${_ENTRY}"
    fi
    
    # Only create recovery entry if we have a fallback initramfs and recovery is enabled
    if [ "${_HAS_FALLBACK}" = "1" ] && echo "${U_BOOT_ALTERNATIVES}" | grep -q recovery; then
        echo "P: Creating recovery entry for kernel ${_VERSION}"
        
        # Build the recovery entry
        _ENTRY="\
label l${_NUMBER}r
	menu label ${U_BOOT_MENU_LABEL} ${_VERSION} (rescue target)
	linux ${_BOOT_DIRECTORY}/${_KERNEL}
	${_FALLBACK_INITRD}"

        # Only add fdt line if _FDT is not empty
        if [ -n "${_FDT}" ]; then
            _ENTRY="${_ENTRY}
	${_FDT}"
        fi
        
        # Only add fdtoverlays line if _FDTOVERLAYS is not empty
        if [ -n "${_FDTOVERLAYS}" ]; then
            _ENTRY="${_ENTRY}
	${_FDTOVERLAYS}"
        fi
        
        _ENTRY="${_ENTRY}
	append ${U_BOOT_ROOT} $(echo "${U_BOOT_PARAMETERS}" | sed -e 's| quiet||') systemd.unit=rescue.target
	"
        
        # Add the entry to the configuration
        _CONFIG="${_CONFIG}

${_ENTRY}"
    fi
    
    _NUMBER="$((_NUMBER + 1))"
    
    if [ "${U_BOOT_ENTRIES}" != "all" ] && [ "${_NUMBER}" -ge "${U_BOOT_ENTRIES}" ]; then
        break
    fi
done

# If no valid kernels with matching initramfs found
if [ "${_NUMBER}" -eq 0 ]; then
    echo "E: No valid kernels with matching initramfs found. Cannot create u-boot configuration."
    exit 1
fi

Update "${_U_BOOT_DIRECTORY}/extlinux.conf" "${_CONFIG}"

