#!/usr/bin/env bash
# Merge enabled Dragon Q8B overlays into the board DTB and point the Boot Loader
# Specification entries at the result. Invoked from kernel-install plugins and
# dragon-q8b-overlay.
set -Eeuo pipefail

dtb_relpath=qcom/sc8280xp-radxa-dragon-q8b.dtb
overlay_dir=/usr/share/dragon-q8b/overlays
overlay_conf=/etc/dragon-q8b/overlays.conf
only_version=
best_effort=0

usage() {
    cat <<'EOF'
Usage: dragon-q8b-refresh-boot [--kernel-version VERSION] [--best-effort]

Merge overlays from /etc/dragon-q8b/overlays.conf into the Q8B DTB with
fdtoverlay and point the Boot Loader Specification entries for the affected
kernels at the merged DTB.

GRUB hands the DTB to the kernel through the "devicetree" key of the BLS
entry, so this script has to place that key where GRUB's blscfg parser will
actually reach it. See set_bls_devicetree below.

Without --best-effort, a missing overlay, missing fdtoverlay, merge
failure, or no eligible Q8B DTB is an error. Kernel-install hooks pass
--best-effort.
EOF
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --kernel-version) only_version=${2:?missing kernel version}; shift 2 ;;
        --best-effort) best_effort=1; shift ;;
        -h|--help) usage; exit 0 ;;
        *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;;
    esac
done

find_base_dtb() {
    local version=$1 candidate
    for candidate in \
        "/usr/lib/modules/$version/dtb/$dtb_relpath" \
        "/lib/modules/$version/dtb/$dtb_relpath" \
        "/boot/dtb-$version/$dtb_relpath"; do
        [[ -f "$candidate" ]] && { printf '%s\n' "$candidate"; return 0; }
    done
    return 1
}

read_overlays() {
    [[ -r "$overlay_conf" ]] || return 0
    awk 'NF && $1 !~ /^#/ {print $1}' "$overlay_conf"
}

merge_dtb() {
    local base=$1 dest=$2
    local -a args=() requested=() overlay name
    mkdir -p "$(dirname "$dest")"
    while IFS= read -r name; do
        [[ -n "$name" ]] || continue
        requested+=("$name")
        overlay="$overlay_dir/${name}.dtbo"
        if [[ ! -f "$overlay" ]]; then
            printf '[dragon-q8b-boot] missing overlay %s\n' "$overlay" >&2
            if [[ "$best_effort" -eq 0 ]]; then
                return 1
            fi
            continue
        fi
        args+=("$overlay")
    done < <(read_overlays)
    if [[ ${#requested[@]} -eq 0 || ${#args[@]} -eq 0 ]]; then
        cp -a "$base" "$dest"
        if [[ ${#requested[@]} -gt 0 && ${#args[@]} -eq 0 ]]; then
            printf '[dragon-q8b-boot] no requested overlays found; copied base DTB\n' >&2
            [[ "$best_effort" -eq 1 ]] || return 1
        fi
        return 0
    fi
    command -v fdtoverlay >/dev/null 2>&1 || {
        printf '[dragon-q8b-boot] fdtoverlay not found\n' >&2
        if [[ "$best_effort" -eq 1 ]]; then
            cp -a "$base" "$dest"
            return 0
        fi
        return 1
    }
    if fdtoverlay -i "$base" -o "$dest" "${args[@]}"; then
        return 0
    fi
    printf '[dragon-q8b-boot] fdtoverlay failed\n' >&2
    if [[ "$best_effort" -eq 1 ]]; then
        cp -a "$base" "$dest"
        return 0
    fi
    rm -f "$dest"
    return 1
}

# Rewrite the devicetree key of one BLS entry so that GRUB's blscfg parser
# really sees it.
#
# Fedora's 20-grub.install writes BLS entries from a template whose last line is
# "${devicetree_entry}". When GRUB_DEVICETREE is unset in /etc/default/grub that
# expands to an empty variable and leaves a blank line at the end of the entry.
# GRUB's blscfg read_entry() stops at the first line it cannot split into a key
# and a value ("if (!separator || separator[1] == '\0') break;"), so a key
# appended after that blank line is dropped without a word and GRUB falls back to
# the firmware-provided device tree. The board then boots with a tree that has no
# TC956x MAC nodes and no QSEE TPM, and the missing Ethernet ports and fTPM look
# like kernel driver bugs instead of a boot configuration problem.
#
# So: drop every existing devicetree line and insert exactly one directly after
# the kernel image line, which is always inside the region GRUB parses.
write_bls_devicetree() {
    local entry=$1 dtb_path=$2 tmp
    tmp=$(mktemp "${entry}.XXXXXX") || return 1
    if ! awk -v dt="devicetree ${dtb_path}" '
            /^[[:space:]]*devicetree[[:space:]]/ { next }
            { line[NR] = $0 }
            END {
                if (NR == 0) { print dt; exit }
                pos = 0
                for (i = 1; i <= NR; i++)
                    if (line[i] ~ /^[[:space:]]*(linux|linux16|efi|image)[[:space:]]/) { pos = i + 1; break }
                if (!pos)
                    for (i = 1; i <= NR; i++)
                        if (line[i] ~ /^[[:space:]]*$/) { pos = i; break }
                if (!pos) pos = NR + 1
                for (i = 1; i <= NR + 1; i++) {
                    if (i == pos) print dt
                    if (i <= NR) print line[i]
                }
            }
        ' "$entry" > "$tmp"; then
        rm -f "$tmp"
        return 1
    fi
    # BLS entries are mode 0600 root; keep whatever the file already had.
    if ! chmod --reference="$entry" "$tmp" 2>/dev/null &&
            ! chmod "$(stat -c '%a' "$entry")" "$tmp" 2>/dev/null; then
        rm -f "$tmp"
        return 1
    fi
    mv -f "$tmp" "$entry"
}

# A key placed past the first blank line is invisible to GRUB; report it here
# rather than leaving the board silently on the firmware device tree.
bls_devicetree_is_visible() {
    local entry=$1
    awk '
        /^[[:space:]]*$/ { blank = NR; next }
        !blank && /^[[:space:]]*devicetree[[:space:]]/ { found = 1; exit }
        END { exit found ? 0 : 1 }
    ' "$entry"
}

set_bls_devicetree() {
    local version=$1
    local relpath entry status
    relpath="/dtb-$version/$dtb_relpath"
    status=0
    shopt -s nullglob
    for entry in \
        /boot/loader/entries/*"${version}"*.conf \
        /boot/efi/loader/entries/*"${version}"*.conf \
        /loader/entries/*"${version}"*.conf; do
        [[ -f "$entry" ]] || continue
        if write_bls_devicetree "$entry" "$relpath" &&
                bls_devicetree_is_visible "$entry"; then
            printf '[dragon-q8b-boot] %s: devicetree %s\n' "$entry" "$relpath" >&2
        else
            printf '[dragon-q8b-boot] could not set a GRUB-visible devicetree in %s\n' "$entry" >&2
            status=1
        fi
    done
    shopt -u nullglob
    if [[ "$status" -ne 0 && "$best_effort" -eq 0 ]]; then
        return 1
    fi
    return 0
}

status=0
processed=0
declare -A seen_dirs
for module_dir in /usr/lib/modules/* /lib/modules/*; do
    [[ -d "$module_dir" ]] || continue
    version=${module_dir##*/}
    if [[ -n "$only_version" && "$version" != "$only_version" ]]; then
        continue
    fi
    real_dir=$(realpath "$module_dir" 2>/dev/null || echo "$module_dir")
    [[ -z "${seen_dirs[$real_dir]:-}" ]] || continue
    seen_dirs["$real_dir"]=1
    [[ -f "$module_dir/vmlinuz" ]] || continue
    base_dtb=$(find_base_dtb "$version" 2>/dev/null) || {
        printf '[dragon-q8b-boot] no Q8B DTB for %s\n' "$version" >&2
        continue
    }
    dest="/boot/dtb-$version/$dtb_relpath"
    if merge_dtb "$base_dtb" "$dest"; then
        processed=$((processed + 1))
        # grubby has no --devtree option, so the BLS entry is the only place the
        # DTB path can live. If this fails the kernel boots on the firmware tree.
        set_bls_devicetree "$version" || status=1
    else
        status=1
        [[ "$best_effort" -eq 1 ]] || continue
    fi
done
if [[ "$best_effort" -eq 0 && "$processed" -eq 0 ]]; then
    printf '[dragon-q8b-boot] no eligible Q8B kernel DTB to refresh\n' >&2
    status=1
fi
exit "$status"
