#!/usr/bin/env bash
# spdk-loader.sh — declarative version manager for the long-lived SPDK target.
#
# Runs as a long-lived compose service alongside the nvmeof target. On (re)start:
#   1. STAGE the CURRENT release (${IMAGE_VERSION}) from the mounted artifact
#      image (/in) into the shared release store (/releases/<ver>/) and ADD it to
#      the running target ONLINE (no target restart).
#   2. Then stay up and REAP: periodically poll connection counts and remove any
#      OTHER (non-current) version once it has been drained (0 connected
#      controllers) for a sustained grace window — i.e. once every host that was
#      running it has rebooted onto the new UKI.
#
# Migration is therefore: bump IMAGE_VERSION in .env, `docker compose up -d`.
# Compose recreates THIS loader (its artifact mount + env changed) and the
# http/dnsmasq UKI servers (image volume changed), but NOT the nvmeof target
# (its definition is unchanged) — so live connections are never dropped. The new
# loader adds the new rootfs online and drains the old one as hosts roll over.
#
# SAFETY: with a single served UKI a host can only ever boot the CURRENT version,
# so a non-current version at 0 connections is genuinely orphaned (no powered-off
# host can return to it — on next boot it fetches the current UKI). Removal is
# still drain-gated in spdk-release (refuses >0), and gated here by a sustained-
# zero grace window (>= REAP_GRACE_SECONDS) to ride out transient nvme-tcp
# reconnects. Set REAP_GRACE_SECONDS comfortably above the host ctrl_loss_tmo.
set -euo pipefail

CURRENT="${IMAGE_VERSION:?spdk-loader needs IMAGE_VERSION}"
RELEASES_DIR="${RELEASES_DIR:-/releases}"
IN="${ARTIFACT_IN:-/in}"
POLL="${REAP_POLL_SECONDS:-30}"
GRACE="${REAP_GRACE_SECONDS:-300}"
REAP="${REAP_OLD:-true}"
RELEASE_CTL=$(command -v spdk-release || echo /usr/local/bin/spdk-release)
# spdk-release reads RPC_SOCK from the environment (shared socket to the target).

echo "==> spdk-loader: current=${CURRENT} reap=${REAP} poll=${POLL}s grace=${GRACE}s"

# --- 1) stage + add the current release (idempotent) -----------------------
dest="$RELEASES_DIR/$CURRENT"
mkdir -p "$dest"
if [ -s "$IN/verity.json" ]; then
    # Skip the (large) copy if already staged with identical verity.json.
    if [ -s "$dest/verity.json" ] && cmp -s "$IN/verity.json" "$dest/verity.json"; then
        echo "==> $CURRENT already staged in $dest (verity.json identical) — skipping copy"
    else
        echo "==> staging $CURRENT from $IN -> $dest"
        cp "$IN/rootfs.erofs" "$dest/rootfs.erofs"
        cp "$IN/rootfs.hash"  "$dest/rootfs.hash"
        cp "$IN/verity.json"  "$dest/verity.json"
    fi
else
    echo "WARNING: no $IN/verity.json (artifact image not mounted at $IN?) — relying on pre-staged $dest" >&2
fi
bash "$RELEASE_CTL" add "$CURRENT"
bash "$RELEASE_CTL" list

if [ "$REAP" != "true" ]; then
    echo "==> REAP_OLD=$REAP — staying up for status only; old versions are NOT auto-removed."
fi

# --- 2) reap loop: a non-current version drained for >= GRACE is removed ----
declare -A zero_since
while true; do
    sleep "$POLL"
    [ "$REAP" = "true" ] || continue
    [ -d "$RELEASES_DIR" ] || continue
    for d in "$RELEASES_DIR"/*/; do
        [ -e "${d}verity.json" ] || continue
        v=$(basename "$d")
        [ "$v" = "$CURRENT" ] && continue          # never reap the current version
        n=$(bash "$RELEASE_CTL" count "$v" 2>/dev/null || echo -1)
        if [ "$n" = "0" ]; then
            ts=${zero_since[$v]:-}
            if [ -z "$ts" ]; then
                ts=$(date +%s); zero_since[$v]=$ts
                echo "==> $v drained (0 conns) — starting ${GRACE}s grace before removal"
            fi
            age=$(( $(date +%s) - ts ))
            if [ "$age" -ge "$GRACE" ]; then
                echo "==> reaping $v (drained ${age}s >= ${GRACE}s grace)"
                if bash "$RELEASE_CTL" remove "$v"; then
                    rm -rf "${RELEASES_DIR:?}/$v"
                    unset 'zero_since[$v]'
                    echo "==> $v removed + staged files deleted"
                else
                    # A connection reappeared between the poll and the remove —
                    # spdk-release refused (drain-gated). Reset and retry later.
                    echo "WARNING: remove $v refused (reconnected?) — grace reset, will retry" >&2
                    unset 'zero_since[$v]'
                fi
            fi
        else
            # >0 connections, or RPC error (-1): reset the grace timer.
            if [ -n "${zero_since[$v]:-}" ]; then
                echo "==> $v no longer drained (conns=$n) — grace reset"
                unset 'zero_since[$v]'
            fi
        fi
    done
done