#!/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"
BUILD_WORK_DIR="${DIST_DIR}/live-build-work"
OVERLAY_STAGE_DIR="${DIST_DIR}/overlay-stage"
CACHE_ROOT="${BEE_CACHE_DIR:-${DIST_DIR}/cache}"
AUTH_KEYS=""

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

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

# 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
}

AUDIT_VERSION_EFFECTIVE="$(resolve_audit_version)"
ISO_VERSION_EFFECTIVE="$(resolve_iso_version)"

# 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 ==="
echo "Debian: ${DEBIAN_VERSION}, Kernel ABI: ${DEBIAN_KERNEL_ABI}, Go: ${GO_VERSION}"
echo "Audit version: ${AUDIT_VERSION_EFFECTIVE}, ISO version: ${ISO_VERSION_EFFECTIVE}"
echo ""

echo "=== syncing git submodules ==="
git -C "${REPO_ROOT}" submodule update --init --recursive

# --- compile bee binary (static, Linux amd64) ---
BEE_BIN="${DIST_DIR}/bee-linux-amd64"
GPU_STRESS_BIN="${DIST_DIR}/bee-gpu-stress-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
    echo "=== building bee binary ==="
    cd "${REPO_ROOT}/audit"
    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

echo ""
echo "=== downloading cuBLAS/cuBLASLt/cudart ${NCCL_CUDA_VERSION} userspace ==="
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_STRESS_BIN" ] && [ "${BUILDER_DIR}/bee-gpu-stress.c" -ot "$GPU_STRESS_BIN" ]; then
    GPU_STRESS_NEED_BUILD=0
fi

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

echo "=== preparing staged overlay ==="
# Sync builder config into work dir, preserving lb cache (chroot + packages).
# We do NOT rm -rf BUILD_WORK_DIR so lb can reuse its chroot on repeat builds.
mkdir -p "${BUILD_WORK_DIR}" "${OVERLAY_STAGE_DIR}"
rsync -a --delete \
    --exclude='cache/' \
    --exclude='chroot/' \
    --exclude='.build/' \
    --exclude='*.iso' \
    --exclude='*.packages' \
    --exclude='*.contents' \
    --exclude='*.files' \
    "${BUILDER_DIR}/" "${BUILD_WORK_DIR}/"
# Also persist package cache to CACHE_ROOT so it survives a manual wipe of BUILD_WORK_DIR.
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/bee-gpu-stress" \
    "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-smoketest" \
    "${OVERLAY_STAGE_DIR}/usr/local/bin/all_reduce_perf"

# --- 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"
cp "${GPU_STRESS_BIN}" "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-gpu-stress"
chmod +x "${OVERLAY_STAGE_DIR}/usr/local/bin/bee-gpu-stress"

# --- 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

# --- build NVIDIA kernel modules ---
echo ""
echo "=== building NVIDIA ${NVIDIA_DRIVER_VERSION} modules ==="
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_DIR}/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

# 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 ---
echo ""
echo "=== downloading NCCL ${NCCL_VERSION}+cuda${NCCL_CUDA_VERSION} ==="
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 bee-gpu-stress 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 ---
echo ""
echo "=== building nccl-tests ${NCCL_TESTS_VERSION} ==="
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"
echo "=== all_reduce_perf injected ==="

# --- 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)"
cat > "${OVERLAY_STAGE_DIR}/etc/bee-release" <<EOF
BEE_ISO_VERSION=${ISO_VERSION_EFFECTIVE}
BEE_AUDIT_VERSION=${AUDIT_VERSION_EFFECTIVE}
BUILD_DATE=${BUILD_DATE}
GIT_COMMIT=${GIT_COMMIT}
DEBIAN_VERSION=${DEBIAN_VERSION}
DEBIAN_KERNEL_ABI=${DEBIAN_KERNEL_ABI}
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}
EOF

# Patch motd with build info
BEE_BUILD_INFO="${BUILD_DATE} git:${GIT_COMMIT} debian:${DEBIAN_VERSION} nvidia:${NVIDIA_DRIVER_VERSION}"
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

# --- substitute version placeholders in package list ---
sed -i \
    -e "s/%%DCGM_VERSION%%/${DCGM_VERSION}/g" \
    -e "s/%%ROCM_VERSION%%/${ROCM_VERSION}/g" \
    -e "s/%%ROCM_SMI_VERSION%%/${ROCM_SMI_VERSION}/g" \
    "${BUILD_WORK_DIR}/config/package-lists/bee.list.chroot" \
    "${BUILD_WORK_DIR}/config/archives/rocm.list.chroot"

# --- 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) ==="

cd "${LB_DIR}"
lb clean 2>&1 | tail -3
lb config 2>&1 | tail -5
lb build 2>&1

# live-build outputs live-image-amd64.hybrid.iso in LB_DIR
ISO_RAW="${LB_DIR}/live-image-amd64.hybrid.iso"
ISO_OUT="${DIST_DIR}/bee-debian${DEBIAN_VERSION}-v${ISO_VERSION_EFFECTIVE}-amd64.iso"
if [ -f "$ISO_RAW" ]; then
    cp "$ISO_RAW" "$ISO_OUT"
    echo ""
    echo "=== done ==="
    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."
