#!/bin/sh
# YABS probe for the node board.
#   sh probe.sh                    install: register, keep the token, one cron line
#   sh probe.sh schedule weekly    benchmark rate only (daily|weekly|monthly)
#   sh probe.sh update hourly      self-update rate only
#                                  (hourly|daily|weekly|monthly|off)
#   sh probe.sh refresh            download the newest probe.sh and re-install it
#   sh probe.sh now                one run now, output on screen; installs nothing
#   sh probe.sh status             running or not, when it started, last log line
# A bare `sh probe.sh update` is the old spelling of `refresh`: the cron lines
# installed before 2026-10-27 call it that way.
# Install writes under ~/.yabs-probe and is the only command that touches cron.
# Re-running is safe: the token and the schedules are kept.
set -eu

# The customer domain serves the same docroot at the root path; the PMSS panel
# URL (/public-vasilves/yabs/) would put the account name in every log and link.
BOARD=${BOARD:-https://05877bc1b477317f.mcx.fi/yabs/}
DIR=$HOME/.yabs-probe
CONF=$DIR/conf
NAME=$(hostname -s)

usage() {
  echo "usage: sh probe.sh install | refresh | now | status" >&2
  echo "       sh probe.sh schedule <daily|weekly|monthly>" >&2
  echo "       sh probe.sh update <hourly|daily|weekly|monthly|off>" >&2
  exit 2
}

# One command per cron line, so neither can move the other by accident.
cmd=${1:-install}
bench=; upd=
case $cmd in
  install)          [ $# -le 1 ] || usage ;;
  schedule)         [ $# -eq 2 ] || usage; bench=$2; cmd=install ;;
  update)           [ $# -le 2 ] || usage
                    if [ $# -eq 2 ]; then upd=$2; cmd=install; else cmd=refresh; fi ;;
  refresh|now|status) [ $# -le 1 ] || usage ;;
  *) usage ;;
esac
# bench: the benchmark cron line. hourly is refused: a run takes 10-20 minutes.
case $bench in
  ''|daily|weekly|monthly) ;;
  hourly) echo "hourly is for the update line: sh probe.sh update hourly" >&2; exit 2 ;;
  *) echo "unknown benchmark schedule: $bench (daily|weekly|monthly)" >&2; exit 2 ;;
esac
# upd: the probe.sh update line. Empty leaves whatever is installed alone.
case $upd in
  ''|hourly|daily|weekly|monthly|off) ;;
  *) echo "unknown update schedule: $upd (hourly|daily|weekly|monthly|off)" >&2; exit 2 ;;
esac

# cron_time <minute> <hour> <schedule> -> five crontab time fields
cron_time() {
  case $3 in
    hourly)  echo "$1 * * * *" ;;
    daily)   echo "$1 $2 * * *" ;;
    weekly)  echo "$1 $2 * * 1" ;;
    monthly) echo "$1 $2 1 * *" ;;
  esac
}

# The benchmark slot is drawn at install time inside these hours, so the nodes do
# not all hit their disks at the same minute. 01-02 is the quiet part of the
# night: clear of the 00:13 disk probe, the 03:00 backup tarball and the 04:00
# configarr run, and far from the hours someone watches something.
BENCH_FROM=${BENCH_FROM:-1}
BENCH_TO=${BENCH_TO:-3}

rand_slot() {  # one "minute hour" pair inside BENCH_FROM..BENCH_TO
  span=$(( BENCH_TO - BENCH_FROM ))
  [ "$span" -gt 0 ] || span=1
  m=$(od -An -N1 -tu1 /dev/urandom | tr -d ' ')
  h=$(od -An -N1 -tu1 /dev/urandom | tr -d ' ')
  printf '%d %d\n' "$(( m % 60 ))" "$(( BENCH_FROM + h % span ))"
}

# set_cron <text that identifies the line> <line to keep, empty to remove it>
set_cron() {
  cur=$(crontab -l 2>/dev/null || true)
  old=$(printf '%s\n' "$cur" | grep -F "$1" || true)
  if [ -z "$2" ] && [ -z "$old" ]; then
    if [ -t 1 ]; then echo "cron entry already absent"; fi
    return 0
  fi
  if [ "$old" = "$2" ]; then
    if [ -t 1 ]; then echo "cron entry already present: $2"; fi
    return 0
  fi
  cur=$(printf '%s\n' "$cur" | grep -vF "$1" | grep -v '^[[:space:]]*$' || true)
  printf '%s\n%s\n' "$cur" "$2" | sed '/^[[:space:]]*$/d' | crontab -
  if [ -z "$2" ]; then
    echo "cron entry removed: $old"
  elif [ -n "$old" ]; then
    echo "cron entry replaced: $old -> $2"
  else
    echo "cron entry added: $2"
  fi
}

if [ "$cmd" = status ]; then
  first=
  if [ -f "$DIR/progress" ]; then
    first=$(sed -n '1p' "$DIR/progress")
    echo "progress : $first"
    last=$(sed -n '$p' "$DIR/progress")
    if [ "$last" != "$first" ]; then echo "result   : $last"; fi
    started=$(sed -n 's/^started //p' "$DIR/progress" | head -1)
    if now=$(date -d "$started" +%s 2>/dev/null); then
      echo "elapsed  : $(( $(date +%s) - now ))s since start"
    fi
  else
    echo "progress : never run"
  fi
  pid=$(pgrep -f "$DIR/run.sh" 2>/dev/null | head -1 || true)
  if [ -n "${pid:-}" ]; then
    echo "state    : running (pid $pid)"
  else
    echo "state    : not running"
  fi
  echo "node     : $(sed -n 's/^NAME=//p' "$CONF" 2>/dev/null || echo 'not registered')"
  echo "log      : $DIR/last.log"
  if [ -f "$DIR/last.log" ]; then
    tr '\r' '\n' < "$DIR/last.log" | sed 's/\033\[[0-9;]*[A-Za-z]//g' | grep -v '^[[:space:]]*$' | tail -3 || true
  fi
  exit 0
fi

# 'now' never installs, so it needs the node already registered.
if [ "$cmd" = now ] && [ ! -f "$CONF" ]; then
  echo "not installed: run 'sh probe.sh' once to register this node" >&2
  exit 1
fi

command -v curl >/dev/null || { echo "curl is required" >&2; exit 1; }
mkdir -p "$DIR"
chmod 700 "$DIR"

# The conf records the board this node registered with; it wins over the default,
# so a node keeps updating after the owner moves the board.
if [ -f "$CONF" ]; then
  conf_board=$(sed -n 's/^BOARD=//p' "$CONF" | head -1 || true)
  if [ -n "$conf_board" ]; then BOARD=$conf_board; fi
fi

if [ "$cmd" = refresh ]; then
  if [ -t 1 ]; then echo "fetching ${BOARD}probe.sh"; fi
  tmp=$DIR/probe.sh.new
  curl -fsSL "${BOARD}probe.sh" -o "$tmp" || { echo "cannot download ${BOARD}probe.sh" >&2; exit 1; }
  if ! sh -n "$tmp" 2>/dev/null; then
    rm -f "$tmp"
    echo "the downloaded probe.sh is not shell, installed copy kept" >&2
    exit 1
  fi
  if cmp -s "$tmp" "$DIR/probe.sh" 2>/dev/null; then
    rm -f "$tmp"
    # Quiet when nothing changed, so the hourly cron run sends no mail.
    if [ -t 1 ]; then echo "probe.sh already current"; fi
  else
    mv -f "$tmp" "$DIR/probe.sh"
    echo "probe.sh replaced"
  fi
  chmod 700 "$DIR/probe.sh"
  # Re-run the new copy with no argument, so run.sh follows it and the two cron
  # lines stay exactly as they are.
  sh "$DIR/probe.sh" || { echo "install from the new copy failed" >&2; exit 1; }
  exit 0
fi

# yabs cannot report a GPU, so the probe passes the render node name along.
gpu=$(ls /dev/dri/renderD* 2>/dev/null | head -1 || true)
gpu=${gpu##*/}

if [ ! -f "$DIR/yabs.sh" ]; then
  echo "fetching yabs.sh"
  curl -fsSL https://yabs.sh -o "$DIR/yabs.sh" || { echo "cannot download yabs.sh" >&2; exit 1; }
  chmod 600 "$DIR/yabs.sh"
fi

if [ ! -f "$CONF" ]; then
  echo "registering $NAME on $BOARD"
  body=$(printf '{"name":"%s","gpu":"%s"}' "$NAME" "$gpu")
  out=$(curl -fsS -X POST -H 'Content-Type: application/json' -d "$body" "${BOARD}?do=register") || {
    echo "registration refused: $out" >&2
    echo "if the name is taken, ask the board owner to delete its row" >&2
    exit 1
  }
  token=$(printf '%s' "$out" | sed -n 's/.*"token":"\([^"]*\)".*/\1/p')
  [ -n "$token" ] || { echo "no token in the answer: $out" >&2; exit 1; }
  umask 077
  printf 'BOARD=%s\nNAME=%s\nTOKEN=%s\n' "$BOARD" "$NAME" "$token" > "$CONF"
  chmod 600 "$CONF"
fi

cat > "$DIR/run.sh" <<'RUN'
#!/usr/bin/env bash
# One YABS run, posted to the board. Generated by probe.sh.
# --fg streams yabs output to the terminal instead of last.log.
set -u
. "$HOME/.yabs-probe/conf"
cd "$HOME/.yabs-probe"
fg=${1:-}
gpu=$(ls /dev/dri/renderD* 2>/dev/null | head -1 || true)
gpu=${gpu##*/}

# Disk media behind $HOME, as the kernel reports it. /sys/block/<dev> names the
# physical queue, so the rotational flag is the disk's own; lsblk adds the
# transport (nvme is flash by definition) and the kernel's ATA rotation rate
# overrides a flag that a virtio layer sometimes gets wrong. A virtio disk
# reports `virt=1`: QEMU leaves the rotational flag at 1 whatever the backing
# store is, so the board must show `virtio` instead of reading that as HDD.
# YABS reports none of this, and the board never guesses it from the IOPS.
disk_media() {
  dev=$(df -P "$HOME" 2>/dev/null | awk 'NR==2{print $1}')
  base=${dev##*/}
  if [ -n "$base" ] && [ ! -d "/sys/block/$base" ] && command -v lsblk >/dev/null 2>&1; then
    base=$(lsblk -no PKNAME "$dev" 2>/dev/null | head -1)
  fi
  [ -n "$base" ] || return 0
  rota=$(cat "/sys/block/$base/queue/rotational" 2>/dev/null || true)
  tran=
  if command -v lsblk >/dev/null 2>&1; then
    tran=$(lsblk -dno TRAN "/dev/$base" 2>/dev/null | tr -d '[:space:]' || true)
  fi
  rpm=$(udevadm info -q property -n "/dev/$base" 2>/dev/null | sed -n 's/^ID_ATA_ROTATION_RATE_RPM=//p' | head -1)
  if [ -n "$rpm" ] && [ "$rpm" -gt 0 ] 2>/dev/null; then rota=1; fi
  # virtio-blk and virtio-scsi both sit under a /sys/devices/.../virtioN/ node.
  virt=
  case $(readlink -f "/sys/class/block/$base" 2>/dev/null) in
    */virtio*) virt=1 ;;
  esac
  printf '%s %s %s\n' "${rota:-}" "${tran:-}" "${virt:-}"
}

# IO pressure over the run: /proc/pressure/io counts the microseconds a task
# spent stalled, so the difference between two reads over the elapsed time is
# the average stall of the window. High means the fio number measured a
# congested disk, not the disk's own speed. Absent when CONFIG_PSI is off.
psi_us() {  # <some|full> -> cumulative stall microseconds
  [ -r /proc/pressure/io ] || return 0
  sed -n "s/^$1 .*total=\([0-9][0-9]*\).*/\1/p" /proc/pressure/io 2>/dev/null | head -1
}

psi_pct() {  # <start_us> <end_us> <elapsed_s> -> share of the window stalled
  [ -n "$1" ] && [ -n "$2" ] && [ "${3:-0}" -gt 0 ] 2>/dev/null || return 0
  # LC_ALL=C: a comma decimal separator would make the value unparsable remotely.
  LC_ALL=C awk -v a="$1" -v b="$2" -v s="$3" 'BEGIN {
    p = (b - a) / 1e6 / s * 100
    if (p < 0) p = 0
    if (p > 100) p = 100
    printf "%.1f", p
  }'
}

echo "started $(date -u +%FT%TZ)" > progress
set -- $(disk_media); rota=${1:-}; tran=${2:-}; virt=${3:-}
t0=$(date +%s); s0=$(psi_us some); f0=$(psi_us full)
rc=0
# yabs.sh is bash, not POSIX, hence the explicit bash call. -g skips Geekbench:
# it downloads and unpacks several GB for one score.
if [ "$fg" = "--fg" ]; then
  bash ./yabs.sh -w ./last.json -g 2>&1 | tee ./last.log || rc=$?
else
  bash ./yabs.sh -w ./last.json -g > ./last.log 2>&1 || rc=$?
fi
t1=$(date +%s); s1=$(psi_us some); f1=$(psi_us full)
ps=$(psi_pct "$s0" "$s1" $((t1 - t0)))
pf=$(psi_pct "$f0" "$f1" $((t1 - t0)))
q=
[ -n "$ps" ] && q="$q&psi_some=$ps"
[ -n "$pf" ] && q="$q&psi_full=$pf"
[ -n "$rota" ] && q="$q&rota=$rota"
[ -n "$tran" ] && q="$q&tran=$tran"
[ -n "$virt" ] && q="$q&virt=$virt"
if [ ! -s ./last.json ]; then
  echo "failed $(date -u +%FT%TZ) yabs exited $rc without JSON, see last.log" > progress
  echo "no JSON produced, yabs exited $rc, see last.log" >&2
  exit 1
fi
if out=$(curl -fsS -X POST --data-binary @./last.json \
    "$BOARD?token=$TOKEN&node=$NAME&gpu=$gpu$q"); then
  echo "posted $(date -u +%FT%TZ) $out" > progress
  echo "$out"
  rm -f ./last.json
else
  # Keep the JSON so a manual re-post is possible without a second 10-minute run.
  echo "failed $(date -u +%FT%TZ) upload refused, JSON kept in last.json" > progress
  echo "upload refused, JSON kept in last.json" >&2
  exit 1
fi
RUN
chmod 700 "$DIR/run.sh"

if [ "$cmd" = install ]; then
  cur=$(crontab -l 2>/dev/null || true)
  old=$(printf '%s\n' "$cur" | grep -F "$DIR/run.sh" || true)
  # No benchmark argument with a line already installed: keep its rate and slot.
  if [ -n "$bench" ]; then
    slot=$(rand_slot)
    when=$(cron_time "${slot% *}" "${slot#* }" "$bench")
  elif [ -n "$old" ]; then when=$(printf '%s\n' "$old" | cut -d' ' -f1-5)
  else slot=$(rand_slot)
       when=$(cron_time "${slot% *}" "${slot#* }" monthly)
  fi
  # 04:23, so a daily run does not pile onto the 04:00 update cron.
  set_cron "$DIR/run.sh" "$when $DIR/run.sh"
  # The update line only moves when the second word says so.
  old=$(crontab -l 2>/dev/null | grep -F "$DIR/probe.sh update" || true)
  upd_when=$(printf '%s\n' "${old:-none}" | cut -d' ' -f1-5)
  case $upd in
    '') ;;
    off) set_cron "$DIR/probe.sh update" ''; upd_when=none ;;
    *) upd_when=$(cron_time 11 5 "$upd")
       set_cron "$DIR/probe.sh update" "$upd_when $DIR/probe.sh update" ;;
  esac

  if [ -t 1 ]; then
    echo
    echo "node     : $NAME"
    echo "gpu      : ${gpu:-none detected}"
    echo "schedule : $when   (benchmark, $DIR/run.sh)"
    echo "update   : $upd_when   ($DIR/probe.sh update)"
    echo "token    : $(sed -n 's/^TOKEN=//p' "$CONF")"
    echo "config   : $CONF (0600)"
    echo "log      : $DIR/last.log"
    echo "run now  : sh probe.sh now"
    echo "progress : sh probe.sh status"
  fi
fi

[ "$cmd" = now ] || exit 0
echo
exec "$DIR/run.sh" --fg