#!/bin/sh
# build.sh — build bee ISO
#
# Single build script. Produces a bootable live ISO with SSH access, TUI, NVIDIA drivers.
#
# Run on Alpine 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"
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 KERNEL_PKG_VERSION
export PATH="$PATH:/usr/local/go/bin"

# NOTE: lz4 compression for modloop is disabled — Alpine initramfs may not support lz4 squashfs.
# Default xz compression is used until lz4 support is confirmed.

echo "=== bee ISO build ==="
echo "Alpine: ${ALPINE_VERSION}, Go: ${GO_VERSION}"
echo ""

# --- compile audit binary (static, Linux amd64) ---
# Skip rebuild if binary is newer than all Go source files.
AUDIT_BIN="${DIST_DIR}/bee-audit-linux-amd64"
NEED_BUILD=1
if [ -f "$AUDIT_BIN" ]; then
    NEWEST_SRC=$(find "${REPO_ROOT}/audit" -name '*.go' -newer "$AUDIT_BIN" | head -1)
    [ -z "$NEWEST_SRC" ] && NEED_BUILD=0
fi

if [ "$NEED_BUILD" = "1" ]; then
    echo "=== building audit 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 "$AUDIT_BIN" \
        ./cmd/audit
    echo "binary: $AUDIT_BIN"
    echo "size:   $(du -sh "$AUDIT_BIN" | cut -f1)"
else
    echo "=== audit binary up to date, skipping build ==="
fi

# --- inject authorized_keys for SSH access ---
# Uses the same Ed25519 keys as release signing (from git.mchus.pro/mchus/keys).
# SSH public keys are stored alongside signing keys as ~/.keys/<name>.key.pub
AUTHORIZED_KEYS_FILE="${OVERLAY_DIR}/root/.ssh/authorized_keys"
mkdir -p "${OVERLAY_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
    # auto-collect all developer SSH public keys from ~/.keys/*.key.pub
    > "$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 (user created by bee-sshsetup at boot)"
        echo "  (generate a key with: sh keys/scripts/keygen.sh <your-name>)"
        USE_PASSWORD_FALLBACK=1
    fi
fi

# --- password fallback: write marker file read by init script ---
if [ "${USE_PASSWORD_FALLBACK:-0}" = "1" ]; then
    touch "${OVERLAY_DIR}/etc/bee-ssh-password-fallback"
fi

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

# --- 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_DIR}/usr/local/bin/${tool}"
        chmod +x "${OVERLAY_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 and inject into overlay ---
echo ""
echo "=== building NVIDIA ${NVIDIA_DRIVER_VERSION} modules ==="
sh "${BUILDER_DIR}/build-nvidia-module.sh" "${NVIDIA_DRIVER_VERSION}" "${DIST_DIR}" "${KERNEL_PKG_VERSION}" "${ALPINE_VERSION}"

# Determine kernel version from installed headers
KVER=$(ls /usr/src/ 2>/dev/null | grep '^linux-headers-' | sed 's/linux-headers-//' | sort -V | tail -1)

# Build-time verification: headers must match the repo version we detected.
PINNED_KVER="$(echo "${KERNEL_PKG_VERSION}" | sed 's/-r[0-9]*//')"
RUNNING_KVER="$(echo "${KVER}" | sed 's/-[0-9]*-lts//')"
if [ "${PINNED_KVER}" != "${RUNNING_KVER}" ]; then
    echo "ERROR: kernel version mismatch!"
    echo "  Repo version:      ${KERNEL_PKG_VERSION} (numeric: ${PINNED_KVER})"
    echo "  Installed headers: ${KVER} (numeric: ${RUNNING_KVER})"
    echo "  This should not happen — apk should have installed the repo version."
    exit 1
fi
echo "=== kernel version OK: ${KVER} ==="

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

# Inject .ko files into overlay at /usr/local/lib/nvidia/ (not /lib/modules/ — modloop squashfs
# mounts over that path at boot and makes it read-only, so overlay content there is inaccessible)
OVERLAY_KMOD_DIR="${OVERLAY_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_DIR}/usr/local/bin" "${OVERLAY_DIR}/usr/lib"
cp "${NVIDIA_CACHE}/bin/nvidia-smi" "${OVERLAY_DIR}/usr/local/bin/"
chmod +x "${OVERLAY_DIR}/usr/local/bin/nvidia-smi"
cp "${NVIDIA_CACHE}/bin/nvidia-bug-report.sh" "${OVERLAY_DIR}/usr/local/bin/" 2>/dev/null || true
chmod +x "${OVERLAY_DIR}/usr/local/bin/nvidia-bug-report.sh" 2>/dev/null || true
cp "${NVIDIA_CACHE}/lib/"* "${OVERLAY_DIR}/usr/lib/" 2>/dev/null || true


# --- embed build metadata ---
mkdir -p "${OVERLAY_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_DIR}/etc/bee-release" <<EOF
BEE_ISO_VERSION=${AUDIT_VERSION}
BEE_AUDIT_VERSION=${AUDIT_VERSION}
BUILD_DATE=${BUILD_DATE}
GIT_COMMIT=${GIT_COMMIT}
ALPINE_VERSION=${ALPINE_VERSION}
NVIDIA_DRIVER_VERSION=${NVIDIA_DRIVER_VERSION}
EOF

# --- export build info for genapkovl to inject into motd ---
BUILD_DATE=$(date +%Y-%m-%d)
GIT_COMMIT=$(git -C "${REPO_ROOT}" rev-parse --short HEAD 2>/dev/null || echo "unknown")
export BEE_BUILD_INFO="${BUILD_DATE} git:${GIT_COMMIT} alpine:${ALPINE_VERSION} nvidia:${NVIDIA_DRIVER_VERSION}"

# --- build ISO using mkimage ---
mkdir -p "${DIST_DIR}"
echo ""
echo "=== building ISO ==="

# Install our mkimage profile where mkimage.sh can find it.
# ~/.mkimage is the user plugin directory loaded by mkimage.sh.
# Clear ~/.mkimage to avoid stale profiles from previous builds being picked up
rm -rf "${HOME}/.mkimage"
mkdir -p "${HOME}/.mkimage"
cp "${BUILDER_DIR}/mkimg.bee.sh" "${HOME}/.mkimage/"
cp "${BUILDER_DIR}/genapkovl-bee.sh" "${HOME}/.mkimage/"

# Export overlay dir so the profile script can find it regardless of SRCDIR.
export BEE_OVERLAY_DIR="${OVERLAY_DIR}"

# Clean workdir selectively: remove everything except apks cache so packages aren't re-downloaded.
# mkimage stores each section in a hash-named subdir; apks_* dirs contain downloaded packages.
if [ -d /var/tmp/bee-iso-work ]; then
    find /var/tmp/bee-iso-work -maxdepth 1 -mindepth 1 \
        -not -name 'apks_*' -not -name 'kernel_*' \
        -not -name 'syslinux_*' -not -name 'grub_*' \
        -exec rm -rf {} + 2>/dev/null || true
fi

# Run from /var/tmp: mkimage.sh calls git internally; running from inside /root/bee causes
# "outside repository" errors. /var/tmp is outside the git repo and has enough scratch space.
# genapkovl-bee.sh is found by mkimage via ~/.mkimage/ — no need to copy it to /var/tmp.
export TMPDIR=/var/tmp
cd /var/tmp
sh /usr/share/aports/scripts/mkimage.sh \
    --tag "v${ALPINE_VERSION}" \
    --outdir "${DIST_DIR}" \
    --arch x86_64 \
    --repository "https://dl-cdn.alpinelinux.org/alpine/v${ALPINE_VERSION}/main" \
    --repository "https://dl-cdn.alpinelinux.org/alpine/v${ALPINE_VERSION}/community" \
    --workdir /var/tmp/bee-iso-work \
    --profile bee

ISO="${DIST_DIR}/alpine-bee-${ALPINE_VERSION}-x86_64.iso"
echo ""
echo "=== done ==="
echo "ISO: $ISO"
echo "Size: $(du -sh "$ISO" 2>/dev/null | cut -f1 || echo 'not found')"
echo ""
echo "Boot via BMC virtual media and SSH to the server IP on port 22 as root."