#!/usr/bin/env bash
#
# One-command entry point for the RX 580 prompt-processing benchmark.
#
# Runs from your workstation. Copies the pod-side harness into the pod and
# execs it there over kubectl. All the real work happens on the pod; this is
# a thin wrapper so there is exactly one command to remember.
#
# Examples:
#   ./scripts/rx580-bench/run.sh --build /root/llama.cpp/build --label fork-before
#   ./scripts/rx580-bench/run.sh --build /root/arms/new-clean/build --label new-clean \
#                               --build-b /root/arms/new-fork/build --label-b new-fork
#   ./scripts/rx580-bench/run.sh --summarize
#   ./scripts/rx580-bench/run.sh --install-only
#
# Every argument other than the wrapper-only flags below is passed straight
# through to bench.sh on the pod. See bench.sh --help.
#
set -euo pipefail

KUBECONFIG_PATH="${KUBECONFIG:-/home/user/Projects/klaster/talos/generated/kubeconfig}"
NS="${RX580_NS:-llama}"
DEPLOY="${RX580_DEPLOY:-deploy/supervisord}"
POD_BENCH_DIR=/root/bench
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

export KUBECONFIG="$KUBECONFIG_PATH"

kx() { kubectl -n "$NS" exec -i "$DEPLOY" -- sh -c "$1"; }

install_harness() {
  echo "[run] installing harness into $POD_BENCH_DIR"
  kx "mkdir -p $POD_BENCH_DIR"
  for f in bench.sh ppbench.py lbsweep.sh summarize.py; do
    kx "cat > $POD_BENCH_DIR/$f" < "$HERE/$f"
  done
  kx "chmod +x $POD_BENCH_DIR/bench.sh $POD_BENCH_DIR/lbsweep.sh $POD_BENCH_DIR/ppbench.py $POD_BENCH_DIR/summarize.py"
  # The corpus is large and rarely changes; only push it if missing.
  if ! kx "test -s $POD_BENCH_DIR/pan-tadeusz.txt" 2>/dev/null; then
    if [ -f "$HERE/pan-tadeusz.txt" ]; then
      echo "[run] uploading corpus"
      kx "cat > $POD_BENCH_DIR/pan-tadeusz.txt" < "$HERE/pan-tadeusz.txt"
    else
      echo "[run] ERROR: $POD_BENCH_DIR/pan-tadeusz.txt missing on the pod and" >&2
      echo "      no local copy at $HERE/pan-tadeusz.txt to upload." >&2
      exit 4
    fi
  fi
}

ARGS=()
INSTALL_ONLY=0
SUMMARIZE=0
for a in "$@"; do
  case "$a" in
    --install-only) INSTALL_ONLY=1 ;;
    --summarize)    SUMMARIZE=1 ;;
    *) ARGS+=("$a") ;;
  esac
done

install_harness

if [ "$INSTALL_ONLY" -eq 1 ]; then
  echo "[run] harness installed; not running anything"
  exit 0
fi

if [ "$SUMMARIZE" -eq 1 ]; then
  kx "python3 $POD_BENCH_DIR/summarize.py ${ARGS[*]:-}"
  exit 0
fi

if [ ${#ARGS[@]} -eq 0 ]; then
  kx "sh $POD_BENCH_DIR/bench.sh --help"
  exit 2
fi

# A full 3-round A/B run takes hours; run detached on the pod and tail it, so a
# dropped kubectl connection cannot orphan a half-finished run with llama-swap
# still down. bench.sh restarts llama-swap from its own EXIT trap either way.
STAMP="$(date -u +%Y%m%d-%H%M%S)"
LOG="$POD_BENCH_DIR/run-$STAMP.log"
echo "[run] starting detached run on the pod, log: $LOG"
kx "cd $POD_BENCH_DIR && nohup sh $POD_BENCH_DIR/bench.sh ${ARGS[*]} > $LOG 2>&1 & echo started"

echo "[run] tailing until ALL_RUNS_DONE (safe to Ctrl-C: the pod keeps running)"
kx "i=0; while [ \$i -lt 100000 ]; do
      if grep -qE 'ALL_RUNS_DONE|Traceback|error:' $LOG; then break; fi
      sleep 15; i=\$((i+1));
    done; cat $LOG"