#!/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)"
PROJECT_BUILD_COMMIT="$(git -C "${REPO_ROOT}" rev-parse --short=12 HEAD 2>/dev/null || echo unknown)"
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=""
BUILD_VARIANT="nvidia"
BEE_GPU_VENDOR="nvidia"
BEE_NVIDIA_MODULE_FLAVOR="open"

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

case "$BUILD_VARIANT" in
    nvidia)
        BEE_GPU_VENDOR="nvidia"
        BEE_NVIDIA_MODULE_FLAVOR="open"
        ;;
    nvidia-legacy)
        BEE_GPU_VENDOR="nvidia"
        BEE_NVIDIA_MODULE_FLAVOR="proprietary"
        ;;
    amd)
        BEE_GPU_VENDOR="amd"
        BEE_NVIDIA_MODULE_FLAVOR=""
        ;;
    nogpu)
        BEE_GPU_VENDOR="nogpu"
        BEE_NVIDIA_MODULE_FLAVOR=""
        ;;
    *)
        echo "unknown variant: $BUILD_VARIANT (expected nvidia, nvidia-legacy, amd, or nogpu)" >&2
        exit 1
        ;;
esac

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

export BEE_GPU_VENDOR BEE_NVIDIA_MODULE_FLAVOR BUILD_VARIANT
if [ -n "${BEE_MIRROR:-}" ]; then
    export BEE_MIRROR
fi

. "${BUILDER_DIR}/VERSIONS"
export MEMTEST_VERSION
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}/cache" "${DIST_DIR}/release"
mkdir -p "${CACHE_ROOT}"
: "${GOCACHE:=${CACHE_ROOT}/go-build}"
: "${GOMODCACHE:=${CACHE_ROOT}/go-mod}"
export GOCACHE GOMODCACHE

# Keep the entrypoint focused on build orchestration. These libraries only
# define helpers; stateful logging is loaded after its output paths are known.
. "${BUILDER_DIR}/lib/iso-validation.sh"
. "${BUILDER_DIR}/lib/template.sh"
. "${BUILDER_DIR}/lib/bootloader.sh"
. "${BUILDER_DIR}/lib/fast-path.sh"

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

    if [ -n "${BEE_AUDIT_VERSION:-}" ] && [ -n "${BEE_ISO_VERSION:-}" ] && [ "${BEE_AUDIT_VERSION}" != "${BEE_ISO_VERSION}" ]; then
        echo "ERROR: BEE_AUDIT_VERSION (${BEE_AUDIT_VERSION}) and BEE_ISO_VERSION (${BEE_ISO_VERSION}) differ; versioning must stay synchronized" >&2
        exit 1
    fi

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

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

    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
            ;;
        "")
            ;;
        *)
            echo "${tag}"
            return 0
            ;;
    esac

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

    date +%Y%m%d
}

sync_builder_workdir() {
    src_dir="$1"
    dst_dir="$2"

    mkdir -p "$dst_dir"

    # Historical bug: old workdirs could keep config/bootloaders/grub-pc even
    # after the source tree moved to grub-efi only. Remove bootloaders eagerly
    # so reused workdirs cannot leak stale templates into a new ISO build.
    rm -rf "$dst_dir/config/bootloaders"

    rsync -a --delete \
        --exclude='cache/' \
        --exclude='chroot/' \
        --exclude='binary/' \
        --exclude='.build/' \
        --exclude='.stage/' \
        --exclude='*.iso' \
        --exclude='*.packages' \
        --exclude='*.contents' \
        --exclude='*.files' \
        "$src_dir/" "$dst_dir/"

    if [ ! -f "$dst_dir/config/bootloaders/grub-efi/grub.cfg" ]; then
        echo "ERROR: staged workdir is missing config/bootloaders/grub-efi/grub.cfg" >&2
        exit 1
    fi
    if [ -e "$dst_dir/config/bootloaders/grub-pc" ]; then
        echo "ERROR: stale config/bootloaders/grub-pc remained in staged workdir" >&2
        exit 1
    fi
}


PROJECT_VERSION_EFFECTIVE="$(resolve_project_version)"
if ! validate_project_version "${PROJECT_VERSION_EFFECTIVE}"; then
    echo "ERROR: project version must start and end with an ASCII letter or digit and contain only letters, digits, dots, and hyphens: ${PROJECT_VERSION_EFFECTIVE}" >&2
    exit 1
fi
SQUASHFS_FILENAME="filesystem-v${PROJECT_VERSION_EFFECTIVE}.squashfs"
ISO_BASENAME="easy-bee-${BUILD_VARIANT}-v${PROJECT_VERSION_EFFECTIVE}-amd64"
# Hostname baked into the live image: same product+version scheme as
# ISO_BASENAME/SQUASHFS_FILENAME, but hostname labels don't allow dots.
BEE_HOSTNAME="easy-bee-${BUILD_VARIANT}-v$(printf '%s' "${PROJECT_VERSION_EFFECTIVE}" | tr '.' '-')"
if [ "${#BEE_HOSTNAME}" -gt 63 ]; then
    echo "ERROR: derived hostname exceeds the 63-byte DNS label limit: ${BEE_HOSTNAME}" >&2
    exit 1
fi
# Versioned output directory: dist/easy-bee-v4.1/ - all final artifacts live here.
OUT_DIR="${DIST_DIR}/release/easy-bee-v${PROJECT_VERSION_EFFECTIVE}"
ISO_VERSION_LABEL_TOKEN="$(printf '%s' "${PROJECT_VERSION_EFFECTIVE}" | tr '[:lower:].-' '[:upper:]__')"
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"

. "${BUILDER_DIR}/lib/build-log.sh"

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 || echo "WARNING: apt-get update failed, trying cached index"
    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
        echo "Hint: set DEBIAN_KERNEL_ABI=x.y.z-N in iso/builder/VERSIONS to skip auto-detection" >&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: ${BUILD_VARIANT}) ==="
echo "Debian: ${DEBIAN_VERSION}, Kernel ABI: ${DEBIAN_KERNEL_ABI}, Go: ${GO_VERSION}"
echo "Project version: ${PROJECT_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) ---
# Always rebuild: version/commit ldflags, go.mod/go.sum changes, and deleted Go
# files are not represented reliably by source mtimes.
BEE_BIN="${DIST_DIR}/cache/bee-linux-amd64"
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=${PROJECT_VERSION_EFFECTIVE} -X main.BuildCommit=${PROJECT_BUILD_COMMIT}' \
        -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

# --- NVIDIA-only build steps ---
GPU_BURN_WORKER_BIN="${DIST_DIR}/cache/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}/cache"

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

    echo "=== bee-gpu-burn FP4 header probe ==="
    fp4_type_match="$(grep -Rsnm 1 'CUDA_R_4F_E2M1' "${CUBLAS_CACHE}/include" 2>/dev/null || true)"
    fp4_scale_match="$(grep -Rsnm 1 'CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3' "${CUBLAS_CACHE}/include" 2>/dev/null || true)"
    if [ -n "$fp4_type_match" ]; then
        echo "fp4_header_symbol=present"
        echo "$fp4_type_match"
    else
        echo "fp4_header_symbol=missing"
    fi
    if [ -n "$fp4_scale_match" ]; then
        echo "fp4_scale_mode_symbol=present"
        echo "$fp4_scale_match"
    else
        echo "fp4_scale_mode_symbol=missing"
    fi

    GPU_STRESS_NEED_BUILD=1
    if [ -f "$GPU_BURN_WORKER_BIN" ]; then
        GPU_STRESS_NEED_BUILD=0
        for dep in \
            "${BUILDER_DIR}/bee-gpu-stress.c" \
            "${BUILDER_DIR}/bee-gpu-stress-cuda.inc" \
            "${BUILDER_DIR}/bee-gpu-stress-cublaslt.inc" \
            "${BUILDER_DIR}/bee-gpu-stress-main.inc" \
            "${BUILDER_DIR}/VERSIONS"; do
            if [ "$dep" -nt "$GPU_BURN_WORKER_BIN" ]; then
                GPU_STRESS_NEED_BUILD=1
                break
            fi
        done
        if [ "$GPU_STRESS_NEED_BUILD" = "0" ] && \
            find "${CUBLAS_CACHE}/include" "${CUBLAS_CACHE}/lib" -type f -newer "$GPU_BURN_WORKER_BIN" | grep -q .; then
            GPU_STRESS_NEED_BUILD=1
        fi
    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
    echo "=== bee-gpu-burn compiled profile probe ==="
    if grep -aq 'fp4_e2m1' "$GPU_BURN_WORKER_BIN"; then
        echo "fp4_profile_string=present"
    else
        echo "fp4_profile_string=missing"
    fi
fi

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

# Sync builder config into variant work dir, preserving lb cache.
sync_builder_workdir "${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 --delete "${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"
    rm -f "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-gpu-burn"
    rm -f "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-john-gpu-stress"
    rm -f "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-nccl-gpu-stress"
    rm -f "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-nvidia-recover"
    rm -f "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-dcgmproftester-staggered"
    rm -f "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-check-nvswitch"
    rm -rf "${OVERLAY_STAGE_DIR}/etc/systemd/system/nvidia-fabricmanager.service.d"
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 "$BEE_BIN" "${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 storcli2 sas2ircu sas3ircu arcconf ssacli saa; 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

# saa companion directories - saa searches for these relative to CWD (/usr/local/bin)
for saa_subdir in acpica_bin ExternalData tool stunnel GO_SNMP; do
    if [ -d "${VENDOR_DIR}/${saa_subdir}" ]; then
        cp -r "${VENDOR_DIR}/${saa_subdir}" "${OVERLAY_STAGE_DIR}/usr/local/bin/"
        find "${OVERLAY_STAGE_DIR}/usr/local/bin/${saa_subdir}" -type f -exec chmod +x {} \; 2>/dev/null || true
        echo "vendor saa: ${saa_subdir}/ (included)"
    else
        echo "vendor saa: ${saa_subdir}/ (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}/cache" "${DEBIAN_KERNEL_ABI}" "${BEE_NVIDIA_MODULE_FLAVOR}"

    KVER="${DEBIAN_KERNEL_ABI}-amd64"
    NVIDIA_CACHE="${DIST_DIR}/cache/nvidia-${BEE_NVIDIA_MODULE_FLAVOR}-${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 /usr/lib/firmware/nvidia/<version>/.
    #
    # It MUST be staged under usr/lib, never a bare lib/. On a merged-usr Debian
    # root /lib is a symlink to usr/lib; staging a real lib/ directory here makes
    # `rsync -a` (fast-path repack) replace that symlink with a plain directory,
    # which orphans /lib/systemd/systemd, /lib/x86_64-linux-gnu/* and friends and
    # panics the boot with "run-init: can't execute '/sbin/init'".
    if [ -d "${NVIDIA_CACHE}/firmware" ] && [ "$(ls -A "${NVIDIA_CACHE}/firmware" 2>/dev/null)" ]; then
        mkdir -p "${OVERLAY_STAGE_DIR}/usr/lib/firmware/nvidia/${NVIDIA_DRIVER_VERSION}"
        cp "${NVIDIA_CACHE}/firmware/"* "${OVERLAY_STAGE_DIR}/usr/lib/firmware/nvidia/${NVIDIA_DRIVER_VERSION}/"
        echo "=== firmware: $(ls "${OVERLAY_STAGE_DIR}/usr/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}/cache" "${NCCL_SHA256:-}"

    NCCL_CACHE="${DIST_DIR}/cache/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}/cache" \
        "${NVCC_VERSION}" \
        "${DEBIAN_VERSION}"

    NCCL_TESTS_CACHE="${DIST_DIR}/cache/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}/cache"
    JOHN_CACHE="${DIST_DIR}/cache/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}
NVIDIA_KERNEL_MODULES_FLAVOR=${BEE_NVIDIA_MODULE_FLAVOR}
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-${BEE_NVIDIA_MODULE_FLAVOR}:${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_VERSION=${PROJECT_VERSION_EFFECTIVE}
export BEE_VERSION
BEE_ISO_VERSION=${PROJECT_VERSION_EFFECTIVE}
BEE_AUDIT_VERSION=${PROJECT_VERSION_EFFECTIVE}
BEE_BUILD_VARIANT=${BUILD_VARIANT}
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"
if [ "$BEE_GPU_VENDOR" = "nvidia" ]; then
    echo "${BEE_NVIDIA_MODULE_FLAVOR}" > "${OVERLAY_STAGE_DIR}/etc/bee-nvidia-modules-flavor"
else
    rm -f "${OVERLAY_STAGE_DIR}/etc/bee-nvidia-modules-flavor"
fi

# 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
    replace_literal_in_file "${OVERLAY_STAGE_DIR}/etc/motd" "%%BUILD_INFO%%" "${BEE_BUILD_INFO}"
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
    replace_literal_in_file "${BUILD_WORK_DIR}/config/package-lists/bee-gpu.list.chroot" \
        "%%NVIDIA_FABRICMANAGER_VERSION%%" "${NVIDIA_FABRICMANAGER_VERSION}"
    replace_literal_in_file "${BUILD_WORK_DIR}/config/package-lists/bee-gpu.list.chroot" \
        "%%DCGM_VERSION%%" "${DCGM_VERSION}"
elif [ "$BEE_GPU_VENDOR" = "amd" ]; then
    gpu_package_list="${BUILD_WORK_DIR}/config/package-lists/bee-gpu.list.chroot"
    replace_literal_in_file "$gpu_package_list" "%%ROCM_VERSION%%" "${ROCM_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%ROCM_SMI_VERSION%%" "${ROCM_SMI_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%ROCM_BANDWIDTH_TEST_VERSION%%" "${ROCM_BANDWIDTH_TEST_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%ROCM_VALIDATION_SUITE_VERSION%%" "${ROCM_VALIDATION_SUITE_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%ROCBLAS_VERSION%%" "${ROCBLAS_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%ROCRAND_VERSION%%" "${ROCRAND_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%HIP_RUNTIME_AMD_VERSION%%" "${HIP_RUNTIME_AMD_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%HIPBLASLT_VERSION%%" "${HIPBLASLT_VERSION}"
    replace_literal_in_file "$gpu_package_list" "%%COMGR_VERSION%%" "${COMGR_VERSION}"
    if [ -f "${BUILD_WORK_DIR}/config/archives/rocm.list.chroot" ]; then
        replace_literal_in_file "${BUILD_WORK_DIR}/config/archives/rocm.list.chroot" \
            "%%ROCM_VERSION%%" "${ROCM_VERSION}"
    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 --delete "${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

# Export the release-specific outer ISO identity before either build path.
# The fast path regenerates bootloader assets and the ISO container too.
BEE_GPU_VENDOR_UPPER="$(echo "${BUILD_VARIANT}" | tr 'a-z-' 'A-Z_')"
_vol_prefix="EASY_BEE_${BEE_GPU_VENDOR_UPPER}_V"
_max_token=$(( 32 - ${#_vol_prefix} ))
_vol_token="$(printf '%s' "${ISO_VERSION_LABEL_TOKEN}" | cut -c1-${_max_token})"
BEE_ISO_VOLUME="${_vol_prefix}${_vol_token}"
unset _vol_prefix _max_token _vol_token
export BEE_GPU_VENDOR_UPPER BEE_ISO_VOLUME BEE_HOSTNAME

# --- auto fast-path: squashfs surgery if only light files changed ---
if ! needs_full_build; then
    echo "=== fast-path build (no heavy config changes since last full build) ==="
    fast_path_repack_squashfs
    # The reusable squashfs now reflects this overlay revision. Advance the
    # manifest immediately, even if outer-ISO assembly later fails, so a file
    # added by this build and removed before a retry cannot linger unnoticed.
    write_overlay_manifest "${FULL_BUILD_OVERLAY_MANIFEST}.new"
    mv "${FULL_BUILD_OVERLAY_MANIFEST}.new" "${FULL_BUILD_OVERLAY_MANIFEST}"
    cd "${LB_DIR}"
    run_step_sh "fast-path live-build config" "81-fast-lb-config" "lb config 2>&1 | tail -5"
    echo "=== fast-path: enforcing canonical bootloader assets ==="
    enforce_live_build_bootloader_assets "${LB_DIR}"
    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_step_sh "fast-path rebuild checksums" "82-fast-lb-checksums" "lb binary_checksums 2>&1"
    run_step_sh "fast-path rebuild ISO" "83-fast-lb-binary-iso" "lb binary_iso 2>&1"
    run_step_sh "fast-path rebuild zsync" "84-fast-lb-zsync" "lb binary_zsync 2>&1"
    ISO_RAW="${LB_DIR}/live-image-amd64.hybrid.iso"
    validate_iso_memtest "$ISO_RAW"
    validate_iso_live_boot_entries "$ISO_RAW"
    validate_iso_grub_assets "$ISO_RAW"
    validate_iso_nvidia_runtime "$ISO_RAW"
    validate_iso_rootfs_layout "$ISO_RAW"
    cp "$ISO_RAW" "$ISO_OUT"
    echo ""
    echo "=== done (${BUILD_VARIANT}, fast-path) ==="
    echo "ISO: $ISO_OUT"
    exit 0
fi

# --- build ISO using live-build ---
echo ""
echo "=== building ISO (variant: ${BUILD_VARIANT}) ==="
rm -f "${FULL_BUILD_MARKER}"

cd "${LB_DIR}"
run_step_sh "live-build clean" "80-lb-clean" "lb clean --all 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}"
export MKSQUASHFS_OPTIONS="-no-xattrs"
run_step_sh "live-build build" "90-lb-build" "lb build 2>&1"
echo "=== enforcing canonical bootloader assets ==="
enforce_live_build_bootloader_assets "${LB_DIR}"
# Rename lb's default filesystem.squashfs to the versioned filename so the
# ISO contains a version-stamped squashfs (e.g. filesystem-v10.15.squashfs).
_std_sq="${LB_DIR}/binary/live/filesystem.squashfs"
_ver_sq="${LB_DIR}/binary/live/${SQUASHFS_FILENAME}"
if [ -f "${_std_sq}" ] && [ "${_std_sq}" != "${_ver_sq}" ]; then
    mv "${_std_sq}" "${_ver_sq}"
    echo "=== squashfs renamed: filesystem.squashfs -> ${SQUASHFS_FILENAME} ==="
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_step_sh "rebuild live-build checksums after bootloader sync" "91b-lb-checksums" "lb binary_checksums 2>&1"
run_step_sh "rebuild ISO after bootloader sync" "91c-lb-binary-iso" "lb binary_iso 2>&1"
run_step_sh "rebuild zsync after bootloader sync" "91d-lb-zsync" "lb binary_zsync 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"
    validate_iso_live_boot_entries "$ISO_RAW"
    validate_iso_grub_assets "$ISO_RAW"
    validate_iso_nvidia_runtime "$ISO_RAW"
    validate_iso_rootfs_layout "$ISO_RAW"
    cp "$ISO_RAW" "$ISO_OUT"
    hash_heavy_config > "${FULL_BUILD_HASH_FILE}.new"
    printf '%s\n' "${DEBIAN_KERNEL_ABI}" > "${FULL_BUILD_ABI_FILE}.new"
    write_overlay_manifest "${FULL_BUILD_OVERLAY_MANIFEST}.new"
    mv "${FULL_BUILD_HASH_FILE}.new" "${FULL_BUILD_HASH_FILE}"
    mv "${FULL_BUILD_ABI_FILE}.new" "${FULL_BUILD_ABI_FILE}"
    mv "${FULL_BUILD_OVERLAY_MANIFEST}.new" "${FULL_BUILD_OVERLAY_MANIFEST}"
    touch "${FULL_BUILD_MARKER}"
    echo ""
    echo "=== done (${BUILD_VARIANT}) ==="
    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."