#!/bin/sh
# Ensure memtest is present in the final ISO even if live-build's built-in
# memtest stage does not copy the binaries or expose menu entries.
set -e

: "${BEE_REQUIRE_MEMTEST:=0}"

MEMTEST_FILES="memtest86+x64.bin memtest86+x64.efi"
BINARY_BOOT_DIR="binary/boot"
GRUB_CFG="binary/boot/grub/grub.cfg"
ISOLINUX_CFG="binary/isolinux/live.cfg"

log() {
    echo "memtest hook: $*"
}

fail_or_warn() {
    msg="$1"
    if [ "${BEE_REQUIRE_MEMTEST}" = "1" ]; then
        log "ERROR: ${msg}"
        exit 1
    fi
    log "WARNING: ${msg}"
    return 0
}

copy_memtest_file() {
    src="$1"
    base="$(basename "$src")"
    dst="${BINARY_BOOT_DIR}/${base}"

    [ -f "$src" ] || return 1
    mkdir -p "${BINARY_BOOT_DIR}"
    cp "$src" "$dst"
    log "copied ${base} from ${src}"
}

extract_memtest_from_deb() {
    deb="$1"
    tmpdir="$(mktemp -d)"

    log "extracting memtest payload from ${deb}"
    dpkg-deb -x "$deb" "$tmpdir"
    for f in ${MEMTEST_FILES}; do
        if [ -f "${tmpdir}/boot/${f}" ]; then
            copy_memtest_file "${tmpdir}/boot/${f}"
        fi
    done
    rm -rf "$tmpdir"
}

ensure_memtest_binaries() {
    missing=0
    for f in ${MEMTEST_FILES}; do
        [ -f "${BINARY_BOOT_DIR}/${f}" ] || missing=1
    done
    [ "$missing" -eq 1 ] || return 0

    for root in chroot/boot /boot; do
        for f in ${MEMTEST_FILES}; do
            [ -f "${BINARY_BOOT_DIR}/${f}" ] || copy_memtest_file "${root}/${f}" || true
        done
    done

    missing=0
    for f in ${MEMTEST_FILES}; do
        [ -f "${BINARY_BOOT_DIR}/${f}" ] || missing=1
    done
    [ "$missing" -eq 1 ] || return 0

    for root in cache chroot/var/cache/apt/archives /var/cache/apt/archives; do
        [ -d "$root" ] || continue
        deb="$(find "$root" -type f \( -name 'memtest86+_*.deb' -o -name 'memtest86+*.deb' \) 2>/dev/null | head -1)"
        [ -n "$deb" ] || continue
        extract_memtest_from_deb "$deb"
        break
    done

    missing=0
    for f in ${MEMTEST_FILES}; do
        if [ ! -f "${BINARY_BOOT_DIR}/${f}" ]; then
            fail_or_warn "missing ${BINARY_BOOT_DIR}/${f}"
            missing=1
        fi
    done
    [ "$missing" -eq 0 ] || return 0
}

ensure_grub_entry() {
    [ -f "$GRUB_CFG" ] || {
        fail_or_warn "missing ${GRUB_CFG}"
        return 0
    }

    grep -q '### BEE MEMTEST ###' "$GRUB_CFG" && return 0

    cat >> "$GRUB_CFG" <<'EOF'

### BEE MEMTEST ###
if [ "${grub_platform}" = "efi" ]; then
    menuentry "Memory Test (memtest86+)" {
        chainloader /boot/memtest86+x64.efi
    }
else
    menuentry "Memory Test (memtest86+)" {
        linux16 /boot/memtest86+x64.bin
    }
fi
### /BEE MEMTEST ###
EOF

    log "appended memtest entry to ${GRUB_CFG}"
}

ensure_isolinux_entry() {
    [ -f "$ISOLINUX_CFG" ] || {
        fail_or_warn "missing ${ISOLINUX_CFG}"
        return 0
    }

    grep -q '### BEE MEMTEST ###' "$ISOLINUX_CFG" && return 0

    cat >> "$ISOLINUX_CFG" <<'EOF'

# ### BEE MEMTEST ###
label memtest
    menu label ^Memory Test (memtest86+)
    linux /boot/memtest86+x64.bin
# ### /BEE MEMTEST ###
EOF

    log "appended memtest entry to ${ISOLINUX_CFG}"
}

log "ensuring memtest binaries and menu entries in binary image"
ensure_memtest_binaries
ensure_grub_entry
ensure_isolinux_entry
log "memtest assets ready"