#!/usr/bin/env bash
# KAE Stage-0 benchmark: LZ4 hardware (libkaelz4 via UADK) vs software
# (system liblz4) at zfile's per-block granularity.
#
# What it answers (go/no-go for a `kae` crate + async compression pool):
#   1. Does this host have KAE hardware at all? (Kunpeng 920 new model / 950,
#      ZIP PCI device, /dev/uacce, libkaelz4 installed)
#   2. Does hardware LZ4 beat software LZ4 at the block size zfile actually
#      uses (4 KiB per LZ4_compress_default call)? At which block size does
#      it start winning? (informs async batching and block-size questions)
#   3. Does hardware offload CPU even when it does not win on throughput?
#   4. Is hardware-produced LZ4 block format decodable by stock liblz4?
#      (zfile read path / resume keeps using software decompression)
#   5. Which async symbols does libkaelz4 export? (block-format async
#      coverage question)
#
# Usage:
#   scripts/kae_stage0_lz4_bench.sh                     # synthetic data, 256 MiB
#   scripts/kae_stage0_lz4_bench.sh --input /path/to/memory-snapshot-file
#   scripts/kae_stage0_lz4_bench.sh --size 64           # 64 MiB synthetic
#   KAE_LIB=/other/path/libkaelz4.so scripts/kae_stage0_lz4_bench.sh
#
# Requirements: bash, gcc/cc, python3 (optional, for page-interleaved mixed
# data), system liblz4. Run on the production Kunpeng bare-metal host.
# Results: printed table + TSV file (default ./kae_stage0_results.tsv).
set -euo pipefail

SIZE_MIB=256
INPUT=""
OUT_TSV="kae_stage0_results.tsv"
BLOCKS="4096 8192 65536 262144 1048576"
DATASETS="zeros mixed random"

usage() { sed -n '2,30p' "$0" | sed 's/^# \{0,1\}//'; exit "${1:-0}"; }
while [[ $# -gt 0 ]]; do
  case "$1" in
    --input) INPUT="$2"; shift 2 ;;
    --size) SIZE_MIB="$2"; shift 2 ;;
    --out) OUT_TSV="$2"; shift 2 ;;
    --blocks) BLOCKS="$2"; shift 2 ;;
    --data) DATASETS="$2"; shift 2 ;;
    -h|--help) usage 0 ;;
    *) echo "unknown arg: $1" >&2; usage 1 ;;
  esac
done

TMP="$(mktemp -d /tmp/kae-stage0.XXXXXX)"
trap 'rm -rf "$TMP"' EXIT

CC_BIN="${CC:-$(command -v gcc || command -v cc || true)}"
[[ -n "$CC_BIN" ]] || { echo "ERROR: no gcc/cc found; block-granularity test needs a compiler" >&2; exit 1; }

echo "== 1. Environment / hardware checks =="
CPU_MODEL="$(lscpu 2>/dev/null | awk -F': +' '/^(BIOS Model name|Model name|Vendor ID):/{v=$2} END{print v}')"
echo "CPU:            ${CPU_MODEL:-unknown}"
echo "Kernel:         $(uname -r) ($(uname -m))"
ZIP_PCI="$(lspci 2>/dev/null | grep -iE 'zip|hisi' || true)"
echo "ZIP PCI dev:    ${ZIP_PCI:-none found (lspci missing or no ZIP engine)}"
UACCE="$(ls /dev/uacce* 2>/dev/null || true)"
echo "uacce devices:  ${UACCE:-none}"

SW_LIB="$(ldconfig -p 2>/dev/null | awk '$1=="liblz4.so.1"{print $NF; exit}')"
[[ -n "$SW_LIB" ]] || { echo "ERROR: system liblz4.so.1 not found" >&2; exit 1; }
echo "SW lib:         $SW_LIB"

KAE_LIB="${KAE_LIB:-}"
if [[ -z "$KAE_LIB" ]]; then
  # KAELz4 ships a patched liblz4.so.1.9.4 exposing the standard LZ4 ABI
  # (LZ4_compress_default etc.) and forwarding to libkaelz4/libwd: that is
  # the sync drop-in path zfile would use, so bench it by default.
  for c in /usr/local/kaelz4/lib/liblz4.so.1 /usr/local/kaelz4/lib/libkaelz4.so /usr/local/lib/libkaelz4.so; do
    [[ -e "$c" ]] && KAE_LIB="$c" && break
  done
  [[ -z "$KAE_LIB" ]] && KAE_LIB="$(find /usr/local /usr/lib64 \( -name 'liblz4.so.1.9.4' -o -name 'libkaelz4.so*' \) 2>/dev/null | head -1 || true)"
fi
if [[ -n "$KAE_LIB" && -e "$KAE_LIB" ]]; then
  echo "KAE lib:        $KAE_LIB"
  echo "KAE deps:       $(ldd "$KAE_LIB" 2>/dev/null | grep -oE 'lib(wd|kaelz4)[^ ]*\.so[^ ]*' | tr '\n' ' ' || echo 'no libwd linkage?!')"
else
  echo "KAE lib:        NOT FOUND (hw tests skipped; install KAE2.0 or set KAE_LIB)"
  KAE_LIB=""
fi
echo

echo "== 2. Input data =="
if [[ -n "$INPUT" ]]; then
  [[ -f "$INPUT" ]] || { echo "ERROR: --input file not found: $INPUT" >&2; exit 1; }
  DATA_FILES="input"
  echo "using provided input: $INPUT ($(stat -c%s "$INPUT") bytes)"
else
  DATA_FILES=""
  HAVE_PY=0; command -v python3 >/dev/null 2>&1 && HAVE_PY=1
  for d in $DATASETS; do
    f="$TMP/$d.bin"
    if [[ "$d" == input ]]; then continue; fi
    if [[ "$HAVE_PY" == 1 ]]; then
      python3 - "$f" "$SIZE_MIB" "$d" <<'PYEOF'
import os, sys
path, mib, kind = sys.argv[1], int(sys.argv[2]), sys.argv[3]
size = mib * 1024 * 1024
page = 4096
text = b"agentenv stage0 sample payload " + b"0123456789abcdef" * 16
with open(path, "wb") as f:
    for i in range(size // page):
        c = i % 10
        if kind == "zeros" or (kind == "mixed" and c < 5):
            buf = b"\x00" * page
        elif kind == "random" or (kind == "mixed" and c < 7):
            buf = os.urandom(page)
        else:
            buf = ((b"snapshot-page-%d " % (i % 97)) * (page // 18) + text)[:page]
        f.write(buf)
PYEOF
    else
      # coarse fallback: contiguous thirds, not page-interleaved
      case "$d" in
        zeros)  dd if=/dev/zero    of="$f" bs=1M count="$SIZE_MIB" status=none ;;
        random) dd if=/dev/urandom of="$f" bs=1M count="$SIZE_MIB" status=none ;;
        mixed)
          dd if=/dev/zero    of="$TMP/mz.bin" bs=1M count=$((SIZE_MIB/3+1)) status=none
          dd if=/dev/urandom of="$TMP/mr.bin" bs=1M count=$((SIZE_MIB/3+1)) status=none
          head -c $((SIZE_MIB*1024*1024/3+1)) /dev/urandom | base64 | head -c $((SIZE_MIB*1024*1024/3+1)) > "$TMP/mt.bin" || true
          cat "$TMP/mz.bin" "$TMP/mr.bin" "$TMP/mt.bin" | head -c $((SIZE_MIB*1024*1024)) > "$f"
          ;;
      esac
    fi
    DATA_FILES="$DATA_FILES $d"
    echo "generated $f ($(stat -c%s "$f") bytes)"
  done
  [[ "$HAVE_PY" == 1 ]] || echo "note: python3 missing -> mixed data is coarse contiguous thirds" >&2
fi
echo

echo "== 3. Build harness =="
cat > "$TMP/bench.c" <<'CEOF'
/* dlopen-based LZ4 block benchmark: compresses input per --block-byte blocks
 * via LZ4_compress_default from a chosen library, times an adaptive pass
 * loop (rusage + monotonic), then verifies every block decodes bit-exact
 * with the verify-lib LZ4_decompress_safe. */
#define _GNU_SOURCE
#include <dlfcn.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>

typedef int (*comp_fn)(const char *, char *, int, int);
typedef int (*decomp_fn)(const char *, char *, int, int);
typedef int (*bound_fn)(int);

static double now_s(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec + ts.tv_nsec * 1e-9;
}

static void *must_dlopen(const char *path) {
    void *h = dlopen(path, RTLD_NOW);
    if (!h) {
        fprintf(stderr, "OPEN_FAIL %s: %s\n", path, dlerror());
        exit(2);
    }
    return h;
}

static void usage(void) {
    fprintf(stderr, "usage: bench --lib SO [--verify-lib SO] --input FILE --block N | --probe\n");
    exit(1);
}

int main(int argc, char **argv) {
    const char *lib = NULL, *vlib = NULL, *input = NULL;
    long block = 0;
    int probe = 0;
    for (int i = 1; i < argc; i++) {
        if (!strcmp(argv[i], "--probe")) probe = 1;
        else if (!strcmp(argv[i], "--lib") && i + 1 < argc) lib = argv[++i];
        else if (!strcmp(argv[i], "--verify-lib") && i + 1 < argc) vlib = argv[++i];
        else if (!strcmp(argv[i], "--input") && i + 1 < argc) input = argv[++i];
        else if (!strcmp(argv[i], "--block") && i + 1 < argc) block = strtol(argv[++i], NULL, 10);
        else usage();
    }
    if (!lib) usage();
    void *h = must_dlopen(lib);
    comp_fn comp = (comp_fn)dlsym(h, "LZ4_compress_default");
    decomp_fn dec = (decomp_fn)dlsym(h, "LZ4_decompress_safe");
    bound_fn bound = (bound_fn)dlsym(h, "LZ4_compressBound");

    if (probe) {
        static const char *names[] = {
            "LZ4_compress_default", "LZ4_compress_fast", "LZ4_compress_HC",
            "LZ4_decompress_safe", "LZ4_compressBound", "LZ4F_compressFrame",
            /* async surface (coverage question for the kae crate); the
             * KAELZ4_* block-format async API lives in libkaelz4.so */
            "LZ4_async_compress_init", "LZ4F_compressFrame_async",
            "KAELZ4_compress_async", "KAELZ4_compress_async_in_session",
            "KAELZ4_create_async_compress_session",
            "kaelz4_async_compress_polling", "KAELZ4_decompress_async", NULL};
        printf("PROBE lib=%s\n", lib);
        for (int i = 0; names[i]; i++)
            printf("SYMBOL %s %s\n", names[i], dlsym(h, names[i]) ? "found" : "missing");
        return 0;
    }

    if (!input || block <= 0) usage();
    if (!comp || !dec) {
        printf("RESULT status=SYMBOL_MISSING lib=%s LZ4_compress_default=%p LZ4_decompress_safe=%p\n",
               lib, (void *)comp, (void *)dec);
        return 3;
    }
    int fd = open(input, O_RDONLY);
    if (fd < 0) { perror("open input"); return 1; }
    struct stat st;
    fstat(fd, &st);
    size_t in_size = (size_t)st.st_size;
    char *src = malloc(in_size);
    if (!src || read(fd, src, in_size) != (ssize_t)in_size) { perror("read input"); return 1; }
    close(fd);
    long nblocks = (long)(in_size / (size_t)block);
    if (nblocks == 0) { fprintf(stderr, "input smaller than block size\n"); return 1; }
    int blk = (int)block;
    int bound_v = bound ? bound(blk) : blk + blk / 255 + 16;
    char *dst = malloc((size_t)bound_v);
    char *raw = malloc((size_t)blk);

    /* warmup + sanity: one full pass, also faults pages in */
    for (long b = 0; b < nblocks; b++) {
        if (comp(src + b * block, dst, blk, bound_v) <= 0) {
            printf("RESULT status=COMPRESS_FAIL lib=%s block=%ld at=%ld\n", lib, block, b);
            return 4;
        }
    }

    /* verify: decode with the verify lib, bit-exact compare (format compat) */
    void *vh = (vlib && strcmp(vlib, lib)) ? must_dlopen(vlib) : h;
    decomp_fn vdec = (decomp_fn)dlsym(vh, "LZ4_decompress_safe");
    long vfail = 0;
    for (long b = 0; b < nblocks; b++) {
        int n = comp(src + b * block, dst, blk, bound_v);
        int d = vdec ? vdec(dst, raw, n, blk) : -1;
        if (n <= 0 || d != blk || memcmp(raw, src + b * block, (size_t)blk)) vfail++;
    }

    /* timed adaptive loop: >=3 passes, >=0.8 s, <=99 passes */
    struct rusage r0, r1;
    getrusage(RUSAGE_SELF, &r0);
    double t0 = now_s();
    long passes = 0, out_total = 0;
    for (;;) {
        for (long b = 0; b < nblocks; b++)
            out_total += comp(src + b * block, dst, blk, bound_v);
        passes++;
        double el = now_s() - t0;
        if ((passes >= 3 && el >= 0.8) || passes >= 99) break;
    }
    double wall = now_s() - t0;
    getrusage(RUSAGE_SELF, &r1);
    double user = (r1.ru_utime.tv_sec - r0.ru_utime.tv_sec) +
                  (r1.ru_utime.tv_usec - r0.ru_utime.tv_usec) * 1e-6;
    double sys = (r1.ru_stime.tv_sec - r0.ru_stime.tv_sec) +
                 (r1.ru_stime.tv_usec - r0.ru_stime.tv_usec) * 1e-6;
    double in_bytes = (double)block * (double)nblocks * (double)passes;

    printf("RESULT status=OK lib=%s input=%s block=%ld blocks=%ld passes=%ld "
           "in_bytes=%.0f out_bytes=%.0f ratio_pct=%.2f wall=%.3f gbps=%.2f "
           "user_s=%.3f sys_s=%.3f verify=%s\n",
           lib, input, block, nblocks, passes, in_bytes, (double)out_total,
           (double)out_total / in_bytes * 100.0, wall, in_bytes / wall / 1e9,
           user, sys, vfail ? "FAIL" : "OK");
    return vfail ? 5 : 0;
}
CEOF
"$CC_BIN" -O2 -o "$TMP/bench" "$TMP/bench.c" -ldl
echo "harness built: $TMP/bench"
echo

echo "== 4. Symbol probe =="
"$TMP/bench" --probe --lib "$SW_LIB" | sed 's/^/  /'
if [[ -n "$KAE_LIB" ]]; then
  KAE_DIR="$(dirname "$KAE_LIB")"
  export LD_LIBRARY_PATH="${KAE_DIR}:/usr/local/kaelz4/lib:/usr/local/uadk/lib:${LD_LIBRARY_PATH:-}"
  "$TMP/bench" --probe --lib "$KAE_LIB" | sed 's/^/  /'
  # libkaelz4 carries the async surface (block-format async + polling)
  KAE_ASYNC_LIB="$KAE_DIR/libkaelz4.so"
  [[ -e "$KAE_ASYNC_LIB" && "$KAE_LIB" != "$KAE_ASYNC_LIB" ]] && "$TMP/bench" --probe --lib "$KAE_ASYNC_LIB" | sed 's/^/  /'
else
  echo "  (hw probe skipped)"
fi
echo

echo "== 5. Sweep (adaptive passes per cell, ~1s each) =="
{
  echo "# kae stage0 lz4 bench  date=$(date -Is) host=$(uname -n)"
  echo "# cpu=${CPU_MODEL:-unknown} kernel=$(uname -r) uacce=${UACCE:-none} kae_lib=${KAE_LIB:-none}"
  printf 'data\tblock\tlib\tgbps\tratio_pct\tcpu_s_per_gb\tverify\n'
} > "$OUT_TSV"

run_one() { # $1=data $2=block $3=libpath $4=label
  local out line
  out="$("$TMP/bench" --lib "$3" --verify-lib "$SW_LIB" --input "$5" --block "$2" 2>&1)" || true
  line="$(printf '%s\n' "$out" | grep '^RESULT' || true)"
  local gbps ratio cpu vfy
  gbps="$(printf '%s\n' "$line" | sed -n 's/.*gbps=\([0-9.]*\).*/\1/p')"
  ratio="$(printf '%s\n' "$line" | sed -n 's/.*ratio_pct=\([0-9.]*\).*/\1/p')"
  vfy="$(printf '%s\n' "$line" | sed -n 's/.*verify=\([A-Z]*\).*/\1/p')"
  cpu=""
  if [[ -n "$line" && -n "$gbps" ]]; then
    cpu="$(printf '%s\n' "$line" | awk '{u=($0~/user_s=/)?0:0; if(match($0,/user_s=[0-9.]+/)){u=substr($0,RSTART+7,RLENGTH-7)+0}; s=0; if(match($0,/sys_s=[0-9.]+/)){s=substr($0,RSTART+6,RLENGTH-6)+0}; i=0; if(match($0,/in_bytes=[0-9]+/)){i=substr($0,RSTART+9,RLENGTH-9)+0}; printf "%.3f",(i>0)?(u+s)*1e9/i:0}')"
  fi
  printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\n' "$1" "$2" "$4" "${gbps:-SKIPPED}" "${ratio:--}" "${cpu:--}" "${vfy:-$(printf '%s' "$out" | head -c 80)}"
  printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\n' "$1" "$2" "$4" "${gbps:-SKIPPED}" "${ratio:--}" "${cpu:--}" "${vfy:-$(printf '%s' "$out" | head -c 80)}" >> "$OUT_TSV"
}

for d in $DATA_FILES; do
  if [[ "$d" == input ]]; then f="$INPUT"; else f="$TMP/$d.bin"; fi
  for blk in $BLOCKS; do
    run_one "$d" "$blk" "$SW_LIB" "sw" "$f"
    [[ -n "$KAE_LIB" ]] && run_one "$d" "$blk" "$KAE_LIB" "hw" "$f"
  done
done
echo

echo "== 6. Summary =="
column -t -s $'\t' "$OUT_TSV" 2>/dev/null || cat "$OUT_TSV"
echo
echo "results written to: $OUT_TSV"
echo

echo "== 7. Verdict (mixed data) =="
awk -F'\t' '
  $1=="mixed" && $7=="OK" && $4+0>0 {
    gbps[$2"_"$3]=$4; cpu[$2"_"$3]=$6
    if (!seen[$2]++) blks[++nb]=$2
  }
  END {
    if (!gbps["4096_sw"]) { print "  no usable mixed/4K rows; check verify column"; exit }
    for (i = 1; i <= nb; i++) {
      b = blks[i]; s = gbps[b"_sw"]; h = gbps[b"_hw"]
      if (!s) continue
      if (h) {
        printf "  block %-7d speedup hw/sw = %.2fx", b, h/s
        if (cpu[b"_sw"]+0 > 0) printf ", hw cpu = %.0f%% of sw", cpu[b"_hw"]*100/cpu[b"_sw"]
        print ""
      } else print "  block " b ": hw missing/skipped"
    }
    print ""
    h4 = gbps["4096_hw"]; s4 = gbps["4096_sw"]
    if (h4+0 >= s4*1.1) print "  => 4K DIRECT WIN: hardware beats software at zfile granularity; sync drop-in (Stage 1) is worth it."
    else if (h4 && cpu["4096_hw"]+0 <= cpu["4096_sw"]*0.5) print "  => 4K CPU OFFLOAD: throughput parity-or-worse but CPU cost much lower; worth it iff host CPU contention is the goal."
    else if (gbps["262144_hw"]+0 >= gbps["262144_sw"]*1.5 || gbps["1048576_hw"]+0 >= gbps["1048576_sw"]*1.5)
      print "  => SMALL BLOCK LOSES, BIG BLOCK WINS: only a batched async pool (Stage 2) or larger blocks could help; re-evaluate with those numbers."
    else print "  => NO-GO on this hardware at this granularity."
  }' "$OUT_TSV"