#!/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}"

# memtest86+ 6.x uses memtest86+.bin (no x64 suffix) for the BIOS binary,
# while 5.x used memtest86+x64.bin. We normalise both to x64 names in the ISO.
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
}

# grub.cfg and live.cfg may not exist yet when binary hooks run — live-build
# creates them after this hook (lb binary_grub-efi / lb binary_syslinux).
# The template already has memtest entries hardcoded, so a missing config file
# here is not an error; validate_iso_memtest() checks the final ISO instead.
warn_only() {
    log "WARNING: $1"
}

copy_memtest_file() {
    src="$1"
    dst_name="${2:-$(basename "$src")}"
    dst="${BINARY_BOOT_DIR}/${dst_name}"

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

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

    log "extracting memtest payload from ${deb}"
    dpkg-deb -x "$deb" "$tmpdir"

    # EFI binary: both 5.x and 6.x use memtest86+x64.efi
    if [ -f "${tmpdir}/boot/memtest86+x64.efi" ]; then
        copy_memtest_file "${tmpdir}/boot/memtest86+x64.efi"
    fi

    # BIOS binary: 5.x = memtest86+x64.bin, 6.x = memtest86+.bin
    if [ -f "${tmpdir}/boot/memtest86+x64.bin" ]; then
        copy_memtest_file "${tmpdir}/boot/memtest86+x64.bin"
    elif [ -f "${tmpdir}/boot/memtest86+.bin" ]; then
        copy_memtest_file "${tmpdir}/boot/memtest86+.bin" "memtest86+x64.bin"
    fi

    rm -rf "$tmpdir"
}

download_and_extract_memtest() {
    tmpdl="$(mktemp -d)"
    if [ -n "${MEMTEST_VERSION:-}" ]; then
        pkg_spec="memtest86+=${MEMTEST_VERSION}"
    else
        pkg_spec="memtest86+"
    fi
    log "downloading ${pkg_spec} from apt"
    if ! ( cd "$tmpdl" && apt-get download "$pkg_spec" 2>/dev/null ); then
        log "apt download failed, retrying after apt-get update"
        apt-get update -qq >/dev/null 2>&1 || true
        ( cd "$tmpdl" && apt-get download "$pkg_spec" 2>/dev/null ) || true
    fi
    deb="$(find "$tmpdl" -maxdepth 1 -type f -name 'memtest86+*.deb' 2>/dev/null | head -1)"
    if [ -n "$deb" ]; then
        extract_memtest_from_deb "$deb"
    else
        log "apt download of memtest86+ failed"
    fi
    rm -rf "$tmpdl"
}

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

    # 1. Try files already placed by lb binary_memtest or chroot
    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
        # 6.x BIOS binary may lack x64 in name — copy with normalised name
        if [ ! -f "${BINARY_BOOT_DIR}/memtest86+x64.bin" ]; then
            copy_memtest_file "${root}/memtest86+.bin" "memtest86+x64.bin" || true
        fi
    done

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

    # 2. Try apt package cache (may be empty if lb binary_memtest already purged)
    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
        [ -f "${BINARY_BOOT_DIR}/${f}" ] || missing=1
    done
    [ "$missing" -eq 1 ] || return 0

    # 3. Fallback: download fresh from apt (lb binary_memtest purges the cache)
    download_and_extract_memtest

    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" ] || {
        warn_only "missing ${GRUB_CFG} (will be created by lb binary_grub-efi from template)"
        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" ] || {
        warn_only "missing ${ISOLINUX_CFG} (will be created by lb binary_syslinux from template)"
        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"