#!/bin/sh
# build.sh — build bee ISO (Debian 12 / live-build)
#
# Single build script. Produces a bootable live ISO with SSH access, TUI, NVIDIA drivers.
#
# Run on Debian 12 builder VM as root after setup-builder.sh.
# Usage:
#   sh iso/builder/build.sh [--authorized-keys /path/to/authorized_keys]

set -e

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

echo "=== bee ISO build ==="
echo "Debian: ${DEBIAN_VERSION}, Kernel ABI: ${DEBIAN_KERNEL_ABI}, Go: ${GO_VERSION}"
echo ""

# --- compile bee binary (static, Linux amd64) ---
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
    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:-$(date +%Y%m%d)}" \
        -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 "=== preparing staged overlay ==="
rm -rf "${BUILD_WORK_DIR}" "${OVERLAY_STAGE_DIR}"
mkdir -p "${BUILD_WORK_DIR}" "${OVERLAY_STAGE_DIR}"
rsync -a "${BUILDER_DIR}/" "${BUILD_WORK_DIR}/"
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-smoketest"

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

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

# --- 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=${AUDIT_VERSION}
BEE_AUDIT_VERSION=${AUDIT_VERSION}
BUILD_DATE=${BUILD_DATE}
GIT_COMMIT=${GIT_COMMIT}
DEBIAN_VERSION=${DEBIAN_VERSION}
DEBIAN_KERNEL_ABI=${DEBIAN_KERNEL_ABI}
NVIDIA_DRIVER_VERSION=${NVIDIA_DRIVER_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

# --- 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${AUDIT_VERSION}-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."