#!/bin/sh
# build.sh — internal ISO build entrypoint executed inside the builder container.

set -e

if [ "${BEE_CONTAINER_BUILD:-0}" != "1" ]; then
    echo "build.sh must run inside iso/builder/build-in-container.sh" >&2
    exit 1
fi

REPO_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
BUILDER_DIR="${REPO_ROOT}/iso/builder"
OVERLAY_DIR="${REPO_ROOT}/iso/overlay"
DIST_DIR="${REPO_ROOT}/dist"
VENDOR_DIR="${REPO_ROOT}/iso/vendor"
CACHE_ROOT="${BEE_CACHE_DIR:-${DIST_DIR}/cache}"
AUTH_KEYS=""
BEE_GPU_VENDOR="nvidia"

# parse args
while [ $# -gt 0 ]; do
    case "$1" in
        --authorized-keys) AUTH_KEYS="$2"; shift 2 ;;
        --variant) BEE_GPU_VENDOR="$2"; shift 2 ;;
        *) echo "unknown arg: $1"; exit 1 ;;
    esac
done

case "$BEE_GPU_VENDOR" in
    nvidia|amd|nogpu) ;;
    *) echo "unknown variant: $BEE_GPU_VENDOR (expected nvidia, amd, or nogpu)" >&2; exit 1 ;;
esac

BUILD_WORK_DIR="${DIST_DIR}/live-build-work-${BEE_GPU_VENDOR}"
OVERLAY_STAGE_DIR="${DIST_DIR}/overlay-stage-${BEE_GPU_VENDOR}"

export BEE_GPU_VENDOR

. "${BUILDER_DIR}/VERSIONS"
export PATH="$PATH:/usr/local/go/bin"
: "${BEE_REQUIRE_MEMTEST:=0}"

# Allow git to read the bind-mounted repo (different UID inside container).
git config --global safe.directory "${REPO_ROOT}"
mkdir -p "${DIST_DIR}"
mkdir -p "${CACHE_ROOT}"
: "${GOCACHE:=${CACHE_ROOT}/go-build}"
: "${GOMODCACHE:=${CACHE_ROOT}/go-mod}"
export GOCACHE GOMODCACHE

resolve_audit_version() {
    if [ -n "${BEE_AUDIT_VERSION:-}" ]; then
        echo "${BEE_AUDIT_VERSION}"
        return 0
    fi

    tag="$(git -C "${REPO_ROOT}" describe --tags --match 'audit/v*' --abbrev=7 --dirty 2>/dev/null || true)"
    if [ -z "${tag}" ]; then
        tag="$(git -C "${REPO_ROOT}" describe --tags --match 'v[0-9]*' --abbrev=7 --dirty 2>/dev/null || true)"
    fi
    case "${tag}" in
        audit/v*)
            echo "${tag#audit/v}"
            return 0
            ;;
        v*)
            echo "${tag#v}"
            return 0
            ;;
        "")
            ;;
        *)
            echo "${tag}"
            return 0
            ;;
    esac

    if [ -n "${AUDIT_VERSION:-}" ]; then
        echo "${AUDIT_VERSION}"
        return 0
    fi

    date +%Y%m%d
}

# ISO image versioned separately from the audit binary (iso/v* tags).
resolve_iso_version() {
    if [ -n "${BEE_ISO_VERSION:-}" ]; then
        echo "${BEE_ISO_VERSION}"
        return 0
    fi

    # Plain v* tags (e.g. v2.7) take priority — this is the current tagging scheme
    tag="$(git -C "${REPO_ROOT}" describe --tags --match 'v[0-9]*' --abbrev=7 --dirty 2>/dev/null || true)"
    case "${tag}" in
        v*)
            echo "${tag#v}"
            return 0
            ;;
    esac

    # Legacy iso/v* tags fallback
    tag="$(git -C "${REPO_ROOT}" describe --tags --match 'iso/v*' --abbrev=7 --dirty 2>/dev/null || true)"
    case "${tag}" in
        iso/v*)
            echo "${tag#iso/v}"
            return 0
            ;;
    esac

    # Fall back to audit version so the name is still meaningful
    resolve_audit_version
}

iso_list_files() {
    iso_path="$1"

    if command -v bsdtar >/dev/null 2>&1; then
        bsdtar -tf "$iso_path"
        return $?
    fi

    if command -v xorriso >/dev/null 2>&1; then
        xorriso -indev "$iso_path" -find / -type f -print 2>/dev/null | sed 's#^/##'
        return $?
    fi

    return 127
}

iso_extract_file() {
    iso_path="$1"
    iso_member="$2"

    if command -v bsdtar >/dev/null 2>&1; then
        bsdtar -xOf "$iso_path" "$iso_member"
        return $?
    fi

    if command -v xorriso >/dev/null 2>&1; then
        xorriso -osirrox on -indev "$iso_path" -cat "/$iso_member" 2>/dev/null
        return $?
    fi

    return 127
}

iso_read_file_list() {
    iso_path="$1"
    out_path="$2"

    iso_list_files "$iso_path" > "$out_path" || return 1
    [ -s "$out_path" ] || return 1
    return 0
}

iso_read_member() {
    iso_path="$1"
    iso_member="$2"
    out_path="$3"

    iso_extract_file "$iso_path" "$iso_member" > "$out_path" || return 1
    [ -s "$out_path" ] || return 1
    return 0
}

require_iso_reader() {
    command -v bsdtar >/dev/null 2>&1 && return 0
    command -v xorriso >/dev/null 2>&1 && return 0
    memtest_fail "ISO reader is required for validation/debug (expected bsdtar or xorriso)" "${1:-}"
}

dump_memtest_debug() {
    phase="$1"
    lb_dir="${2:-}"
    iso_path="${3:-}"
    phase_slug="$(printf '%s' "${phase}" | tr ' /' '__')"
    memtest_log="${LOG_DIR:-}/memtest-${phase_slug}.log"

    (
        echo "=== memtest debug: ${phase} ==="

        echo "-- auto/config --"
        if [ -f "${BUILDER_DIR}/auto/config" ]; then
            grep -n -- '--memtest' "${BUILDER_DIR}/auto/config" || echo "  (no --memtest line found)"
        else
            echo "  (missing ${BUILDER_DIR}/auto/config)"
        fi

        echo "-- source bootloader templates --"
        for cfg in \
            "${BUILDER_DIR}/config/bootloaders/grub-pc/grub.cfg" \
            "${BUILDER_DIR}/config/bootloaders/isolinux/live.cfg.in"; do
            if [ -f "$cfg" ]; then
                echo "  file: $cfg"
                grep -n 'Memory Test\|memtest' "$cfg" || echo "    (no memtest lines)"
            fi
        done

        echo "-- source binary hooks --"
        for hook in \
            "${BUILDER_DIR}/config/hooks/normal/9100-memtest.hook.binary"; do
            if [ -f "$hook" ]; then
                echo "  hook: $hook"
            else
                echo "  (missing $hook)"
            fi
        done

        if [ -n "$lb_dir" ] && [ -d "$lb_dir" ]; then
            echo "-- live-build workdir package lists --"
            for pkg in \
                "$lb_dir/config/package-lists/bee.list.chroot" \
                "$lb_dir/config/package-lists/bee-gpu.list.chroot" \
                "$lb_dir/config/package-lists/bee-nvidia.list.chroot"; do
                if [ -f "$pkg" ]; then
                    echo "  file: $pkg"
                    grep -n 'memtest' "$pkg" || echo "    (no memtest lines)"
                fi
            done

            echo "-- live-build chroot/boot --"
            if [ -d "$lb_dir/chroot/boot" ]; then
                find "$lb_dir/chroot/boot" -maxdepth 1 -name 'memtest*' -print | sed 's/^/  /' || true
            else
                echo "  (missing $lb_dir/chroot/boot)"
            fi

            echo "-- live-build binary/boot --"
            if [ -d "$lb_dir/binary/boot" ]; then
                find "$lb_dir/binary/boot" -maxdepth 1 -name 'memtest*' -print | sed 's/^/  /' || true
            else
                echo "  (missing $lb_dir/binary/boot)"
            fi

            echo "-- live-build binary grub cfg --"
            if [ -f "$lb_dir/binary/boot/grub/grub.cfg" ]; then
                grep -n 'Memory Test\|memtest' "$lb_dir/binary/boot/grub/grub.cfg" || echo "  (no memtest lines)"
            else
                echo "  (missing $lb_dir/binary/boot/grub/grub.cfg)"
            fi

            echo "-- live-build binary isolinux cfg --"
            if [ -f "$lb_dir/binary/isolinux/live.cfg" ]; then
                grep -n 'Memory Test\|memtest' "$lb_dir/binary/isolinux/live.cfg" || echo "  (no memtest lines)"
            else
                echo "  (missing $lb_dir/binary/isolinux/live.cfg)"
            fi

            echo "-- live-build package cache --"
            if [ -d "$lb_dir/cache/packages.chroot" ]; then
                find "$lb_dir/cache/packages.chroot" -maxdepth 1 -name 'memtest86+*.deb' -print | sed 's/^/  /' || true
            else
                echo "  (missing $lb_dir/cache/packages.chroot)"
            fi
        fi

        if [ -n "$iso_path" ] && [ -f "$iso_path" ]; then
            iso_files="$(mktemp)"
            iso_grub_cfg="$(mktemp)"
            iso_isolinux_cfg="$(mktemp)"

            echo "-- ISO memtest files --"
            if iso_read_file_list "$iso_path" "$iso_files"; then
                grep 'memtest' "$iso_files" | sed 's/^/  /' || echo "  (no memtest files in ISO)"
            else
                echo "  (failed to list ISO contents)"
            fi

            echo "-- ISO GRUB memtest lines --"
            if iso_read_member "$iso_path" boot/grub/grub.cfg "$iso_grub_cfg"; then
                grep -n 'Memory Test\|memtest' "$iso_grub_cfg" || echo "  (no memtest lines in boot/grub/grub.cfg)"
            else
                echo "  (failed to read boot/grub/grub.cfg from ISO)"
            fi

            echo "-- ISO isolinux memtest lines --"
            if iso_read_member "$iso_path" isolinux/live.cfg "$iso_isolinux_cfg"; then
                grep -n 'Memory Test\|memtest' "$iso_isolinux_cfg" || echo "  (no memtest lines in isolinux/live.cfg)"
            else
                echo "  (failed to read isolinux/live.cfg from ISO)"
            fi

            rm -f "$iso_files" "$iso_grub_cfg" "$iso_isolinux_cfg"
        fi

        echo "=== end memtest debug: ${phase} ==="
    ) | {
        if [ -n "${LOG_DIR:-}" ] && [ -d "${LOG_DIR}" ]; then
            tee "${memtest_log}"
        else
            cat
        fi
    }
}

memtest_fail() {
    msg="$1"
    iso_path="${2:-}"
    level="WARNING"
    if [ "${BEE_REQUIRE_MEMTEST:-0}" = "1" ]; then
        level="ERROR"
    fi
    echo "${level}: ${msg}" >&2
    dump_memtest_debug "failure" "${LB_DIR:-}" "$iso_path" >&2
    if [ "${BEE_REQUIRE_MEMTEST:-0}" = "1" ]; then
        exit 1
    fi
    return 0
}

iso_memtest_present() {
    iso_path="$1"
    iso_files="$(mktemp)"

    [ -f "$iso_path" ] || return 1

    if command -v bsdtar >/dev/null 2>&1; then
        :
    elif command -v xorriso >/dev/null 2>&1; then
        :
    else
        return 2
    fi

    iso_read_file_list "$iso_path" "$iso_files" || {
        rm -f "$iso_files"
        return 2
    }

    grep -q '^boot/memtest86+x64\.bin$' "$iso_files" || {
        rm -f "$iso_files"
        return 1
    }
    grep -q '^boot/memtest86+x64\.efi$' "$iso_files" || {
        rm -f "$iso_files"
        return 1
    }

    grub_cfg="$(mktemp)"
    isolinux_cfg="$(mktemp)"

    iso_read_member "$iso_path" boot/grub/grub.cfg "$grub_cfg" || {
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 2
    }
    iso_read_member "$iso_path" isolinux/live.cfg "$isolinux_cfg" || {
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 2
    }

    grep -q 'Memory Test (memtest86+)' "$grub_cfg" || {
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 1
    }
    grep -q '/boot/memtest86+x64\.efi' "$grub_cfg" || {
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 1
    }
    grep -q '/boot/memtest86+x64\.bin' "$grub_cfg" || {
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 1
    }
    grep -q 'Memory Test (memtest86+)' "$isolinux_cfg" || {
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 1
    }
    grep -q '/boot/memtest86+x64\.bin' "$isolinux_cfg" || {
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 1
    }

    rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
    return 0
}

validate_iso_memtest() {
    iso_path="$1"
    echo "=== validating memtest in ISO ==="

    [ -f "$iso_path" ] || {
        memtest_fail "ISO not found for validation: $iso_path" "$iso_path"
        return 0
    }
    require_iso_reader "$iso_path" || return 0

    iso_files="$(mktemp)"
    iso_read_file_list "$iso_path" "$iso_files" || {
        memtest_fail "failed to list ISO contents while validating memtest" "$iso_path"
        rm -f "$iso_files"
        return 0
    }

    grep -q '^boot/memtest86+x64\.bin$' "$iso_files" || {
        memtest_fail "memtest BIOS binary missing in ISO: boot/memtest86+x64.bin" "$iso_path"
        rm -f "$iso_files"
        return 0
    }
    grep -q '^boot/memtest86+x64\.efi$' "$iso_files" || {
        memtest_fail "memtest EFI binary missing in ISO: boot/memtest86+x64.efi" "$iso_path"
        rm -f "$iso_files"
        return 0
    }

    grub_cfg="$(mktemp)"
    isolinux_cfg="$(mktemp)"

    iso_read_member "$iso_path" boot/grub/grub.cfg "$grub_cfg" || {
        memtest_fail "failed to read boot/grub/grub.cfg from ISO" "$iso_path"
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 0
    }
    iso_read_member "$iso_path" isolinux/live.cfg "$isolinux_cfg" || {
        memtest_fail "failed to read isolinux/live.cfg from ISO" "$iso_path"
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 0
    }

    grep -q 'Memory Test (memtest86+)' "$grub_cfg" || {
        memtest_fail "GRUB menu entry for memtest is missing" "$iso_path"
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 0
    }
    grep -q '/boot/memtest86+x64\.efi' "$grub_cfg" || {
        memtest_fail "GRUB memtest EFI path is missing" "$iso_path"
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 0
    }
    grep -q '/boot/memtest86+x64\.bin' "$grub_cfg" || {
        memtest_fail "GRUB memtest BIOS path is missing" "$iso_path"
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 0
    }
    grep -q 'Memory Test (memtest86+)' "$isolinux_cfg" || {
        memtest_fail "isolinux menu entry for memtest is missing" "$iso_path"
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 0
    }
    grep -q '/boot/memtest86+x64\.bin' "$isolinux_cfg" || {
        memtest_fail "isolinux memtest path is missing" "$iso_path"
        rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
        return 0
    }

    rm -f "$iso_files" "$grub_cfg" "$isolinux_cfg"
    echo "=== memtest validation OK ==="
}

append_memtest_grub_entry() {
    grub_cfg="$1"
    [ -f "$grub_cfg" ] || return 1
    grep -q 'Memory Test (memtest86+)' "$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
}

append_memtest_isolinux_entry() {
    isolinux_cfg="$1"
    [ -f "$isolinux_cfg" ] || return 1
    grep -q 'Memory Test (memtest86+)' "$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
}

copy_memtest_from_deb() {
    deb="$1"
    dst_boot="$2"
    tmpdir="$(mktemp -d)"

    dpkg-deb -x "$deb" "$tmpdir"
    for f in memtest86+x64.bin memtest86+x64.efi; do
        if [ -f "$tmpdir/boot/$f" ]; then
            cp "$tmpdir/boot/$f" "$dst_boot/$f"
        fi
    done
    rm -rf "$tmpdir"
}

reset_live_build_stage() {
    lb_dir="$1"
    stage="$2"

    for root in \
        "$lb_dir/.build" \
        "$lb_dir/.stage" \
        "$lb_dir/auto"; do
        [ -d "$root" ] || continue
        find "$root" -maxdepth 1 \( -name "${stage}" -o -name "${stage}.*" -o -name "*${stage}*" \) -exec rm -rf {} + 2>/dev/null || true
    done
}

recover_iso_memtest() {
    lb_dir="$1"
    iso_path="$2"
    binary_boot="$lb_dir/binary/boot"
    grub_cfg="$lb_dir/binary/boot/grub/grub.cfg"
    isolinux_cfg="$lb_dir/binary/isolinux/live.cfg"

    echo "=== attempting memtest recovery in binary tree ==="

    mkdir -p "$binary_boot"

    for root in \
        "$lb_dir/chroot/boot" \
        "/boot"; do
        for f in memtest86+x64.bin memtest86+x64.efi; do
            if [ ! -f "$binary_boot/$f" ] && [ -f "$root/$f" ]; then
                cp "$root/$f" "$binary_boot/$f"
                echo "memtest recovery: copied $f from $root"
            fi
        done
    done

    if [ ! -f "$binary_boot/memtest86+x64.bin" ] || [ ! -f "$binary_boot/memtest86+x64.efi" ]; then
        for dir in \
            "$lb_dir/cache/packages.binary" \
            "$lb_dir/cache/packages.chroot" \
            "$lb_dir/chroot/var/cache/apt/archives" \
            "${BEE_CACHE_DIR:-${DIST_DIR}/cache}/lb-packages" \
            "/var/cache/apt/archives"; do
            [ -d "$dir" ] || continue
            deb="$(find "$dir" -maxdepth 1 -type f -name 'memtest86+*.deb' 2>/dev/null | head -1)"
            [ -n "$deb" ] || continue
            echo "memtest recovery: extracting payload from $deb"
            copy_memtest_from_deb "$deb" "$binary_boot"
            break
        done
    fi

    if [ ! -f "$binary_boot/memtest86+x64.bin" ] || [ ! -f "$binary_boot/memtest86+x64.efi" ]; then
        tmpdl="$(mktemp -d)"
        if (
            cd "$tmpdl" && apt-get download memtest86+ >/dev/null 2>&1
        ); then
            deb="$(find "$tmpdl" -maxdepth 1 -type f -name 'memtest86+*.deb' 2>/dev/null | head -1)"
            if [ -n "$deb" ]; then
                echo "memtest recovery: downloaded $deb"
                copy_memtest_from_deb "$deb" "$binary_boot"
            fi
        fi
        rm -rf "$tmpdl"
    fi

    if [ -f "$grub_cfg" ]; then
        append_memtest_grub_entry "$grub_cfg" && echo "memtest recovery: ensured GRUB entry"
    else
        echo "memtest recovery: WARNING: missing $grub_cfg"
    fi

    if [ -f "$isolinux_cfg" ]; then
        append_memtest_isolinux_entry "$isolinux_cfg" && echo "memtest recovery: ensured isolinux entry"
    else
        echo "memtest recovery: WARNING: missing $isolinux_cfg"
    fi

    reset_live_build_stage "$lb_dir" "binary_checksums"
    reset_live_build_stage "$lb_dir" "binary_iso"
    reset_live_build_stage "$lb_dir" "binary_zsync"

    run_optional_step_sh "rebuild live-build checksums after memtest recovery" "91-lb-checksums" "lb binary_checksums 2>&1"
    run_optional_step_sh "rebuild ISO after memtest recovery" "92-lb-binary-iso" "rm -f '$iso_path' && lb binary_iso 2>&1"
    run_optional_step_sh "rebuild zsync after memtest recovery" "93-lb-zsync" "lb binary_zsync 2>&1"

    if [ ! -f "$iso_path" ]; then
        memtest_fail "ISO rebuild was skipped or failed after memtest recovery: $iso_path" "$iso_path"
    fi
}

AUDIT_VERSION_EFFECTIVE="$(resolve_audit_version)"
ISO_VERSION_EFFECTIVE="$(resolve_iso_version)"
ISO_BASENAME="easy-bee-${BEE_GPU_VENDOR}-v${ISO_VERSION_EFFECTIVE}-amd64"
# Versioned output directory: dist/easy-bee-v4.1/ — all final artefacts live here.
OUT_DIR="${DIST_DIR}/easy-bee-v${ISO_VERSION_EFFECTIVE}"
mkdir -p "${OUT_DIR}"
LOG_DIR="${OUT_DIR}/${ISO_BASENAME}.logs"
LOG_ARCHIVE="${OUT_DIR}/${ISO_BASENAME}.logs.tar.gz"
ISO_OUT="${OUT_DIR}/${ISO_BASENAME}.iso"
LOG_OUT="${LOG_DIR}/build.log"

cleanup_build_log() {
    status="${1:-$?}"
    trap - EXIT INT TERM HUP

    if [ "${STEP_LOG_ACTIVE:-0}" = "1" ]; then
        cleanup_step_log "${status}" || true
    fi

    if [ "${BUILD_LOG_ACTIVE:-0}" = "1" ]; then
        BUILD_LOG_ACTIVE=0
        exec 1>&3 2>&4
        exec 3>&- 4>&-
        if [ -n "${BUILD_TEE_PID:-}" ]; then
            wait "${BUILD_TEE_PID}" 2>/dev/null || true
        fi
        rm -f "${BUILD_LOG_PIPE}"
    fi

    if [ -n "${LOG_DIR:-}" ] && [ -d "${LOG_DIR}" ] && command -v tar >/dev/null 2>&1; then
        rm -f "${LOG_ARCHIVE}"
        tar -czf "${LOG_ARCHIVE}" -C "$(dirname "${LOG_DIR}")" "$(basename "${LOG_DIR}")" 2>/dev/null || true
        rm -rf "${LOG_DIR}"
    fi

    exit "${status}"
}

start_build_log() {
    command -v tee >/dev/null 2>&1 || {
        echo "ERROR: tee is required for build logging" >&2
        exit 1
    }

    rm -rf "${LOG_DIR}"
    rm -f "${LOG_ARCHIVE}"
    mkdir -p "${LOG_DIR}"
    BUILD_LOG_PIPE="$(mktemp -u "${TMPDIR:-/tmp}/bee-build-log.XXXXXX")"
    mkfifo "${BUILD_LOG_PIPE}"

    exec 3>&1 4>&2
    tee "${LOG_OUT}" < "${BUILD_LOG_PIPE}" &
    BUILD_TEE_PID=$!
    exec > "${BUILD_LOG_PIPE}" 2>&1
    BUILD_LOG_ACTIVE=1

    trap 'cleanup_build_log "$?"' EXIT INT TERM HUP

    echo "=== build log dir: ${LOG_DIR} ==="
    echo "=== build log: ${LOG_OUT} ==="
    echo "=== build log archive: ${LOG_ARCHIVE} ==="
}

cleanup_step_log() {
    status="${1:-$?}"

    if [ "${STEP_LOG_ACTIVE:-0}" = "1" ]; then
        STEP_LOG_ACTIVE=0
        exec 1>&5 2>&6
        exec 5>&- 6>&-
        if [ -n "${STEP_TEE_PID:-}" ]; then
            wait "${STEP_TEE_PID}" 2>/dev/null || true
        fi
        rm -f "${STEP_LOG_PIPE}"
    fi

    return "${status}"
}

run_step() {
    step_name="$1"
    step_slug="$2"
    shift 2

    step_log="${LOG_DIR}/${step_slug}.log"
    echo ""
    echo "=== step: ${step_name} ==="
    echo "=== step log: ${step_log} ==="

    STEP_LOG_PIPE="$(mktemp -u "${TMPDIR:-/tmp}/bee-step-log.XXXXXX")"
    mkfifo "${STEP_LOG_PIPE}"

    exec 5>&1 6>&2
    tee "${step_log}" < "${STEP_LOG_PIPE}" >&5 &
    STEP_TEE_PID=$!
    exec > "${STEP_LOG_PIPE}" 2>&1
    STEP_LOG_ACTIVE=1

    set +e
    "$@"
    step_status=$?
    set -e

    cleanup_step_log "${step_status}"
    if [ "${step_status}" -ne 0 ]; then
        echo "ERROR: step failed: ${step_name} (see ${step_log})" >&2
        exit "${step_status}"
    fi

    echo "=== step OK: ${step_name} ==="
}

run_step_sh() {
    step_name="$1"
    step_slug="$2"
    step_script="$3"

    run_step "${step_name}" "${step_slug}" sh -c "${step_script}"
}

run_optional_step_sh() {
    step_name="$1"
    step_slug="$2"
    step_script="$3"

    if [ "${BEE_REQUIRE_MEMTEST:-0}" = "1" ]; then
        run_step_sh "${step_name}" "${step_slug}" "${step_script}"
        return 0
    fi

    step_log="${LOG_DIR}/${step_slug}.log"
    echo ""
    echo "=== optional step: ${step_name} ==="
    echo "=== optional step log: ${step_log} ==="
    set +e
    sh -c "${step_script}" > "${step_log}" 2>&1
    step_status=$?
    set -e
    cat "${step_log}"
    if [ "${step_status}" -ne 0 ]; then
        echo "WARNING: optional step failed: ${step_name} (see ${step_log})" >&2
    else
        echo "=== optional step OK: ${step_name} ==="
    fi
}

start_build_log

# Auto-detect kernel ABI: refresh apt index, then query current linux-image-amd64 dependency.
# If headers for the detected ABI are not yet installed (kernel updated since image build),
# install them on the fly so NVIDIA modules and ISO kernel always match.
if [ -z "${DEBIAN_KERNEL_ABI}" ] || [ "${DEBIAN_KERNEL_ABI}" = "auto" ]; then
    echo "=== refreshing apt index to detect current kernel ABI ==="
    apt-get update -qq
    DEBIAN_KERNEL_ABI=$(apt-cache depends linux-image-amd64 2>/dev/null \
        | awk '/Depends:.*linux-image-[0-9]/{print $2}' \
        | grep -oE '[0-9]+\.[0-9]+\.[0-9]+-[0-9]+' \
        | head -1)
    if [ -z "${DEBIAN_KERNEL_ABI}" ]; then
        echo "ERROR: could not auto-detect kernel ABI from apt-cache" >&2
        exit 1
    fi
    echo "=== kernel ABI: ${DEBIAN_KERNEL_ABI} ==="
fi

# Export detected ABI so that auto/config can pin the exact kernel package
# (prevents NVIDIA module/kernel mismatch if linux-image-amd64 meta-package
# gets updated between build.sh start and lb build chroot step)
export BEE_KERNEL_ABI="${DEBIAN_KERNEL_ABI}"

KVER="${DEBIAN_KERNEL_ABI}-amd64"
if [ ! -d "/usr/src/linux-headers-${KVER}" ]; then
    echo "=== installing linux-headers-${KVER} (kernel updated since image build) ==="
    apt-get install -y "linux-headers-${KVER}"
fi

echo "=== bee ISO build (variant: ${BEE_GPU_VENDOR}) ==="
echo "Debian: ${DEBIAN_VERSION}, Kernel ABI: ${DEBIAN_KERNEL_ABI}, Go: ${GO_VERSION}"
echo "Audit version: ${AUDIT_VERSION_EFFECTIVE}, ISO version: ${ISO_VERSION_EFFECTIVE}"
echo ""

run_step "sync git submodules" "05-git-submodules" \
    git -C "${REPO_ROOT}" submodule update --init --recursive

# --- compile bee binary (static, Linux amd64) ---
# Shared between variants — built once, reused on second pass.
BEE_BIN="${DIST_DIR}/bee-linux-amd64"
NEED_BUILD=1
if [ -f "$BEE_BIN" ]; then
    NEWEST_SRC=$(find "${REPO_ROOT}/audit" -name '*.go' -newer "$BEE_BIN" | head -1)
    [ -z "$NEWEST_SRC" ] && NEED_BUILD=0
fi

if [ "$NEED_BUILD" = "1" ]; then
    run_step_sh "build bee binary" "10-build-bee" \
        "cd '${REPO_ROOT}/audit' && \
        env GOOS=linux GOARCH=amd64 CGO_ENABLED=0 \
            go build \
            -ldflags '-s -w -X main.Version=${AUDIT_VERSION_EFFECTIVE}' \
            -o '${BEE_BIN}' \
            ./cmd/bee"
    echo "binary: $BEE_BIN"
    if command -v stat >/dev/null 2>&1; then
        BEE_SIZE_BYTES="$(stat -c '%s' "$BEE_BIN" 2>/dev/null || stat -f '%z' "$BEE_BIN")"
    else
        BEE_SIZE_BYTES="$(wc -c < "$BEE_BIN" | tr -d ' ')"
    fi
    if command -v numfmt >/dev/null 2>&1; then
        echo "size:   $(numfmt --to=iec --suffix=B "$BEE_SIZE_BYTES")"
    else
        echo "size:   ${BEE_SIZE_BYTES} bytes"
    fi
else
    echo "=== bee binary up to date, skipping build ==="
fi

# --- NVIDIA-only build steps ---
GPU_BURN_WORKER_BIN="${DIST_DIR}/bee-gpu-burn-worker-linux-amd64"
if [ "$BEE_GPU_VENDOR" = "nvidia" ]; then
    run_step "download cuBLAS/cuBLASLt/cudart ${NCCL_CUDA_VERSION} userspace" "20-cublas" \
        sh "${BUILDER_DIR}/build-cublas.sh" \
        "${CUBLAS_VERSION}" \
        "${CUDA_USERSPACE_VERSION}" \
        "${NCCL_CUDA_VERSION}" \
        "${DIST_DIR}"

    CUBLAS_CACHE="${DIST_DIR}/cublas-${CUBLAS_VERSION}+cuda${NCCL_CUDA_VERSION}"

    GPU_STRESS_NEED_BUILD=1
    if [ -f "$GPU_BURN_WORKER_BIN" ] && [ "${BUILDER_DIR}/bee-gpu-stress.c" -ot "$GPU_BURN_WORKER_BIN" ]; then
        GPU_STRESS_NEED_BUILD=0
    fi

    if [ "$GPU_STRESS_NEED_BUILD" = "1" ]; then
        run_step "build bee-gpu-burn worker" "21-gpu-burn-worker" \
            gcc -O2 -s -Wall -Wextra \
            -I"${CUBLAS_CACHE}/include" \
            -o "$GPU_BURN_WORKER_BIN" \
            "${BUILDER_DIR}/bee-gpu-stress.c" \
            -ldl -lm
        echo "binary: $GPU_BURN_WORKER_BIN"
    else
        echo "=== bee-gpu-burn worker up to date, skipping build ==="
    fi
fi

echo "=== preparing staged overlay (${BEE_GPU_VENDOR}) ==="
mkdir -p "${BUILD_WORK_DIR}" "${OVERLAY_STAGE_DIR}"

# Sync builder config into variant work dir, preserving lb cache.
rsync -a --delete \
    --exclude='cache/' \
    --exclude='chroot/' \
    --exclude='.build/' \
    --exclude='*.iso' \
    --exclude='*.packages' \
    --exclude='*.contents' \
    --exclude='*.files' \
    "${BUILDER_DIR}/" "${BUILD_WORK_DIR}/"

# Share deb package cache across variants.
# Restore: populate work dir cache from shared cache before build.
# Persist: sync back after build (done after lb build below).
LB_PKG_CACHE="${CACHE_ROOT}/lb-packages"
mkdir -p "${LB_PKG_CACHE}"
if [ -d "${BUILD_WORK_DIR}/cache/packages.chroot" ]; then
    rsync -a --delete "${BUILD_WORK_DIR}/cache/packages.chroot/" "${LB_PKG_CACHE}/"
elif [ -d "${LB_PKG_CACHE}" ] && [ "$(ls -A "${LB_PKG_CACHE}" 2>/dev/null)" ]; then
    mkdir -p "${BUILD_WORK_DIR}/cache/packages.chroot"
    rsync -a "${LB_PKG_CACHE}/" "${BUILD_WORK_DIR}/cache/packages.chroot/"
fi

rsync -a "${OVERLAY_DIR}/" "${OVERLAY_STAGE_DIR}/"
rm -f \
    "${OVERLAY_STAGE_DIR}/etc/bee-ssh-password-fallback" \
    "${OVERLAY_STAGE_DIR}/etc/bee-release" \
    "${OVERLAY_STAGE_DIR}/root/.ssh/authorized_keys" \
    "${OVERLAY_STAGE_DIR}/usr/local/bin/bee" \
    "${OVERLAY_STAGE_DIR}/usr/local/bin/john" \
    "${OVERLAY_STAGE_DIR}/usr/local/lib/bee/bee-gpu-burn-worker" \
    "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-smoketest" \
    "${OVERLAY_STAGE_DIR}/usr/local/bin/all_reduce_perf"
rm -rf \
    "${OVERLAY_STAGE_DIR}/usr/local/lib/bee/john"

# Remove NVIDIA-specific overlay files for non-nvidia variants
if [ "$BEE_GPU_VENDOR" != "nvidia" ]; then
    rm -f "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-nvidia-load"
    rm -f "${OVERLAY_STAGE_DIR}/etc/systemd/system/bee-nvidia.service"
fi

# --- inject authorized_keys for SSH access ---
AUTHORIZED_KEYS_FILE="${OVERLAY_STAGE_DIR}/root/.ssh/authorized_keys"
mkdir -p "${OVERLAY_STAGE_DIR}/root/.ssh"

if [ -n "$AUTH_KEYS" ]; then
    cp "$AUTH_KEYS" "$AUTHORIZED_KEYS_FILE"
    chmod 600 "$AUTHORIZED_KEYS_FILE"
    echo "SSH authorized_keys: installed from $AUTH_KEYS"
else
    > "$AUTHORIZED_KEYS_FILE"
    FOUND=0
    for ssh_pub in "$HOME"/.keys/*.key.pub; do
        [ -f "$ssh_pub" ] || continue
        cat "$ssh_pub" >> "$AUTHORIZED_KEYS_FILE"
        echo "SSH: added $(basename "$ssh_pub" .key.pub)"
        FOUND=$((FOUND + 1))
    done
    if [ "$FOUND" -gt 0 ]; then
        chmod 600 "$AUTHORIZED_KEYS_FILE"
        echo "SSH authorized_keys: $FOUND key(s) from ~/.keys/*.key.pub"
    else
        echo "WARNING: no SSH public keys found — falling back to password auth"
        echo "  SSH login: bee / eeb"
        USE_PASSWORD_FALLBACK=1
    fi
fi

if [ "${USE_PASSWORD_FALLBACK:-0}" = "1" ]; then
    touch "${OVERLAY_STAGE_DIR}/etc/bee-ssh-password-fallback"
else
    rm -f "${OVERLAY_STAGE_DIR}/etc/bee-ssh-password-fallback"
fi

# --- copy bee binary into overlay ---
mkdir -p "${OVERLAY_STAGE_DIR}/usr/local/bin"
cp "${DIST_DIR}/bee-linux-amd64" "${OVERLAY_STAGE_DIR}/usr/local/bin/bee"
chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/bee"

if [ "$BEE_GPU_VENDOR" = "nvidia" ] && [ -f "$GPU_BURN_WORKER_BIN" ]; then
    mkdir -p "${OVERLAY_STAGE_DIR}/usr/local/lib/bee" "${OVERLAY_STAGE_DIR}/usr/local/bin"
    cp "${GPU_BURN_WORKER_BIN}" "${OVERLAY_STAGE_DIR}/usr/local/lib/bee/bee-gpu-burn-worker"
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/lib/bee/bee-gpu-burn-worker"
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-gpu-burn" 2>/dev/null || true
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-john-gpu-stress" 2>/dev/null || true
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-nccl-gpu-stress" 2>/dev/null || true
fi

# --- inject smoketest into overlay so it runs directly on the live CD ---
cp "${BUILDER_DIR}/smoketest.sh" "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-smoketest"
chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-smoketest"

# --- vendor utilities (optional pre-fetched binaries) ---
for tool in storcli64 sas2ircu sas3ircu arcconf ssacli; do
    if [ -f "${VENDOR_DIR}/${tool}" ]; then
        cp "${VENDOR_DIR}/${tool}" "${OVERLAY_STAGE_DIR}/usr/local/bin/${tool}"
        chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/${tool}" || true
        echo "vendor tool: ${tool} (included)"
    else
        echo "vendor tool: ${tool} (not found, skipped)"
    fi
done

# --- NVIDIA kernel modules and userspace libs ---
if [ "$BEE_GPU_VENDOR" = "nvidia" ]; then
    run_step "build NVIDIA ${NVIDIA_DRIVER_VERSION} modules" "40-nvidia-module" \
        sh "${BUILDER_DIR}/build-nvidia-module.sh" "${NVIDIA_DRIVER_VERSION}" "${DIST_DIR}" "${DEBIAN_KERNEL_ABI}"

    KVER="${DEBIAN_KERNEL_ABI}-amd64"
    NVIDIA_CACHE="${DIST_DIR}/nvidia-${NVIDIA_DRIVER_VERSION}-${KVER}"

    # Inject .ko files into overlay at /usr/local/lib/nvidia/
    OVERLAY_KMOD_DIR="${OVERLAY_STAGE_DIR}/usr/local/lib/nvidia"
    mkdir -p "${OVERLAY_KMOD_DIR}"
    cp "${NVIDIA_CACHE}/modules/"*.ko "${OVERLAY_KMOD_DIR}/"

    # Inject nvidia-smi and libnvidia-ml
    mkdir -p "${OVERLAY_STAGE_DIR}/usr/local/bin" "${OVERLAY_STAGE_DIR}/usr/lib"
    cp "${NVIDIA_CACHE}/bin/nvidia-smi" "${OVERLAY_STAGE_DIR}/usr/local/bin/"
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/nvidia-smi"
    cp "${NVIDIA_CACHE}/bin/nvidia-bug-report.sh" "${OVERLAY_STAGE_DIR}/usr/local/bin/" 2>/dev/null || true
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/nvidia-bug-report.sh" 2>/dev/null || true
    cp "${NVIDIA_CACHE}/lib/"* "${OVERLAY_STAGE_DIR}/usr/lib/" 2>/dev/null || true
    mkdir -p "${OVERLAY_STAGE_DIR}/etc/OpenCL/vendors"
    printf 'libnvidia-opencl.so.1\n' > "${OVERLAY_STAGE_DIR}/etc/OpenCL/vendors/nvidia.icd"

    # Inject GSP firmware into /lib/firmware/nvidia/<version>/
    if [ -d "${NVIDIA_CACHE}/firmware" ] && [ "$(ls -A "${NVIDIA_CACHE}/firmware" 2>/dev/null)" ]; then
        mkdir -p "${OVERLAY_STAGE_DIR}/lib/firmware/nvidia/${NVIDIA_DRIVER_VERSION}"
        cp "${NVIDIA_CACHE}/firmware/"* "${OVERLAY_STAGE_DIR}/lib/firmware/nvidia/${NVIDIA_DRIVER_VERSION}/"
        echo "=== firmware: $(ls "${OVERLAY_STAGE_DIR}/lib/firmware/nvidia/${NVIDIA_DRIVER_VERSION}/" | wc -l) files injected ==="
    fi

    # --- build / download NCCL ---
    run_step "download NCCL ${NCCL_VERSION}+cuda${NCCL_CUDA_VERSION}" "50-nccl" \
        sh "${BUILDER_DIR}/build-nccl.sh" "${NCCL_VERSION}" "${NCCL_CUDA_VERSION}" "${DIST_DIR}" "${NCCL_SHA256:-}"

    NCCL_CACHE="${DIST_DIR}/nccl-${NCCL_VERSION}+cuda${NCCL_CUDA_VERSION}"

    # Inject libnccl.so.* into overlay alongside other NVIDIA userspace libs
    cp "${NCCL_CACHE}/lib/"* "${OVERLAY_STAGE_DIR}/usr/lib/"
    echo "=== NCCL: $(ls "${NCCL_CACHE}/lib/" | wc -l) files injected into /usr/lib/ ==="

    # Inject cuBLAS/cuBLASLt/cudart runtime libs used by the bee-gpu-burn worker tensor-core GEMM path
    cp "${CUBLAS_CACHE}/lib/"* "${OVERLAY_STAGE_DIR}/usr/lib/"
    echo "=== cuBLAS: $(ls "${CUBLAS_CACHE}/lib/" | wc -l) files injected into /usr/lib/ ==="

    # --- build nccl-tests ---
    run_step "build nccl-tests ${NCCL_TESTS_VERSION}" "60-nccl-tests" \
        sh "${BUILDER_DIR}/build-nccl-tests.sh" \
        "${NCCL_TESTS_VERSION}" \
        "${NCCL_VERSION}" \
        "${NCCL_CUDA_VERSION}" \
        "${DIST_DIR}" \
        "${NVCC_VERSION}" \
        "${DEBIAN_VERSION}"

    NCCL_TESTS_CACHE="${DIST_DIR}/nccl-tests-${NCCL_TESTS_VERSION}"
    cp "${NCCL_TESTS_CACHE}/bin/all_reduce_perf" "${OVERLAY_STAGE_DIR}/usr/local/bin/all_reduce_perf"
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/all_reduce_perf"
    cp "${NCCL_TESTS_CACHE}/lib/"* "${OVERLAY_STAGE_DIR}/usr/lib/" 2>/dev/null || true
    echo "=== all_reduce_perf injected ==="

    run_step "build john jumbo ${JOHN_JUMBO_COMMIT}" "70-john" \
        sh "${BUILDER_DIR}/build-john.sh" "${JOHN_JUMBO_COMMIT}" "${DIST_DIR}"
    JOHN_CACHE="${DIST_DIR}/john-${JOHN_JUMBO_COMMIT}"
    mkdir -p "${OVERLAY_STAGE_DIR}/usr/local/lib/bee/john"
    rsync -a --delete "${JOHN_CACHE}/run/" "${OVERLAY_STAGE_DIR}/usr/local/lib/bee/john/run/"
    ln -sfn ../lib/bee/john/run/john "${OVERLAY_STAGE_DIR}/usr/local/bin/john"
    chmod +x "${OVERLAY_STAGE_DIR}/usr/local/lib/bee/john/run/john"
    echo "=== john injected ==="
fi

# --- embed build metadata ---
mkdir -p "${OVERLAY_STAGE_DIR}/etc"
BUILD_DATE="$(date +%Y-%m-%d)"
GIT_COMMIT="$(git -C "${REPO_ROOT}" rev-parse --short HEAD 2>/dev/null || echo unknown)"

if [ "$BEE_GPU_VENDOR" = "nvidia" ]; then
    GPU_VERSION_LINE="NVIDIA_DRIVER_VERSION=${NVIDIA_DRIVER_VERSION}
NCCL_VERSION=${NCCL_VERSION}
NCCL_CUDA_VERSION=${NCCL_CUDA_VERSION}
CUBLAS_VERSION=${CUBLAS_VERSION}
CUDA_USERSPACE_VERSION=${CUDA_USERSPACE_VERSION}
NCCL_TESTS_VERSION=${NCCL_TESTS_VERSION}
JOHN_JUMBO_COMMIT=${JOHN_JUMBO_COMMIT}"
    GPU_BUILD_INFO="nvidia:${NVIDIA_DRIVER_VERSION}"
elif [ "$BEE_GPU_VENDOR" = "amd" ]; then
    GPU_VERSION_LINE="ROCM_VERSION=${ROCM_VERSION}"
    GPU_BUILD_INFO="rocm:${ROCM_VERSION}"
else
    GPU_VERSION_LINE=""
    GPU_BUILD_INFO="nogpu"
fi

cat > "${OVERLAY_STAGE_DIR}/etc/bee-release" <<EOF
BEE_ISO_VERSION=${ISO_VERSION_EFFECTIVE}
BEE_AUDIT_VERSION=${AUDIT_VERSION_EFFECTIVE}
BEE_GPU_VENDOR=${BEE_GPU_VENDOR}
BUILD_DATE=${BUILD_DATE}
GIT_COMMIT=${GIT_COMMIT}
DEBIAN_VERSION=${DEBIAN_VERSION}
DEBIAN_KERNEL_ABI=${DEBIAN_KERNEL_ABI}
${GPU_VERSION_LINE}
EOF

# Write GPU vendor marker for hooks
echo "${BEE_GPU_VENDOR}" > "${OVERLAY_STAGE_DIR}/etc/bee-gpu-vendor"

# Patch motd with build info
BEE_BUILD_INFO="${BUILD_DATE} git:${GIT_COMMIT} debian:${DEBIAN_VERSION} ${GPU_BUILD_INFO}"
if [ -f "${OVERLAY_STAGE_DIR}/etc/motd" ]; then
    sed "s/%%BUILD_INFO%%/${BEE_BUILD_INFO}/" "${OVERLAY_STAGE_DIR}/etc/motd" \
        > "${OVERLAY_STAGE_DIR}/etc/motd.patched"
    mv "${OVERLAY_STAGE_DIR}/etc/motd.patched" "${OVERLAY_STAGE_DIR}/etc/motd"
fi

# --- copy variant-specific package list, remove all other variant lists ---
# live-build picks up ALL .list.chroot files — delete other variants to avoid conflicts.
cp "${BUILD_WORK_DIR}/config/package-lists/bee-${BEE_GPU_VENDOR}.list.chroot" \
   "${BUILD_WORK_DIR}/config/package-lists/bee-gpu.list.chroot"
rm -f "${BUILD_WORK_DIR}/config/package-lists/bee-nvidia.list.chroot" \
      "${BUILD_WORK_DIR}/config/package-lists/bee-amd.list.chroot" \
      "${BUILD_WORK_DIR}/config/package-lists/bee-nogpu.list.chroot"

# --- remove archives for the other vendor(s) ---
if [ "$BEE_GPU_VENDOR" = "nvidia" ]; then
    rm -f "${BUILD_WORK_DIR}/config/archives/rocm.list.chroot" \
          "${BUILD_WORK_DIR}/config/archives/rocm.key.chroot"
elif [ "$BEE_GPU_VENDOR" = "amd" ]; then
    rm -f "${BUILD_WORK_DIR}/config/archives/nvidia-cuda.list.chroot" \
          "${BUILD_WORK_DIR}/config/archives/nvidia-cuda.key.chroot"
else
    # nogpu: remove both
    rm -f "${BUILD_WORK_DIR}/config/archives/rocm.list.chroot" \
          "${BUILD_WORK_DIR}/config/archives/rocm.key.chroot" \
          "${BUILD_WORK_DIR}/config/archives/nvidia-cuda.list.chroot" \
          "${BUILD_WORK_DIR}/config/archives/nvidia-cuda.key.chroot"
fi

# --- substitute version placeholders in package list and archive ---
if [ "$BEE_GPU_VENDOR" = "nvidia" ]; then
    sed -i \
        -e "s/%%DCGM_VERSION%%/${DCGM_VERSION}/g" \
        "${BUILD_WORK_DIR}/config/package-lists/bee-gpu.list.chroot"
elif [ "$BEE_GPU_VENDOR" = "amd" ]; then
    sed -i \
        -e "s/%%ROCM_VERSION%%/${ROCM_VERSION}/g" \
        -e "s/%%ROCM_SMI_VERSION%%/${ROCM_SMI_VERSION}/g" \
        -e "s/%%ROCM_BANDWIDTH_TEST_VERSION%%/${ROCM_BANDWIDTH_TEST_VERSION}/g" \
        -e "s/%%ROCM_VALIDATION_SUITE_VERSION%%/${ROCM_VALIDATION_SUITE_VERSION}/g" \
        -e "s/%%ROCBLAS_VERSION%%/${ROCBLAS_VERSION}/g" \
        -e "s/%%ROCRAND_VERSION%%/${ROCRAND_VERSION}/g" \
        -e "s/%%HIP_RUNTIME_AMD_VERSION%%/${HIP_RUNTIME_AMD_VERSION}/g" \
        -e "s/%%HIPBLASLT_VERSION%%/${HIPBLASLT_VERSION}/g" \
        -e "s/%%COMGR_VERSION%%/${COMGR_VERSION}/g" \
        "${BUILD_WORK_DIR}/config/package-lists/bee-gpu.list.chroot"
    if [ -f "${BUILD_WORK_DIR}/config/archives/rocm.list.chroot" ]; then
        sed -i \
            -e "s/%%ROCM_VERSION%%/${ROCM_VERSION}/g" \
            "${BUILD_WORK_DIR}/config/archives/rocm.list.chroot"
    fi
fi

# --- sync overlay into live-build includes.chroot ---
LB_DIR="${BUILD_WORK_DIR}"
LB_INCLUDES="${LB_DIR}/config/includes.chroot"
mkdir -p "${LB_INCLUDES}"
rsync -a "${OVERLAY_STAGE_DIR}/" "${LB_INCLUDES}/"

# Ensure SSH authorized_keys perms are correct (rsync may alter)
if [ -f "${LB_INCLUDES}/root/.ssh/authorized_keys" ]; then
    chmod 700 "${LB_INCLUDES}/root/.ssh"
    chmod 600 "${LB_INCLUDES}/root/.ssh/authorized_keys"
fi

# --- build ISO using live-build ---
echo ""
echo "=== building ISO (live-build, variant: ${BEE_GPU_VENDOR}) ==="

# Export for auto/config
BEE_GPU_VENDOR_UPPER="$(echo "${BEE_GPU_VENDOR}" | tr 'a-z' 'A-Z')"
export BEE_GPU_VENDOR_UPPER

cd "${LB_DIR}"
run_step_sh "live-build clean" "80-lb-clean" "lb clean 2>&1 | tail -3"
run_step_sh "live-build config" "81-lb-config" "lb config 2>&1 | tail -5"
dump_memtest_debug "pre-build" "${LB_DIR}"
run_step_sh "live-build build" "90-lb-build" "lb build 2>&1"

# --- persist deb package cache back to shared location ---
# This allows the second variant to reuse all downloaded packages.
if [ -d "${BUILD_WORK_DIR}/cache/packages.chroot" ]; then
    rsync -a "${BUILD_WORK_DIR}/cache/packages.chroot/" "${LB_PKG_CACHE}/"
    echo "=== package cache synced to ${LB_PKG_CACHE} ==="
fi

# live-build outputs live-image-amd64.hybrid.iso in LB_DIR
ISO_RAW="${LB_DIR}/live-image-amd64.hybrid.iso"
if [ -f "$ISO_RAW" ]; then
    dump_memtest_debug "post-build" "${LB_DIR}" "$ISO_RAW"
    if iso_memtest_present "$ISO_RAW"; then
        :
    else
        memtest_status=$?
        if [ "$memtest_status" -eq 1 ]; then
            recover_iso_memtest "${LB_DIR}" "$ISO_RAW"
            dump_memtest_debug "post-recovery" "${LB_DIR}" "$ISO_RAW"
        elif [ "$memtest_status" -eq 2 ]; then
            memtest_fail "failed to inspect ISO for memtest before recovery" "$ISO_RAW"
        fi
    fi
    validate_iso_memtest "$ISO_RAW"
    cp "$ISO_RAW" "$ISO_OUT"
    echo ""
    echo "=== done (${BEE_GPU_VENDOR}) ==="
    echo "ISO: $ISO_OUT"
    if command -v stat >/dev/null 2>&1; then
        ISO_SIZE_BYTES="$(stat -c '%s' "$ISO_OUT" 2>/dev/null || stat -f '%z' "$ISO_OUT")"
    else
        ISO_SIZE_BYTES="$(wc -c < "$ISO_OUT" | tr -d ' ')"
    fi
    if command -v numfmt >/dev/null 2>&1; then
        echo "Size: $(numfmt --to=iec --suffix=B "$ISO_SIZE_BYTES")"
    else
        echo "Size: ${ISO_SIZE_BYTES} bytes"
    fi
else
    echo "ERROR: ISO not found at $ISO_RAW"
    exit 1
fi

echo ""
echo "Boot via BMC virtual media and SSH to the server IP on port 22 as root."