#!/bin/sh
# lib/squashfs-layers.sh - deterministic semantic SquashFS layering for the
# bee live medium.
#
# Module: squashfs-layers
# Version: 1.0
#
# Why this exists
# ---------------
# The live medium used to ship one ~2.8 GB filesystem-v<ver>.squashfs. Booting
# through a BMC / IPMI virtual CD reads that single file sequentially during the
# live-boot "toram" copy; if the redirected medium drops off the bus part-way
# through, the whole copy is lost and the boot fails. Splitting the rootfs into
# several self-contained squashfs layers means a mid-copy failure only costs the
# layer in flight, and the retry loop re-reads far less data.
#
# This is a resilience / reduced-re-read mechanism. It is NOT a fix for the
# underlying BMC virtual-media instability.
#
# Layer model (NVIDIA variants)
# -----------------------------
#   00-base            Debian rootfs, kernel + modules, systemd, Bee, networking,
#                      CLI diagnostic tools, dpkg database. Boot-critical.
#   05-firmware        Device firmware (firmware-* packages: NIC, wifi, non-NVIDIA
#                      GPU). Not NVIDIA GSP firmware - that rides with the driver.
#   08-desktop         Local-console GUI stack: X.org, lightdm, openbox, mesa,
#                      GTK, chromium, mupdf, fonts. SSH / headless use never
#                      touches this layer.
#   10-nvidia-driver   NVIDIA kernel modules (.ko), driver userspace libraries
#                      (libnvidia-*, libcuda.so*), nvidia-smi, GSP firmware,
#                      OpenCL ICD, modprobe config, alternatives, bee-nvidia-*.
#   20-nvidia-platform Fabric Manager, libnvidia-nscq, nvlsm, DCGM core and the
#                      non-CUDA proprietary DCGM components plus their units.
#   30-nvidia-cuda-libs CUDA userspace runtime (cuBLAS/cuBLASLt/cudart), NCCL,
#                      nccl-tests, the bee GPU stress worker assets.
#   40-nvidia-dcgm-cuda The CUDA-linked DCGM components (dcgmproftester CUDA
#                      kernels) - split from 30 because they are large.
#
# AMD / nogpu: a single 00-base squashfs, no filesystem.module.
#
# Classification is by dpkg file ownership (var/lib/dpkg/info/*.list), not by
# path substring. Files that build.sh injects directly from the build cache and
# the project overlay (and therefore belong to no .deb) are classified by the
# explicit rules in bee_layer_for_injected_path. A path is never moved to an
# NVIDIA layer merely because it contains the string "nvidia".
#
# Ordering
# --------
# Layer names sort lexically 00 < 10 < 20 < 30 < 40. live-boot reads
# live/filesystem.module (written by bee_layer_write_module_file) verbatim and
# stacks the images so the LAST listed image has the HIGHEST OverlayFS priority.
# bee-install unpacks the same list in order with `unsquashfs -f` (last write
# wins) and the fast-path reconstruction extracts in the same order. A
# higher-numbered layer therefore always wins a conflict, everywhere.

# Compressed-size budget per layer. Soft target from the design; hard ceiling
# fails the build so a layer cannot silently grow back toward the monolith.
BEE_LAYER_TARGET_MIB="${BEE_LAYER_TARGET_MIB:-700}"
BEE_LAYER_MAX_MIB="${BEE_LAYER_MAX_MIB:-800}"

# Deterministic mksquashfs options shared by the monolith and every layer.
BEE_LAYER_MKSQUASHFS_OPTS="-comp zstd -b 1048576 -noappend -no-progress -no-xattrs -processors 1"

bee_layer_slugs_for_variant() {
    # echo the ordered layer slugs for a build variant, one per line.
    case "$1" in
        nvidia|nvidia-legacy)
            printf '%s\n' \
                00-base \
                05-firmware \
                08-desktop \
                10-nvidia-driver \
                20-nvidia-platform \
                30-nvidia-cuda-libs \
                40-nvidia-dcgm-cuda
            ;;
        amd|nogpu)
            printf '%s\n' 00-base
            ;;
        *)
            echo "bee_layer: unknown variant: $1" >&2
            return 1
            ;;
    esac
}

# Map an installed dpkg package name to a layer slug, or empty for "base".
# case is first-match: NVIDIA names are matched before the generic firmware-* and
# desktop families, so firmware-nvidia-* rides with the driver, not 05-firmware.
bee_layer_for_dpkg_pkg() {
    case "$1" in
        nvidia-fabricmanager|libnvidia-nscq|nvlsm|libibumad3|libibumad[0-9]*|\
        datacenter-gpu-manager-4-core|datacenter-gpu-manager-4-proprietary)
            echo 20-nvidia-platform ;;
        datacenter-gpu-manager-4-cuda[0-9]*|\
        datacenter-gpu-manager-4-proprietary-cuda[0-9]*)
            echo 40-nvidia-dcgm-cuda ;;
        nvidia-modprobe|nvidia-kernel-common|nvidia-kernel-support|\
        nvidia-installer-cleanup|glx-alternative-nvidia|nvidia-tesla-*|\
        firmware-nvidia-*|libnvidia-*|\
        nvtop|clinfo|ocl-icd-libopencl1)
            echo 10-nvidia-driver ;;
        firmware-*)
            echo 05-firmware ;;
        xserver-xorg*|xserver-common|xorg|xorg-*|xfonts-*|xinit|\
        x11-apps|x11-utils|x11-xserver-utils|x11-xkb-utils|x11-session-utils|x11-common|\
        lightdm|lightdm-*|liblightdm-*|openbox|obconf|feh|scrot|\
        chromium|chromium-*|mupdf|mupdf-*|\
        libllvm[0-9]*|libgl1-mesa-dri|libglx-mesa0|libglapi-mesa|libegl-mesa0|\
        libglu1-mesa|glx-alternative-mesa|libgbm1|\
        libgtk-3-*|libgtk2.0-*|libgtkmm-3.0-*|libgtksourceview-*|\
        libpango-*|libpangomm-*|libpangocairo-*|libpangoft2-*|libpangoxft-*|\
        libcairo2|libcairo-gobject2|libcairomm-*|libgdk-pixbuf-*|libgdk-pixbuf2.0-*|\
        fonts-*)
            echo 08-desktop ;;
        *)
            echo "" ;;
    esac
}

# Emit "<extended-regex><TAB><layer-slug>" lines for files that build.sh injects
# directly (no owning .deb). Matched against the merged-usr-canonical relative
# path (no leading slash). These override dpkg ownership.
bee_layer_injected_rules() {
    cat <<'RULES'
^usr/local/lib/nvidia/[^/]+\.ko$	10-nvidia-driver
^usr/local/bin/nvidia-smi$	10-nvidia-driver
^usr/local/bin/nvidia-bug-report\.sh$	10-nvidia-driver
^usr/lib/libcuda\.so(\..*)?$	10-nvidia-driver
^usr/lib/libnvidia-[^/]+$	10-nvidia-driver
^usr/lib/firmware/nvidia/	10-nvidia-driver
^etc/OpenCL/vendors/nvidia\.icd$	10-nvidia-driver
^usr/local/bin/bee-nvidia-load$	10-nvidia-driver
^usr/local/bin/bee-nvidia-recover$	10-nvidia-driver
^usr/local/bin/bee-check-nvswitch$	10-nvidia-driver
^etc/systemd/system/bee-nvidia\.service$	10-nvidia-driver
^etc/systemd/system/nvidia-fabricmanager\.service\.d/	20-nvidia-platform
^usr/lib/libnccl\.so(\..*)?$	30-nvidia-cuda-libs
^usr/lib/libcublas\.so(\..*)?$	30-nvidia-cuda-libs
^usr/lib/libcublasLt\.so(\..*)?$	30-nvidia-cuda-libs
^usr/lib/libcudart\.so(\..*)?$	30-nvidia-cuda-libs
^usr/local/bin/all_reduce_perf$	30-nvidia-cuda-libs
^usr/local/lib/bee/bee-gpu-burn-worker$	30-nvidia-cuda-libs
^usr/local/bin/bee-gpu-burn$	30-nvidia-cuda-libs
^usr/local/bin/bee-nccl-gpu-stress$	30-nvidia-cuda-libs
^usr/local/bin/bee-john-gpu-stress$	30-nvidia-cuda-libs
^usr/local/bin/bee-dcgmproftester-staggered$	40-nvidia-dcgm-cuda
RULES
}

# Canonicalise the pre-merged-usr paths dpkg records (/bin, /sbin, /lib,
# /lib64) to their real /usr location and drop the leading slash.
bee_layer_canon_paths() {
    sed -e 's#^/bin/#usr/bin/#' \
        -e 's#^/sbin/#usr/sbin/#' \
        -e 's#^/lib/#usr/lib/#' \
        -e 's#^/lib64/#usr/lib64/#' \
        -e 's#^/##'
}


# bee_layer_classify <rootfs_dir> <variant> <work_dir>
# Partition every regular file and symlink under <rootfs_dir> into exactly one
# layer. Writes <work_dir>/<slug>.files (LC_ALL=C sorted) and
# <work_dir>/classify-report.txt. Runs as a subshell: internal state cannot
# leak into the caller, and any step failure aborts with a non-zero exit.
bee_layer_classify() (
    set -e
    _root="$1"
    _variant="$2"
    _wd="$3"

    [ -d "$_root/var/lib/dpkg/info" ] || {
        echo "bee_layer_classify: no dpkg database under $_root" >&2
        exit 1
    }
    mkdir -p "$_wd"
    _slugs="$(bee_layer_slugs_for_variant "$_variant")"
    _tab="$(printf '\t')"

    # 1. Every file and symlink in the tree.
    ( cd "$_root" && find . -mindepth 1 \( -type f -o -type l \) -printf '%P\n' ) \
        | LC_ALL=C sort > "$_wd/all.files"

    # 2. dpkg-owned files that belong to a non-base layer.
    : > "$_wd/dpkg.map"
    for _list in "$_root"/var/lib/dpkg/info/*.list; do
        [ -f "$_list" ] || continue
        _pkg="$(basename "$_list" .list)"
        _pkg="${_pkg%%:*}"
        _layer="$(bee_layer_for_dpkg_pkg "$_pkg")"
        [ -n "$_layer" ] || continue
        bee_layer_canon_paths < "$_list" \
            | LC_ALL=C sort -u \
            | awk -v l="$_layer" -v t="$_tab" 'NF { print $0 t l }' >> "$_wd/dpkg.map"
    done

    # 3. Injected (no-deb) files, by explicit rule. These override dpkg.
    : > "$_wd/injected.map"
    bee_layer_injected_rules | while IFS="$_tab" read -r _re _layer; do
        [ -n "$_re" ] || continue
        LC_ALL=C grep -E "$_re" "$_wd/all.files" \
            | awk -v l="$_layer" -v t="$_tab" '{ print $0 t l }' >> "$_wd/injected.map" || true
    done

    # 4. Merge: injected first so it wins; keep the first layer seen per path.
    #    Drop the merged-usr compat symlinks (bin/sbin/lib/lib64) unconditionally
    #    - a bare "/lib" entry appears in some dpkg .list files (firmware-*), and
    #    those symlinks must always stay in the base layer.
    #    Then keep only rows that name a real file or symlink in the tree: dpkg
    #    .list files also record bare directories, which must never reach an
    #    rsync --files-from list (rsync would copy the directory recursively).
    cat "$_wd/injected.map" "$_wd/dpkg.map" \
        | awk -F'\t' '$1!="bin" && $1!="sbin" && $1!="lib" && $1!="lib64" && !seen[$1]++ { print }' \
        | LC_ALL=C sort -t "$_tab" -k1,1 > "$_wd/candidate.map"
    LC_ALL=C join -t "$_tab" -j 1 "$_wd/all.files" "$_wd/candidate.map" \
        > "$_wd/assigned.map" || true

    # 5. Split into per-layer file lists; base gets the remainder.
    for _slug in $_slugs; do : > "$_wd/$_slug.files"; done
    awk -F'\t' -v wd="$_wd" '{ print $1 >> (wd "/" $2 ".files") }' "$_wd/assigned.map"
    for _slug in $_slugs; do
        LC_ALL=C sort -o "$_wd/$_slug.files" "$_wd/$_slug.files"
    done
    cut -f1 "$_wd/assigned.map" | LC_ALL=C sort > "$_wd/assigned.paths"
    LC_ALL=C comm -23 "$_wd/all.files" "$_wd/assigned.paths" > "$_wd/00-base.files"

    # 6. Completeness / disjointness.
    _total="$(wc -l < "$_wd/all.files")"
    _sum=0
    : > "$_wd/classify-report.txt"
    for _slug in $_slugs; do
        _n="$(wc -l < "$_wd/$_slug.files")"
        _sum=$((_sum + _n))
        printf '%-20s %8d files\n' "$_slug" "$_n" >> "$_wd/classify-report.txt"
    done
    printf '%-20s %8d files\n' "TOTAL" "$_sum" >> "$_wd/classify-report.txt"
    printf '%-20s %8d files\n' "tree" "$_total" >> "$_wd/classify-report.txt"
    cat "$_wd/classify-report.txt"

    if [ "$_sum" -ne "$_total" ]; then
        echo "bee_layer_classify: partition covers $_sum of $_total files" >&2
        exit 1
    fi
    _dups="$(cat $(for _slug in $_slugs; do echo "$_wd/$_slug.files"; done) \
        | LC_ALL=C sort | LC_ALL=C uniq -d)"
    if [ -n "$_dups" ]; then
        echo "bee_layer_classify: files assigned to more than one layer:" >&2
        printf '%s\n' "$_dups" | head >&2
        exit 1
    fi

    # 7. merged-usr guard.
    for _link in bin sbin lib lib64; do
        [ -L "$_root/$_link" ] || continue
        LC_ALL=C grep -qx "$_link" "$_wd/00-base.files" || {
            echo "bee_layer_classify: merged-usr symlink /$_link is not in the base layer" >&2
            exit 1
        }
    done
    for _slug in $_slugs; do
        [ "$_slug" = 00-base ] && continue
        if LC_ALL=C grep -Eq '^(bin|sbin|lib|lib64)/' "$_wd/$_slug.files"; then
            echo "bee_layer_classify: $_slug contains a top-level bin/sbin/lib/lib64 path (merged-usr trap)" >&2
            exit 1
        fi
    done
)

# bee_layer_build <rootfs_dir> <work_dir> <out_dir> <version> <variant>
# Materialise each layer list into filesystem-v<version>-<slug>.squashfs under
# <out_dir>. Enforces the size ceiling. Leaves the monolith untouched.
bee_layer_build() (
    set -e
    _root="$1"; _wd="$2"; _out="$3"; _ver="$4"; _variant="$5"
    _slugs="$(bee_layer_slugs_for_variant "$_variant")"
    mkdir -p "$_out"

    for _slug in $_slugs; do
        _list="$_wd/$_slug.files"
        [ -s "$_list" ] || { echo "bee_layer_build: empty file list for $_slug" >&2; exit 1; }
        _stage="$_wd/stage-$_slug"
        _sq="$_out/filesystem-v${_ver}-${_slug}.squashfs"
        rm -rf "$_stage"
        mkdir -p "$_stage"
        # --files-from with -a recreates implied parent directories from the
        # source but copies only listed entries. --link-dest hardlinks unchanged
        # files from the source tree (same fs) so a 5 GB rootfs is not physically
        # copied once per layer.
        rsync -a --link-dest="$_root" --files-from="$_list" "$_root/" "$_stage/"
        rm -f "$_sq"
        # shellcheck disable=SC2086
        mksquashfs "$_stage" "$_sq" $BEE_LAYER_MKSQUASHFS_OPTS
        rm -rf "$_stage"

        _bytes="$(stat -c '%s' "$_sq")"
        _mib=$(( _bytes / 1048576 ))
        printf '  %-40s %5d MiB\n' "$(basename "$_sq")" "$_mib"
        if [ "$_mib" -gt "$BEE_LAYER_MAX_MIB" ]; then
            echo "bee_layer_build: $(basename "$_sq") is ${_mib} MiB, over the ${BEE_LAYER_MAX_MIB} MiB ceiling" >&2
            exit 1
        fi
        if [ "$_mib" -gt "$BEE_LAYER_TARGET_MIB" ]; then
            echo "  WARNING: $(basename "$_sq") exceeds the ${BEE_LAYER_TARGET_MIB} MiB target" >&2
        fi
    done
)

# bee_layer_verify_each <out_dir> <version> <variant>
# unsquashfs -s plus a full strict-errors extraction of every layer.
bee_layer_verify_each() (
    set -e
    _out="$1"; _ver="$2"; _variant="$3"
    _slugs="$(bee_layer_slugs_for_variant "$_variant")"
    _tmp="$(mktemp -d)"
    trap 'rm -rf "$_tmp"' EXIT
    for _slug in $_slugs; do
        _sq="$_out/filesystem-v${_ver}-${_slug}.squashfs"
        [ -f "$_sq" ] || { echo "bee_layer_verify_each: missing $_sq" >&2; exit 1; }
        unsquashfs -s "$_sq" >/dev/null || { echo "bee_layer_verify_each: unsquashfs -s failed for $_sq" >&2; exit 1; }
        rm -rf "$_tmp/x"
        unsquashfs -d "$_tmp/x" -strict-errors "$_sq" >/dev/null || {
            echo "bee_layer_verify_each: strict extraction failed for $_sq" >&2
            exit 1
        }
    done
    echo "=== all $(echo "$_slugs" | wc -w) layers pass unsquashfs -s + strict extraction ==="
)

# bee_layer_merge <out_dir> <version> <variant> <dest_root>
# Reconstruct the single rootfs live-boot would present, extracting layers in
# module order (last wins), and assert it is bootable.
bee_layer_merge() (
    set -e
    _out="$1"; _ver="$2"; _variant="$3"; _dest="$4"
    _slugs="$(bee_layer_slugs_for_variant "$_variant")"
    rm -rf "$_dest"
    mkdir -p "$_dest"
    for _slug in $_slugs; do
        _sq="$_out/filesystem-v${_ver}-${_slug}.squashfs"
        unsquashfs -f -d "$_dest" "$_sq" >/dev/null || {
            echo "bee_layer_merge: extraction of $_sq failed" >&2
            exit 1
        }
    done

    for _need in usr/sbin/init usr/lib/systemd/systemd usr/local/bin/bee; do
        [ -e "$_dest/$_need" ] || { echo "bee_layer_merge: merged rootfs is missing /$_need" >&2; exit 1; }
    done
    for _link in bin sbin lib lib64; do
        [ -L "$_dest/$_link" ] || { echo "bee_layer_merge: merged rootfs /$_link is not a symlink" >&2; exit 1; }
    done

    case "$_variant" in
        nvidia|nvidia-legacy)
            for _need in \
                usr/local/bin/nvidia-smi \
                usr/local/lib/nvidia/nvidia.ko \
                usr/bin/dcgmi \
                usr/bin/nv-hostengine \
                etc/systemd/system/bee-nvidia.service; do
                [ -e "$_dest/$_need" ] || { echo "bee_layer_merge: merged NVIDIA rootfs is missing /$_need" >&2; exit 1; }
            done
            [ -n "$(find "$_dest/usr/lib/firmware/nvidia" -name 'gsp_*.bin' 2>/dev/null | head -1)" ] \
                || { echo "bee_layer_merge: no GSP firmware in merged rootfs" >&2; exit 1; }
            ls "$_dest"/usr/lib/libnvidia-ml.so* >/dev/null 2>&1 \
                || { echo "bee_layer_merge: libnvidia-ml missing from merged rootfs" >&2; exit 1; }
            for _need in libcudart.so libcublas.so libnccl.so; do
                ls "$_dest"/usr/lib/${_need}* >/dev/null 2>&1 \
                    || { echo "bee_layer_merge: $_need missing from merged rootfs" >&2; exit 1; }
            done
            ls "$_dest"/usr/bin/dcgmproftester* >/dev/null 2>&1 \
                || { echo "bee_layer_merge: dcgmproftester missing from merged rootfs" >&2; exit 1; }
            ;;
    esac
    echo "=== merged rootfs from $(echo "$_slugs" | wc -w) layers is complete and bootable ==="
)

# bee_layer_write_module_file <live_dir> <version> <variant>
# Write <live_dir>/filesystem.module - the explicit, locale-independent image
# order live-boot 1:20230131 reads (MODULE defaults to "filesystem").
bee_layer_write_module_file() (
    set -e
    _live="$1"; _ver="$2"; _variant="$3"
    _slugs="$(bee_layer_slugs_for_variant "$_variant")"
    _mod="$_live/filesystem.module"
    : > "$_mod"
    for _slug in $_slugs; do
        printf 'filesystem-v%s-%s.squashfs\n' "$_ver" "$_slug" >> "$_mod"
    done
    echo "=== wrote $_mod ==="
    cat "$_mod"
)