#!/bin/bash
# void-mklive postsetup hook (-x flag).
#
# Runs AFTER copy_include_directories (overlay applied) but BEFORE
# generate_initramfs (boot initramfs).  The xbps initramfs-regenerate trigger
# fires during install_packages — BEFORE the overlay — so the squashfs rootfs's
# own /boot/initramfs-KVER.img is missing our etc/dracut.conf.d/10-nvidia.conf.
#
# This script re-runs dracut inside the rootfs chroot for every installed kernel
# so the squashfs initramfs:
#   • includes nvidia.ko (add_drivers from 10-nvidia.conf)
#   • excludes nouveau.ko (omit_drivers from 10-nvidia.conf)
#
# $1 = absolute path to the rootfs directory being built.

set -euo pipefail

ROOTFS="$1"

if [[ ! -d "$ROOTFS/usr/lib/modules" ]]; then
    echo "[postsetup-nvidia] no modules dir — skipping"
    exit 0
fi

if [[ ! -x "$ROOTFS/usr/bin/dracut" ]]; then
    echo "[postsetup-nvidia] dracut not in rootfs — skipping"
    exit 0
fi

if [[ ! -f "$ROOTFS/etc/dracut.conf.d/10-nvidia.conf" ]]; then
    echo "[postsetup-nvidia] WARNING: 10-nvidia.conf not found in rootfs overlay — skipping"
    exit 0
fi

# Mount pseudo-filesystems needed by dracut inside the chroot.
_umount_pseudo() {
    for f in dev proc sys; do
        [[ -d "$ROOTFS/$f" ]] && \
            { umount -R -f "$ROOTFS/$f" 2>/dev/null || umount -R -l "$ROOTFS/$f" 2>/dev/null || true; }
    done
}
trap _umount_pseudo EXIT

for f in dev proc sys; do
    mkdir -p "$ROOTFS/$f"
    mount --rbind "/$f" "$ROOTFS/$f"
done

echo "[postsetup-nvidia] regenerating initramfs with nvidia config"
echo "[postsetup-nvidia] 10-nvidia.conf:"
sed 's/^/    /' "$ROOTFS/etc/dracut.conf.d/10-nvidia.conf"

found=0
for kdir in "$ROOTFS/usr/lib/modules"/*/; do
    [[ -d "${kdir}/kernel" ]] || continue
    kver="${kdir##*/}"
    found=1

    # Verify nvidia.ko is present for this kernel version.
    nv_ko=$(find "$ROOTFS/usr/lib/modules/$kver" -name 'nvidia.ko' -o -name 'nvidia.ko.zst' 2>/dev/null | head -1)
    if [[ -z "$nv_ko" ]]; then
        echo "[postsetup-nvidia] WARNING: nvidia.ko not found for kernel $kver — initramfs regeneration skipped"
        continue
    fi
    echo "[postsetup-nvidia] kernel $kver: nvidia.ko at ${nv_ko#"$ROOTFS"}"

    echo "[postsetup-nvidia] running dracut --force for $kver ..."
    chroot "$ROOTFS" env -i PATH=/usr/sbin:/usr/bin:/sbin:/bin \
        /usr/bin/dracut --force \
        "/boot/initramfs-${kver}.img" \
        "$kver" 2>&1 | grep -v '^$' | tail -5 || {
        echo "[postsetup-nvidia] WARNING: dracut failed for $kver — continuing"
        continue
    }

    sz=$(ls -lh "$ROOTFS/boot/initramfs-${kver}.img" 2>/dev/null | awk '{print $5}')
    echo "[postsetup-nvidia] done: /boot/initramfs-${kver}.img ($sz)"
done

[[ "$found" -eq 1 ]] || echo "[postsetup-nvidia] WARNING: no kernels found in rootfs"
echo "[postsetup-nvidia] complete"