#!/bin/sh
# test-squashfs-layers.sh - regression tests for lib/squashfs-layers.sh.
#
# Builds a synthetic mini-rootfs, splits it, and asserts the semantic layer
# contract holds. Requires mksquashfs / unsquashfs / rsync; skips cleanly when
# they are not installed (e.g. on a developer macOS box).
set -eu

BUILDER_DIR="$(CDPATH= cd -- "$(dirname "$0")" && pwd)"
. "$BUILDER_DIR/lib/squashfs-layers.sh"

for tool in mksquashfs unsquashfs rsync; do
    command -v "$tool" >/dev/null 2>&1 || {
        echo "test-squashfs-layers: $tool not available, skipping"
        exit 0
    }
done

T="$(mktemp -d)"
trap 'rm -rf "$T"' EXIT INT TERM HUP
ROOT="$T/root"

mk() { mkdir -p "$(dirname "$ROOT/$1")"; printf '%s' "${2:-x}" > "$ROOT/$1"; }
mkbig() { mkdir -p "$(dirname "$ROOT/$1")"; head -c "${2:-1048576}" /dev/urandom > "$ROOT/$1"; }
listf() { mkdir -p "$ROOT/var/lib/dpkg/info"; printf '%s\n' "$@" > "$ROOT/var/lib/dpkg/info/$1.list"; }

# --- base rootfs ---
mk usr/lib/systemd/systemd "#!systemd"
( cd "$ROOT/usr" && mkdir -p sbin && ln -s ../lib/systemd/systemd sbin/init )
mk usr/local/bin/bee "#!bee"
mk etc/hostname "easy-bee"
mk usr/bin/true "#!true"
mk usr/bin/dmidecode "#!dmi"
( cd "$ROOT" && ln -s usr/bin bin && ln -s usr/bin sbin && ln -s usr/lib lib && ln -s usr/lib lib64 )
listf base-files /etc/hostname /usr/bin/true
listf dmidecode /usr/bin/dmidecode

# --- 05 firmware: non-NVIDIA device firmware. The bare "/lib" entry mimics the
#     real firmware-* .list files and must NOT drag the /lib symlink out of base.
listf firmware-iwlwifi /lib /lib/firmware /lib/firmware/iwlwifi-cc-a0-77.ucode
mk usr/lib/firmware/iwlwifi-cc-a0-77.ucode
listf firmware-realtek /lib/firmware/rtl_nic/rtl8168h-2.fw
mk usr/lib/firmware/rtl_nic/rtl8168h-2.fw

# --- 08 desktop ---
listf chromium /usr/bin/chromium
mk usr/bin/chromium "#!chromium"
listf xserver-xorg-core /usr/bin/Xorg
mk usr/bin/Xorg "#!Xorg"
listf lightdm /usr/sbin/lightdm
mk usr/sbin/lightdm "#!lightdm"
listf libgl1-mesa-dri /usr/lib/x86_64-linux-gnu/dri/swrast_dri.so
mk usr/lib/x86_64-linux-gnu/dri/swrast_dri.so
listf fonts-dejavu-core /usr/share/fonts/truetype/dejavu/DejaVuSans.ttf
mk usr/share/fonts/truetype/dejavu/DejaVuSans.ttf

# --- 10 driver: dpkg + injected ---
listf firmware-nvidia-gsp /lib /lib/firmware /lib/firmware/nvidia/580.159.03/gsp_ga10x.bin
mk usr/lib/firmware/nvidia/580.159.03/gsp_ga10x.bin
mk usr/lib/firmware/nvidia/580.159.03/gsp_tu10x.bin
listf nvidia-modprobe /usr/bin/nvidia-modprobe
mk usr/bin/nvidia-modprobe
mk usr/local/lib/nvidia/nvidia.ko
mk usr/local/lib/nvidia/nvidia-uvm.ko
mk usr/local/bin/nvidia-smi "#!nvidia-smi"
mk usr/lib/libnvidia-ml.so.1
mk usr/lib/libnvidia-ml.so.580.159.03
mk usr/lib/libcuda.so.1
mk etc/OpenCL/vendors/nvidia.icd "libnvidia-opencl.so.1"
mk etc/systemd/system/bee-nvidia.service "[Unit]"
mk usr/local/bin/bee-nvidia-load "#!load"

# --- 20 platform ---
listf nvidia-fabricmanager /usr/bin/nv-fabricmanager /usr/lib/systemd/system/nvidia-fabricmanager.service
mk usr/bin/nv-fabricmanager
mk usr/lib/systemd/system/nvidia-fabricmanager.service "[Unit]"
listf libnvidia-nscq /usr/lib/x86_64-linux-gnu/libnvidia-nscq.so.580.159.03
mk usr/lib/x86_64-linux-gnu/libnvidia-nscq.so.580.159.03
listf datacenter-gpu-manager-4-core /usr/bin/dcgmi /usr/bin/nv-hostengine
mk usr/bin/dcgmi "#!dcgmi"
mk usr/bin/nv-hostengine "#!nvh"
mk etc/systemd/system/nvidia-fabricmanager.service.d/bee-nvswitch-check.conf "[Service]"

# --- 30 cuda libs (injected) ---
mk usr/lib/libcudart.so.13
mk usr/lib/libcublas.so.13
mk usr/lib/libcublasLt.so.13
mk usr/lib/libnccl.so.2
mk usr/local/bin/all_reduce_perf "#!arp"
mk usr/local/lib/bee/bee-gpu-burn-worker "#!worker"

# --- 40 dcgm cuda ---
listf datacenter-gpu-manager-4-cuda13 /usr/bin/dcgmproftester12
mkbig usr/bin/dcgmproftester12 3145728
mk usr/local/bin/bee-dcgmproftester-staggered "#!stag"

# a plain "nvidia" in the path must NOT be enough on its own: an unowned,
# unruled file under an nvidia dir stays in base.
mk usr/share/doc/nvidia-random/README "notes"

# --- empty directories: no file/symlink anywhere in their subtree, so they
#     never appear in all.files and must be classified/carried separately.
mkdir -p "$ROOT/var/log/nvidia-dcgm"                 # ruled -> 20-nvidia-platform
mkdir -p "$ROOT/var/lib/bee-empty-unruled"            # unruled -> base
mkdir -p "$ROOT/var/lib/bee-empty-nested/still-empty" # nested empty branch -> base

fail() { echo "FAIL: $*" >&2; exit 1; }
in_layer() { LC_ALL=C grep -qx "$2" "$WD/$1.files" || fail "expected $2 in layer $1"; }
in_dir_layer() { LC_ALL=C grep -qx "$2" "$WD/$1.dirs" || fail "expected empty dir $2 in layer $1"; }

MONO="$T/filesystem-v14.99.squashfs"
mksquashfs "$ROOT" "$MONO" -comp zstd -noappend -no-progress -no-xattrs >/dev/null

R2="$T/r2"; unsquashfs -d "$R2" "$MONO" >/dev/null
WD="$T/cls"
bee_layer_classify "$R2" nvidia "$WD"

in_layer 00-base usr/sbin/init
in_layer 00-base usr/lib/systemd/systemd
in_layer 00-base usr/local/bin/bee
in_layer 00-base etc/hostname
in_layer 00-base usr/bin/dmidecode
in_layer 00-base bin
in_layer 00-base sbin
in_layer 00-base lib
in_layer 00-base lib64
in_layer 05-firmware usr/lib/firmware/iwlwifi-cc-a0-77.ucode
in_layer 05-firmware usr/lib/firmware/rtl_nic/rtl8168h-2.fw
in_layer 08-desktop usr/bin/chromium
in_layer 08-desktop usr/bin/Xorg
in_layer 08-desktop usr/sbin/lightdm
in_layer 08-desktop usr/lib/x86_64-linux-gnu/dri/swrast_dri.so
in_layer 08-desktop usr/share/fonts/truetype/dejavu/DejaVuSans.ttf
in_layer 10-nvidia-driver usr/local/lib/nvidia/nvidia.ko
in_layer 10-nvidia-driver usr/local/bin/nvidia-smi
in_layer 10-nvidia-driver usr/lib/libnvidia-ml.so.1
in_layer 10-nvidia-driver usr/lib/libcuda.so.1
in_layer 10-nvidia-driver usr/lib/firmware/nvidia/580.159.03/gsp_ga10x.bin
in_layer 10-nvidia-driver etc/OpenCL/vendors/nvidia.icd
in_layer 10-nvidia-driver etc/systemd/system/bee-nvidia.service
in_layer 10-nvidia-driver usr/bin/nvidia-modprobe
in_layer 20-nvidia-platform usr/bin/nv-fabricmanager
in_layer 20-nvidia-platform usr/lib/x86_64-linux-gnu/libnvidia-nscq.so.580.159.03
in_layer 20-nvidia-platform usr/bin/dcgmi
in_layer 20-nvidia-platform usr/bin/nv-hostengine
in_layer 20-nvidia-platform etc/systemd/system/nvidia-fabricmanager.service.d/bee-nvswitch-check.conf
in_layer 30-nvidia-cuda-libs usr/lib/libcudart.so.13
in_layer 30-nvidia-cuda-libs usr/lib/libcublasLt.so.13
in_layer 30-nvidia-cuda-libs usr/lib/libnccl.so.2
in_layer 30-nvidia-cuda-libs usr/local/bin/all_reduce_perf
in_layer 30-nvidia-cuda-libs usr/local/lib/bee/bee-gpu-burn-worker
in_layer 40-nvidia-dcgm-cuda usr/bin/dcgmproftester12
in_layer 40-nvidia-dcgm-cuda usr/local/bin/bee-dcgmproftester-staggered
in_layer 00-base usr/share/doc/nvidia-random/README

in_dir_layer 20-nvidia-platform var/log/nvidia-dcgm
in_dir_layer 00-base var/lib/bee-empty-unruled
in_dir_layer 00-base var/lib/bee-empty-nested
in_dir_layer 00-base var/lib/bee-empty-nested/still-empty

OUT="$T/out"
bee_layer_build "$R2" "$WD" "$OUT" 14.99 nvidia
bee_layer_verify_each "$OUT" 14.99 nvidia
bee_layer_merge "$OUT" 14.99 nvidia "$T/merged"

[ -d "$T/merged/var/log/nvidia-dcgm" ] || fail "empty dir var/log/nvidia-dcgm missing from merged rootfs"
[ -d "$T/merged/var/lib/bee-empty-unruled" ] || fail "empty dir var/lib/bee-empty-unruled missing from merged rootfs"
[ -d "$T/merged/var/lib/bee-empty-nested/still-empty" ] || fail "nested empty dir missing from merged rootfs"

MODLIVE="$T/live"; mkdir -p "$MODLIVE"
cp "$OUT"/filesystem-v14.99-*.squashfs "$MODLIVE/"
bee_layer_write_module_file "$MODLIVE" 14.99 nvidia
expect_mod="filesystem-v14.99-00-base.squashfs
filesystem-v14.99-05-firmware.squashfs
filesystem-v14.99-08-desktop.squashfs
filesystem-v14.99-10-nvidia-driver.squashfs
filesystem-v14.99-20-nvidia-platform.squashfs
filesystem-v14.99-30-nvidia-cuda-libs.squashfs
filesystem-v14.99-40-nvidia-dcgm-cuda.squashfs"
[ "$(cat "$MODLIVE/filesystem.module")" = "$expect_mod" ] || fail "module file order wrong"

# --- failure modes ---
# missing layer -> merge fails
rm -f "$OUT/filesystem-v14.99-20-nvidia-platform.squashfs"
if bee_layer_merge "$OUT" 14.99 nvidia "$T/m2" >/dev/null 2>&1; then
    fail "merge succeeded with a missing layer"
fi
bee_layer_build "$R2" "$WD" "$OUT" 14.99 nvidia >/dev/null   # restore

# size ceiling -> build fails
if BEE_LAYER_MAX_MIB=0 bee_layer_build "$R2" "$WD" "$T/out0" 14.99 nvidia >/dev/null 2>&1; then
    fail "build ignored the size ceiling"
fi

# incomplete partition -> classify fails
rm -rf "$WD"; bee_layer_classify "$R2" nvidia "$WD"
echo "usr/bin/does-not-exist" >> "$WD/10-nvidia-driver.files"
_total=$(wc -l < "$WD/all.files")
_sum=0
for s in $(bee_layer_slugs_for_variant nvidia); do _sum=$((_sum + $(wc -l < "$WD/$s.files"))); done
[ "$_sum" -ne "$_total" ] || fail "tampered partition should not still sum to the tree size"

echo "squashfs layer tests: OK"