#!/usr/bin/env bash
# bend-autoscaler.sh — bin-pack autopilot for a bend factory.
#
# Watches:
#   - .ready / .bin tier distribution on $LUMBDA_QUEUE_DIR
#   - live RAM / VRAM headroom on this host
#   - current TIER_*_DISPATCH caps (from supervisor.config)
#
# Computes a safe per-tier cap vector that maximizes throughput under a
# memory budget — bin-pack: greedily fill smallest tiers first (highest
# cap, lightest per-dispatch footprint), then ratchet larger tiers until
# predicted RAM/VRAM saturate.
#
# Writes new caps to $LUMBDA_QUEUE_DIR/supervisor.config. Supervisor's
# reload_config picks up our change within $POLL_SECONDS, kills the
# dispatcher, restart loop respawns with new env. No SIGHUP gymnastics.
#
# Tuned for a 24-core / 62 GB RAM / 24 GB VRAM workstation but every
# threshold stays parametric — substrate-grade.
#
# Heartbeat: $LUMBDA_QUEUE_DIR/autoscaler.heartbeat via lib-heartbeat.sh.
# Supervisor restarts us if we go silent.
#
# Exit:
#   - SIGTERM/INT
#   - $LUMBDA_QUEUE_DIR/AUTOSCALER_STOP marker
#
# Env vars consumed: see $LUMBDA_FACTORY_DIR/CONTRACT.md for our canonical
# table (LUMBDA_QUEUE_DIR, LUMBDA_BACKEND_PROC_PATTERN, AUTOSCALER_V2,
# AUTOSCALER_PEAK_INITIAL, AUTOSCALER_CPU_CEILING,
# AUTOSCALER_DLQ_DAMPING_THRESHOLD, TIER_*_DISPATCH_CEILING, ...).
set -u

QUEUE_DIR="${LUMBDA_QUEUE_DIR:-${QUEUE_DIR:-/tmp/lumbda-queue}}"
POLL_SECONDS="${AUTOSCALER_POLL:-20}"
LOG="${AUTOSCALER_LOG:-$QUEUE_DIR/autoscaler.log}"
CONFIG_FILE="$QUEUE_DIR/supervisor.config"
STOP_FILE="$QUEUE_DIR/AUTOSCALER_STOP"

LUMBDA_FACTORY_DIR="${LUMBDA_FACTORY_DIR:-$(dirname "$(readlink -f "$0")")}"
LIB_HEARTBEAT="${LIB_HEARTBEAT:-$LUMBDA_FACTORY_DIR/lib-heartbeat.sh}"
LIB_TIER="${LIB_TIER:-$LUMBDA_FACTORY_DIR/lib-tier.sh}"
# shellcheck source=/dev/null
. "$LIB_HEARTBEAT"
# shellcheck source=/dev/null
. "$LIB_TIER"

# Backend proc pattern — pgrep -f match for counting in-flight backend forks.
BACKEND_PROC_PATTERN="${LUMBDA_BACKEND_PROC_PATTERN:-bend-cuda}"

# V1 (legacy) — fixed budget × total knobs. Kept as fallback via AUTOSCALER_V2=0.
RAM_BUDGET_FRACTION="${RAM_BUDGET_FRACTION:-0.70}"
VRAM_BUDGET_FRACTION="${VRAM_BUDGET_FRACTION:-0.85}"
MIN_DISPATCH_PER_TIER="${MIN_DISPATCH_PER_TIER:-1}"

# ── V2 controller knobs ──────────────────────────────────────────
# V2 replaces V1's "reserve PEAK × allocated SLOT for every tier" model with
# a live-state model: sample actual VRAM free + running backend fork peaks
# + DLQ growth rate, plan = running + headroom_workers. Fall-through
# admission gives single-tier demand the full budget; ±25 % per-poll rate
# damping prevents oscillation; DLQ surge halves a plan next poll.
AUTOSCALER_V2="${AUTOSCALER_V2:-1}"                    # 1 = V2 live model (default), 0 = V1 fallback
AUTOSCALER_PEAK_INITIAL="${AUTOSCALER_PEAK_INITIAL:-4096}"  # MIB per backend fork before first observation
AUTOSCALER_PEAK_MIN="${AUTOSCALER_PEAK_MIN:-2048}"     # MIB floor (never under-estimate)
AUTOSCALER_EWMA_ALPHA="${AUTOSCALER_EWMA_ALPHA:-0.3}"  # weight of new observation (0.3 = slow decay; max() biased anyway)
AUTOSCALER_CPU_CEILING="${AUTOSCALER_CPU_CEILING:-30}" # never plan > N total workers (24 cores + modest oversub)
AUTOSCALER_RAMP_FROM_ZERO="${AUTOSCALER_RAMP_FROM_ZERO:-4}"  # max plan when cur=0 (cold-start step)
AUTOSCALER_DLQ_DAMPING_THRESHOLD="${AUTOSCALER_DLQ_DAMPING_THRESHOLD:-10}"  # DLQ entries/poll above which next plan halves
AUTOSCALER_VRAM_SAFETY_RESERVE_MIB="${AUTOSCALER_VRAM_SAFETY_RESERVE_MIB:-2048}"  # never plan into last N MIB of free VRAM
# RAM model — V2 originally used VRAM EWMA peak for RAM headroom too. Real
# backend RSS runs 1.5-2× VRAM (lib-tier.sh:48-49 calibration). Separate
# EWMA for RAM peak avoids over-estimating headroom + edging into swap.
AUTOSCALER_RAM_PEAK_INITIAL="${AUTOSCALER_RAM_PEAK_INITIAL:-6400}"  # MIB per fork RSS guess pre-observation
AUTOSCALER_RAM_PEAK_MIN="${AUTOSCALER_RAM_PEAK_MIN:-4096}"          # MIB floor on RAM/worker estimate
AUTOSCALER_RAM_SAFETY_RESERVE_MIB="${AUTOSCALER_RAM_SAFETY_RESERVE_MIB:-8192}"  # leave N MIB unallocated for OS + pool emits
# Swap-pressure signal — fires when MemAvailable falls below this fraction
# of MemTotal OR when SwapUsed grows over poll period. Triggers a halving
# of plan capacity AND writes pool advisory file so emit pool can slow down.
AUTOSCALER_RAM_PRESSURE_FRACTION="${AUTOSCALER_RAM_PRESSURE_FRACTION:-0.15}"   # < 15% free MemAvailable = pressure
AUTOSCALER_SWAP_PRESSURE_DELTA_MIB="${AUTOSCALER_SWAP_PRESSURE_DELTA_MIB:-256}" # SwapUsed grew > N MIB in last poll = active swap
# Per-tier hard ceilings — V2 never plans above these regardless of headroom.
# Same shape as lib-tier.sh defaults but explicit so V2 reads them even when
# autoscaler runs against an unchanged supervisor.config.
TIER_MICRO_DISPATCH_CEILING="${TIER_MICRO_DISPATCH_CEILING:-24}"
TIER_SMALL_DISPATCH_CEILING="${TIER_SMALL_DISPATCH_CEILING:-16}"
TIER_MEDIUM_DISPATCH_CEILING="${TIER_MEDIUM_DISPATCH_CEILING:-12}"
TIER_LARGE_DISPATCH_CEILING="${TIER_LARGE_DISPATCH_CEILING:-6}"
# HUGE — solo-dispatch tier. Cap = 1 enforced + V2 drains all other tiers
# before admitting. Largest candidate bins (6-12 GB) land here.
TIER_HUGE_DISPATCH_CEILING="${TIER_HUGE_DISPATCH_CEILING:-1}"

log() {
    local ts msg
    ts=$(date +%H:%M:%S)
    msg="[$ts] $*"
    echo "$msg" | tee -a "$LOG"
}

# Read host capacity. Returns lines:
#   total_ram_mib XXXX
#   total_vram_mib YYYY
read_capacity() {
    local total_ram_kb total_ram_mib total_vram_mib
    total_ram_kb=$(awk '/^MemTotal:/ {print $2}' /proc/meminfo 2>/dev/null || echo 0)
    total_ram_mib=$((total_ram_kb / 1024))
    if command -v nvidia-smi >/dev/null 2>&1; then
        total_vram_mib=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -dc 0-9)
    fi
    [ -z "$total_vram_mib" ] && total_vram_mib=0
    printf 'total_ram_mib %d\ntotal_vram_mib %d\n' "$total_ram_mib" "$total_vram_mib"
}

# Count .ready + .bin files per tier in queue. Single ls + stat per call.
# Emits lines: ready_<tier> N (incl. huge — solo-dispatch tier).
count_ready_by_tier() {
    declare -A counts
    counts[micro]=0; counts[small]=0; counts[medium]=0; counts[large]=0; counts[huge]=0
    local f tag sz tier
    for f in "$QUEUE_DIR"/*.ready "$QUEUE_DIR"/*.bin; do
        [ -f "$f" ] || continue
        case "$f" in
            *.ready) tag=$(basename "$f" .ready) ;;
            *.bin)   tag=$(basename "$f" .bin) ;;
        esac
        # De-dup: prefer .bin (covers .ready + .inflight states)
        local bin="$QUEUE_DIR/${tag}.bin"
        [ -f "$bin" ] || continue
        # Skip cells already done/inflight (in-flight handled by tier_inflight)
        [ -f "$QUEUE_DIR/${tag}.done" ] && continue
        sz=$(stat -c %s "$bin" 2>/dev/null || echo 0)
        tier=$(tier_of_size "$sz")
        case "$tier" in
            micro|small|medium|large|huge) counts[$tier]=$((counts[$tier] + 1)) ;;
        esac
    done
    for t in micro small medium large huge; do
        printf 'ready_%s %d\n' "$t" "${counts[$t]}"
    done
}

# Bin-pack: pick per-tier dispatch counts that maximize sum(count_t × value_t)
# under RAM + VRAM budgets. value_t is throughput proxy (assume 1 per dispatch
# slot; weighting by tier's expected wall-time would refine, but uniform-1 is
# a good first cut — our autopilot's job is admission, not perfect throughput).
#
# Greedy: start with all tiers at min, scan small→large. For each tier,
# allocate as many slots as fit AND as demand (ready count) wants AND
# tier's hard ceiling from lib-tier defaults. Stop when budget exhausted.
#
# Outputs four lines:
#   plan_micro N
#   plan_small N
#   plan_medium N
#   plan_large N
plan_caps() {
    local ram_budget="$1" vram_budget="$2"
    local ready_micro="$3" ready_small="$4" ready_medium="$5" ready_large="$6"
    local ram_used=0 vram_used=0
    declare -A plan demand cap_max cpu vram
    plan[micro]=$MIN_DISPATCH_PER_TIER
    plan[small]=$MIN_DISPATCH_PER_TIER
    plan[medium]=$MIN_DISPATCH_PER_TIER
    plan[large]=$MIN_DISPATCH_PER_TIER
    demand[micro]=$ready_micro
    demand[small]=$ready_small
    demand[medium]=$ready_medium
    demand[large]=$ready_large
    # cap_max = lib-tier HARD CEILINGS (not the live caps autoscaler just
    # wrote). Reading from tier_dispatch would chicken-and-egg: autoscaler
    # writes cap=1, next iter reads cap_max=1, can never grow back up
    # even when budget freed. Hardcode ceilings here so a planner always
    # considers a full feasible range.
    cap_max[micro]=24
    cap_max[small]=16
    cap_max[medium]=12
    cap_max[large]=6
    cpu[micro]=$(tier_cpu_mib micro)
    cpu[small]=$(tier_cpu_mib small)
    cpu[medium]=$(tier_cpu_mib medium)
    cpu[large]=$(tier_cpu_mib large)
    vram[micro]=$(tier_vram_mib micro)
    vram[small]=$(tier_vram_mib small)
    vram[medium]=$(tier_vram_mib medium)
    vram[large]=$(tier_vram_mib large)

    # Pre-charge minimum-per-tier floor.
    local t
    for t in micro small medium large; do
        ram_used=$((ram_used + plan[$t] * cpu[$t]))
        vram_used=$((vram_used + plan[$t] * vram[$t]))
    done

    # Greedy fill LARGEST tier first. Smaller tiers drain fast (10-30 s
    # per cell) so under-allocating them is cheap; large tiers (200+ s
    # per cell) starve catastrophically — at L=1 a 40-cell large queue
    # takes ~2 hours to drain. Fill large→medium→small→micro so big
    # demand gets priority on VRAM budget; small tiers backfill a
    # remainder + still clear quickly.
    for t in large medium small micro; do
        # Skip tiers with no demand.
        [ "${demand[$t]}" -eq 0 ] && continue
        local want=$((demand[$t] < cap_max[$t] ? demand[$t] : cap_max[$t]))
        while [ "${plan[$t]}" -lt "$want" ]; do
            local next_ram=$((ram_used + cpu[$t]))
            local next_vram=$((vram_used + vram[$t]))
            if [ "$next_ram" -gt "$ram_budget" ] || [ "$next_vram" -gt "$vram_budget" ]; then
                break
            fi
            plan[$t]=$((plan[$t] + 1))
            ram_used=$next_ram
            vram_used=$next_vram
        done
    done

    for t in micro small medium large; do
        printf 'plan_%s %d\n' "$t" "${plan[$t]}"
    done
    printf 'ram_used_mib %d ram_budget_mib %d\n' "$ram_used" "$ram_budget"
    printf 'vram_used_mib %d vram_budget_mib %d\n' "$vram_used" "$vram_budget"
}

# Read current TIER_*_DISPATCH values from supervisor.config (if present).
# Emits: cur_<tier> N
read_current_caps() {
    local cur_micro="${TIER_MICRO_DISPATCH:-0}"
    local cur_small="${TIER_SMALL_DISPATCH:-0}"
    local cur_medium="${TIER_MEDIUM_DISPATCH:-0}"
    local cur_large="${TIER_LARGE_DISPATCH:-0}"
    local cur_huge="${TIER_HUGE_DISPATCH:-0}"
    if [ -f "$CONFIG_FILE" ]; then
        # shellcheck source=/dev/null
        local v
        v=$(grep -E '^TIER_MICRO_DISPATCH=' "$CONFIG_FILE" | tail -1 | cut -d= -f2)
        [ -n "$v" ] && cur_micro="$v"
        v=$(grep -E '^TIER_SMALL_DISPATCH=' "$CONFIG_FILE" | tail -1 | cut -d= -f2)
        [ -n "$v" ] && cur_small="$v"
        v=$(grep -E '^TIER_MEDIUM_DISPATCH=' "$CONFIG_FILE" | tail -1 | cut -d= -f2)
        [ -n "$v" ] && cur_medium="$v"
        v=$(grep -E '^TIER_LARGE_DISPATCH=' "$CONFIG_FILE" | tail -1 | cut -d= -f2)
        [ -n "$v" ] && cur_large="$v"
        v=$(grep -E '^TIER_HUGE_DISPATCH=' "$CONFIG_FILE" | tail -1 | cut -d= -f2)
        [ -n "$v" ] && cur_huge="$v"
    fi
    printf 'cur_micro %d\ncur_small %d\ncur_medium %d\ncur_large %d\ncur_huge %d\n' \
        "$cur_micro" "$cur_small" "$cur_medium" "$cur_large" "$cur_huge"
}

# Atomically rewrite supervisor.config preserving non-TIER lines + replacing
# TIER_<NAME>_DISPATCH = plan_<name>. Writes through a tempfile so a partial
# write can't poison a supervisor's source.
write_caps() {
    local micro="$1" small="$2" medium="$3" large="$4" huge="$5"
    local tmp
    tmp="$(mktemp "${CONFIG_FILE}.XXXXXX")"
    if [ -f "$CONFIG_FILE" ]; then
        grep -vE '^TIER_(MICRO|SMALL|MEDIUM|LARGE|HUGE)_DISPATCH=' "$CONFIG_FILE" > "$tmp"
    fi
    printf 'TIER_MICRO_DISPATCH=%d\nTIER_SMALL_DISPATCH=%d\nTIER_MEDIUM_DISPATCH=%d\nTIER_LARGE_DISPATCH=%d\nTIER_HUGE_DISPATCH=%d\n' \
        "$micro" "$small" "$medium" "$large" "$huge" >> "$tmp"
    mv "$tmp" "$CONFIG_FILE"
}

# ── V2 controller ────────────────────────────────────────────────
# All V2 helpers namespaced v2_* so V1 code paths stay readable.

# v2_sample_live — query current VRAM + backend forks + peak per fork +
# count of in-flight HUGE bins + RAM pressure indicators.
# Emits:
#   vram_free_mib N
#   ram_free_mib  N        (MemAvailable in MIB)
#   ram_total_mib N        (MemTotal in MIB — for pressure fraction calc)
#   swap_used_mib N        (SwapTotal - SwapFree in MIB)
#   running_forks N
#   observed_peak_mib N    (max VRAM use among running backend forks)
#   observed_ram_peak_mib N (max RSS among running backend forks — separate
#                            from VRAM since RSS = 1.5-2x VRAM typically)
#   observed_ram_sum_mib N (sum RSS — used to bound headroom_ram + log)
#   inflight_huge N        (count of HUGE bins claimed by a dispatcher worker
#                           but not yet .done)
v2_sample_live() {
    local vram_free=0 ram_free_kb=0 ram_free=0 ram_total_kb=0 ram_total=0
    local swap_total_kb=0 swap_free_kb=0 swap_used=0
    local running=0 peak=0 ram_peak=0 ram_sum=0 inflight_huge=0
    if command -v nvidia-smi >/dev/null 2>&1; then
        vram_free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -dc 0-9)
        [ -z "$vram_free" ] && vram_free=0
    fi
    ram_free_kb=$(awk '/^MemAvailable:/ {print $2}' /proc/meminfo 2>/dev/null || echo 0)
    ram_total_kb=$(awk '/^MemTotal:/ {print $2}' /proc/meminfo 2>/dev/null || echo 0)
    swap_total_kb=$(awk '/^SwapTotal:/ {print $2}' /proc/meminfo 2>/dev/null || echo 0)
    swap_free_kb=$(awk '/^SwapFree:/ {print $2}' /proc/meminfo 2>/dev/null || echo 0)
    ram_free=$((ram_free_kb / 1024))
    ram_total=$((ram_total_kb / 1024))
    swap_used=$(( (swap_total_kb - swap_free_kb) / 1024 ))
    running=$(pgrep -f "$BACKEND_PROC_PATTERN" 2>/dev/null | wc -l)
    if [ "$running" -gt 0 ]; then
        # VRAM peak via nvidia-smi
        if command -v nvidia-smi >/dev/null 2>&1; then
            peak=$(nvidia-smi --query-compute-apps=used_memory --format=csv,noheader,nounits 2>/dev/null \
                   | tr -dc '0-9\n' | grep -v '^$' | sort -n | tail -1)
            [ -z "$peak" ] && peak=0
        fi
        # RAM peak + sum via ps RSS — independent of VRAM since backend RSS
        # = ops table + ancilla buffers + CUDA driver context, scales with bin
        # size differently than GPU side does.
        local rss_list
        rss_list=$(pgrep -f "$BACKEND_PROC_PATTERN" 2>/dev/null | xargs -r -I{} ps -o rss= -p {} 2>/dev/null | tr -d ' ')
        if [ -n "$rss_list" ]; then
            ram_peak=$(echo "$rss_list" | sort -n | tail -1)
            ram_peak=$((ram_peak / 1024))  # kB → MIB
            ram_sum=$(echo  "$rss_list" | awk '{s+=$1} END {print s}')
            ram_sum=$((ram_sum / 1024))
        fi
    fi
    # Count in-flight HUGE bins.
    shopt -s nullglob
    local f tag bin sz
    for f in "$QUEUE_DIR"/*.inflight-*; do
        [ ! -f "$f" ] && continue
        tag=$(basename "$f" | sed 's/\.inflight-[0-9]*$//')
        bin="$QUEUE_DIR/${tag}.bin"
        [ ! -f "$bin" ] && continue
        sz=$(stat -c %s "$bin" 2>/dev/null || echo 0)
        if [ "$sz" -gt "${TIER_LARGE_MAX:-4294967296}" ]; then
            inflight_huge=$((inflight_huge + 1))
        fi
    done
    shopt -u nullglob
    printf 'vram_free_mib %d\nram_free_mib %d\nram_total_mib %d\nswap_used_mib %d\nrunning_forks %d\nobserved_peak_mib %d\nobserved_ram_peak_mib %d\nobserved_ram_sum_mib %d\ninflight_huge %d\n' \
        "$vram_free" "$ram_free" "$ram_total" "$swap_used" "$running" "$peak" "$ram_peak" "$ram_sum" "$inflight_huge"
}

# v2_update_ewma — read prior peak from autoscaler.ewma, fold in new sample.
# Bias up (max) for safety so under-estimate is impossible. Floor at
# AUTOSCALER_PEAK_MIN to avoid runaway shrink when zero forks observed.
# Prints new peak.
v2_update_ewma() {
    local observed_peak="$1"
    local ewma_file="$QUEUE_DIR/autoscaler.ewma"
    local old_peak="$AUTOSCALER_PEAK_INITIAL"
    if [ -f "$ewma_file" ]; then
        local v
        v=$(grep -E '^peak_per_worker_mib=' "$ewma_file" | tail -1 | cut -d= -f2)
        [ -n "$v" ] && old_peak="$v"
    fi
    local new_peak
    if [ "$observed_peak" -gt 0 ]; then
        new_peak=$(awk -v old="$old_peak" -v obs="$observed_peak" -v a="$AUTOSCALER_EWMA_ALPHA" \
            'BEGIN { smoothed = old*(1-a) + obs*a;
                     if (obs > smoothed) smoothed = obs;
                     printf "%d", smoothed }')
    else
        new_peak="$old_peak"
    fi
    [ "$new_peak" -lt "$AUTOSCALER_PEAK_MIN" ] && new_peak="$AUTOSCALER_PEAK_MIN"
    {
        echo "peak_per_worker_mib=$new_peak"
        echo "last_updated=$(date -u +%FT%TZ)"
        echo "last_observed_mib=$observed_peak"
    } > "$ewma_file"
    echo "$new_peak"
}

# v2_update_ram_ewma — same shape as v2_update_ewma but for RAM peak.
# Separate file (autoscaler.ram-ewma) so VRAM + RAM tracks stay independent.
v2_update_ram_ewma() {
    local observed_peak="$1"
    local ewma_file="$QUEUE_DIR/autoscaler.ram-ewma"
    local old_peak="$AUTOSCALER_RAM_PEAK_INITIAL"
    if [ -f "$ewma_file" ]; then
        local v
        v=$(grep -E '^ram_peak_per_worker_mib=' "$ewma_file" | tail -1 | cut -d= -f2)
        [ -n "$v" ] && old_peak="$v"
    fi
    local new_peak
    if [ "$observed_peak" -gt 0 ]; then
        new_peak=$(awk -v old="$old_peak" -v obs="$observed_peak" -v a="$AUTOSCALER_EWMA_ALPHA" \
            'BEGIN { smoothed = old*(1-a) + obs*a;
                     if (obs > smoothed) smoothed = obs;
                     printf "%d", smoothed }')
    else
        new_peak="$old_peak"
    fi
    [ "$new_peak" -lt "$AUTOSCALER_RAM_PEAK_MIN" ] && new_peak="$AUTOSCALER_RAM_PEAK_MIN"
    {
        echo "ram_peak_per_worker_mib=$new_peak"
        echo "last_updated=$(date -u +%FT%TZ)"
        echo "last_observed_mib=$observed_peak"
    } > "$ewma_file"
    echo "$new_peak"
}

# v2_swap_pressure — read prior swap_used from autoscaler.swap-state,
# diff against current. Returns 1 if SwapUsed grew > threshold MIB in
# last poll, else 0. State file rewritten each call.
v2_swap_pressure() {
    local cur_swap_used="$1"
    local state="$QUEUE_DIR/autoscaler.swap-state"
    local prev_swap=0
    if [ -f "$state" ]; then
        prev_swap=$(grep -E '^swap_used_mib=' "$state" | tail -1 | cut -d= -f2)
        [ -z "$prev_swap" ] && prev_swap=0
    fi
    local delta=$((cur_swap_used - prev_swap))
    echo "swap_used_mib=$cur_swap_used" > "$state"
    if [ "$delta" -gt "$AUTOSCALER_SWAP_PRESSURE_DELTA_MIB" ]; then
        echo 1
    else
        echo 0
    fi
}

# v2_write_pool_advisory — emit advisory file the pool reads each iter.
# Format: KEY=VALUE lines. Pool can act on:
#   slow_emit=1|0       — when 1, pool should reduce concurrent emits
#   reason=ram_pressure|swap_active|none
v2_write_pool_advisory() {
    local slow="$1" reason="$2"
    {
        echo "# bend-autoscaler advisory — last update $(date -u +%FT%TZ)"
        echo "slow_emit=$slow"
        echo "reason=$reason"
    } > "$QUEUE_DIR/emit.advisory"
}

# v2_dlq_growth — count DLQ reason files mtime'd in last $1 seconds.
# Used as oversubscription signal: spike in DLQ → halve next plan.
v2_dlq_growth() {
    local poll="$1"
    local mins=$((poll / 60 + 1))
    find "$QUEUE_DIR/dlq" -maxdepth 1 -type f -name '*.reason' -mmin -"$mins" 2>/dev/null | wc -l
}

# v2_damp_change — limit per-tier rate of change to ±25 % per poll.
# Cold start (cur=0) ramps up to AUTOSCALER_RAMP_FROM_ZERO in one step.
v2_damp_change() {
    local cur="$1" plan="$2"
    if [ "$cur" = "0" ]; then
        if [ "$plan" -gt "$AUTOSCALER_RAMP_FROM_ZERO" ]; then
            echo "$AUTOSCALER_RAMP_FROM_ZERO"
        else
            echo "$plan"
        fi
        return
    fi
    local max_step=$((cur / 4))
    [ "$max_step" -lt 1 ] && max_step=1
    if [ "$plan" -gt "$cur" ]; then
        local up=$((cur + max_step))
        [ "$plan" -gt "$up" ] && plan="$up"
    elif [ "$plan" -lt "$cur" ]; then
        local dn=$((cur - max_step))
        [ "$plan" -lt "$dn" ] && plan="$dn"
    fi
    echo "$plan"
}

# v2_plan_caps — live-state planner.
# Args (positional):
#   1 vram_free_mib   2 ram_free_mib   3..7 ready_micro/small/medium/large/huge
#   8..12 cur_micro/small/medium/large/huge   13 peak_per_worker_mib
#   14 running_forks   15 dlq_growth   16 inflight_huge
#   17 ram_peak_per_worker_mib   18 ram_total_mib   19 swap_pressure
# Emits plan_<tier> N + diagnostics (headroom_workers, total_budget, dlq_damped,
# huge_active, ram_damped — 1 when RAM pressure halved plan).
v2_plan_caps() {
    local vram_free="$1" ram_free="$2"
    local r_micro="$3" r_small="$4" r_medium="$5" r_large="$6" r_huge="$7"
    local c_micro="$8" c_small="$9" c_medium="${10}" c_large="${11}" c_huge="${12}"
    local peak="${13}" running="${14}" dlq_growth="${15}" inflight_huge="${16:-0}"
    local ram_peak="${17:-$peak}" ram_total="${18:-0}" swap_pressure="${19:-0}"

    # Headroom workers from VRAM. Safety reserve protects against bursty
    # forks landing all at once + saturating GPU mid-poll-period.
    local vram_after_reserve=$((vram_free - AUTOSCALER_VRAM_SAFETY_RESERVE_MIB))
    [ "$vram_after_reserve" -lt 0 ] && vram_after_reserve=0
    local headroom_vram=$((vram_after_reserve / peak))
    # RAM bound — backend RSS = 1.5-2× VRAM typically (lib-tier.sh:48-49).
    # Use SEPARATE ram_peak EWMA + ram_safety_reserve so headroom_ram stays
    # honest. Old code shared VRAM peak which OVER-estimated RAM headroom
    # + edged into swap on heavy sweep loads.
    local ram_after_reserve=$((ram_free - AUTOSCALER_RAM_SAFETY_RESERVE_MIB))
    [ "$ram_after_reserve" -lt 0 ] && ram_after_reserve=0
    local headroom_ram=$((ram_after_reserve / ram_peak))
    local headroom=$headroom_vram
    [ "$headroom_ram" -lt "$headroom" ] && headroom=$headroom_ram

    # Total budget = workers we COULD have, given current load + headroom.
    local total_budget=$((running + headroom))
    # Cap at CPU ceiling (24 cores + modest oversub).
    [ "$total_budget" -gt "$AUTOSCALER_CPU_CEILING" ] && total_budget=$AUTOSCALER_CPU_CEILING

    # RAM-pressure damping — fires when MemAvailable < AUTOSCALER_RAM_PRESSURE_FRACTION
    # of total OR active swap (SwapUsed grew > delta MIB last poll). Either
    # signal means we sit at swap edge; halving total_budget reduces new
    # admissions so running workers can drain + RAM recovers before swap thrash.
    local ram_damped=0
    local ram_pressure_threshold=$((ram_total * 15 / 100))  # 15% default; matches RAM_PRESSURE_FRACTION
    if [ "$ram_total" -gt 0 ] && [ "$ram_free" -lt "$ram_pressure_threshold" ]; then
        total_budget=$((total_budget / 2))
        ram_damped=1
    fi
    if [ "$swap_pressure" -eq 1 ]; then
        total_budget=$((total_budget / 2))
        ram_damped=1
    fi

    # DLQ damping — surge of DLQ entries means we already over-subscribed
    # OR we have a bad-cell batch. Either way, halving caps next poll
    # protects against amplification while operator triages.
    local dlq_damped=0
    if [ "$dlq_growth" -gt "$AUTOSCALER_DLQ_DAMPING_THRESHOLD" ]; then
        total_budget=$((total_budget / 2))
        dlq_damped=1
    fi

    # HUGE solo-dispatch branch — must run before normal planning. Triggers:
    #   (a) r_huge>0  : huge bin queued, waiting for card to drain.
    #   (b) inflight_huge>0 : a huge already claimed by a dispatcher, still
    #                         executing. Other tiers must stay starved until
    #                         dispatcher releases this slot.
    # cur_huge alone is NOT a trigger — stale supervisor.config can hold
    # TIER_HUGE_DISPATCH=1 indefinitely with no bin actually running.
    local plan_micro=0 plan_small=0 plan_medium=0 plan_large=0 plan_huge=0
    local huge_active=0
    if [ "$r_huge" -gt 0 ] || [ "$inflight_huge" -gt 0 ]; then
        huge_active=1
        # Card-empty? Admit huge. "card-empty" signal is running=0
        # OR inflight_huge>0 (huge already running counts as "card claimed
        # by huge, keep slot open for it to continue").
        if [ "$running" -eq 0 ] || [ "$inflight_huge" -gt 0 ]; then
            plan_huge="$TIER_HUGE_DISPATCH_CEILING"
        fi
        # plan_others stays at 0 regardless — drain mode. Damping at the
        # tail rate-limits descent so already-running cells finish
        # naturally without abrupt kill.
        # Emit + return — skip normal multi-tier planning.
        plan_micro=$(v2_damp_change  "$c_micro"  "$plan_micro")
        plan_small=$(v2_damp_change  "$c_small"  "$plan_small")
        plan_medium=$(v2_damp_change "$c_medium" "$plan_medium")
        plan_large=$(v2_damp_change  "$c_large"  "$plan_large")
        plan_huge=$(v2_damp_change   "$c_huge"   "$plan_huge")
        printf 'plan_micro %d\nplan_small %d\nplan_medium %d\nplan_large %d\nplan_huge %d\n' \
            "$plan_micro" "$plan_small" "$plan_medium" "$plan_large" "$plan_huge"
        printf 'headroom_workers %d running_forks %d total_budget %d peak_per_worker_mib %d ram_peak_mib %d dlq_damped %d ram_damped %d huge_active %d\n' \
            "$headroom" "$running" "$total_budget" "$peak" "$ram_peak" "$dlq_damped" "$ram_damped" "$huge_active"
        return
    fi

    # Tier demand sentinels — V2 zero-floor (V1 forced MIN=1 even on empty).
    local nonzero=0
    [ "$r_micro"  -gt 0 ] && nonzero=$((nonzero + 1))
    [ "$r_small"  -gt 0 ] && nonzero=$((nonzero + 1))
    [ "$r_medium" -gt 0 ] && nonzero=$((nonzero + 1))
    [ "$r_large"  -gt 0 ] && nonzero=$((nonzero + 1))

    if [ "$nonzero" -le 1 ]; then
        # Single-tier (or no-demand) — grant full budget to one tier
        # with work. Capped at that tier's ceiling.
        if   [ "$r_large"  -gt 0 ]; then
            plan_large=$((total_budget < TIER_LARGE_DISPATCH_CEILING ? total_budget : TIER_LARGE_DISPATCH_CEILING))
        elif [ "$r_medium" -gt 0 ]; then
            plan_medium=$((total_budget < TIER_MEDIUM_DISPATCH_CEILING ? total_budget : TIER_MEDIUM_DISPATCH_CEILING))
        elif [ "$r_small"  -gt 0 ]; then
            plan_small=$((total_budget < TIER_SMALL_DISPATCH_CEILING ? total_budget : TIER_SMALL_DISPATCH_CEILING))
        elif [ "$r_micro"  -gt 0 ]; then
            plan_micro=$((total_budget < TIER_MICRO_DISPATCH_CEILING ? total_budget : TIER_MICRO_DISPATCH_CEILING))
        fi
    else
        # Multi-tier — MIN=1 to each tier with demand, then greedy fill
        # largest-tier-first (slow drains get priority on capacity).
        local remaining=$total_budget
        [ "$r_large"  -gt 0 ] && [ "$remaining" -gt 0 ] && { plan_large=1;  remaining=$((remaining - 1)); }
        [ "$r_medium" -gt 0 ] && [ "$remaining" -gt 0 ] && { plan_medium=1; remaining=$((remaining - 1)); }
        [ "$r_small"  -gt 0 ] && [ "$remaining" -gt 0 ] && { plan_small=1;  remaining=$((remaining - 1)); }
        [ "$r_micro"  -gt 0 ] && [ "$remaining" -gt 0 ] && { plan_micro=1;  remaining=$((remaining - 1)); }
        # Greedy: largest tier with demand and below ceiling consumes 1 slot per pass.
        while [ "$remaining" -gt 0 ]; do
            local advanced=0
            if [ "$plan_large"  -lt "$TIER_LARGE_DISPATCH_CEILING"  ] && [ "$r_large"  -gt "$plan_large"  ]; then
                plan_large=$((plan_large + 1));   remaining=$((remaining - 1)); advanced=1
                [ "$remaining" -eq 0 ] && break
            fi
            if [ "$plan_medium" -lt "$TIER_MEDIUM_DISPATCH_CEILING" ] && [ "$r_medium" -gt "$plan_medium" ]; then
                plan_medium=$((plan_medium + 1)); remaining=$((remaining - 1)); advanced=1
                [ "$remaining" -eq 0 ] && break
            fi
            if [ "$plan_small"  -lt "$TIER_SMALL_DISPATCH_CEILING"  ] && [ "$r_small"  -gt "$plan_small"  ]; then
                plan_small=$((plan_small + 1));   remaining=$((remaining - 1)); advanced=1
                [ "$remaining" -eq 0 ] && break
            fi
            if [ "$plan_micro"  -lt "$TIER_MICRO_DISPATCH_CEILING"  ] && [ "$r_micro"  -gt "$plan_micro"  ]; then
                plan_micro=$((plan_micro + 1));   remaining=$((remaining - 1)); advanced=1
                [ "$remaining" -eq 0 ] && break
            fi
            [ "$advanced" -eq 0 ] && break
        done
    fi

    # Rate-of-change damping — ±25 % per poll. Prevents oscillation when
    # demand or VRAM headroom stays bouncy.
    plan_micro=$(v2_damp_change  "$c_micro"  "$plan_micro")
    plan_small=$(v2_damp_change  "$c_small"  "$plan_small")
    plan_medium=$(v2_damp_change "$c_medium" "$plan_medium")
    plan_large=$(v2_damp_change  "$c_large"  "$plan_large")

    # Post-damp CPU ceiling — damping toward a HIGHER cur can keep total
    # above ceiling. Trim largest tiers first (cheapest to lose a slot at
    # slow-drain end). Repeat until within budget.
    plan_sum=$((plan_micro + plan_small + plan_medium + plan_large))
    while [ "$plan_sum" -gt "$AUTOSCALER_CPU_CEILING" ]; do
        if   [ "$plan_large"  -gt 0 ]; then plan_large=$((plan_large - 1))
        elif [ "$plan_medium" -gt 0 ]; then plan_medium=$((plan_medium - 1))
        elif [ "$plan_small"  -gt 0 ]; then plan_small=$((plan_small - 1))
        elif [ "$plan_micro"  -gt 0 ]; then plan_micro=$((plan_micro - 1))
        else break
        fi
        plan_sum=$((plan_micro + plan_small + plan_medium + plan_large))
    done

    # plan_huge always 0 in non-huge branch (handled above + returned early).
    printf 'plan_micro %d\nplan_small %d\nplan_medium %d\nplan_large %d\nplan_huge %d\n' \
        "$plan_micro" "$plan_small" "$plan_medium" "$plan_large" 0
    printf 'headroom_workers %d running_forks %d total_budget %d peak_per_worker_mib %d ram_peak_mib %d dlq_damped %d ram_damped %d huge_active %d\n' \
        "$headroom" "$running" "$total_budget" "$peak" "$ram_peak" "$dlq_damped" "$ram_damped" "$huge_active"
}

# ── exit handlers ────────────────────────────────────────────────
EXIT_REASON=""
on_exit() {
    local r="${EXIT_REASON:-unknown-exit}"
    heartbeat_exit_cause "$QUEUE_DIR" "autoscaler" "$r"
    echo "[$(date +%H:%M:%S)] bend-autoscaler exit reason=$r" | tee -a "$LOG"
}
trap on_exit EXIT
trap 'EXIT_REASON=signal-int; exit 130' INT
trap 'EXIT_REASON=signal-term; exit 143' TERM

log "bend-autoscaler start queue=$QUEUE_DIR poll=${POLL_SECONDS}s v2=$AUTOSCALER_V2"
if [ "$AUTOSCALER_V2" = "1" ]; then
    log "  V2 live model — peak_init=${AUTOSCALER_PEAK_INITIAL}MIB cpu_ceiling=$AUTOSCALER_CPU_CEILING dlq_threshold=$AUTOSCALER_DLQ_DAMPING_THRESHOLD"
else
    log "  V1 fixed-budget — RAM frac=$RAM_BUDGET_FRACTION  VRAM frac=$VRAM_BUDGET_FRACTION"
fi

# Cold-recovery baseline — persist KNOWN-GOOD caps so a cold restart of
# every component (autoscaler crash + supervisor restart with stale
# config) can fall back to a safe operating point. Written each poll
# after a successful plan.
BASELINE_FILE="$QUEUE_DIR/supervisor.config.baseline"

while true; do
    heartbeat_touch "$QUEUE_DIR" "autoscaler"
    if [ -f "$STOP_FILE" ]; then
        log "STOP marker — exiting"
        rm -f "$STOP_FILE"
        EXIT_REASON="stop-marker"
        break
    fi

    ready_out=$(count_ready_by_tier)
    ready_micro=$(echo "$ready_out" | awk '/^ready_micro/ {print $2}')
    ready_small=$(echo "$ready_out" | awk '/^ready_small/ {print $2}')
    ready_medium=$(echo "$ready_out" | awk '/^ready_medium/ {print $2}')
    ready_large=$(echo "$ready_out" | awk '/^ready_large/ {print $2}')
    ready_huge=$(echo  "$ready_out" | awk '/^ready_huge/  {print $2}')
    total_ready=$((ready_micro + ready_small + ready_medium + ready_large + ready_huge))

    cur_out=$(read_current_caps)
    cur_micro=$(echo "$cur_out" | awk '/^cur_micro/ {print $2}')
    cur_small=$(echo "$cur_out" | awk '/^cur_small/ {print $2}')
    cur_medium=$(echo "$cur_out" | awk '/^cur_medium/ {print $2}')
    cur_large=$(echo "$cur_out" | awk '/^cur_large/ {print $2}')
    cur_huge=$(echo  "$cur_out" | awk '/^cur_huge/  {print $2}')

    if [ "$AUTOSCALER_V2" = "1" ]; then
        live_out=$(v2_sample_live)
        vram_free_mib=$(echo "$live_out" | awk '/^vram_free_mib/ {print $2}')
        ram_free_mib=$(echo  "$live_out" | awk '/^ram_free_mib/  {print $2}')
        ram_total_mib=$(echo "$live_out" | awk '/^ram_total_mib/ {print $2}')
        swap_used_mib=$(echo "$live_out" | awk '/^swap_used_mib/ {print $2}')
        running_forks=$(echo "$live_out" | awk '/^running_forks/ {print $2}')
        observed_peak_mib=$(echo "$live_out" | awk '/^observed_peak_mib/ {print $2}')
        observed_ram_peak_mib=$(echo "$live_out" | awk '/^observed_ram_peak_mib/ {print $2}')
        inflight_huge=$(echo "$live_out" | awk '/^inflight_huge/ {print $2}')
        peak_per_worker_mib=$(v2_update_ewma "$observed_peak_mib")
        ram_peak_per_worker_mib=$(v2_update_ram_ewma "$observed_ram_peak_mib")
        dlq_growth=$(v2_dlq_growth "$POLL_SECONDS")
        swap_pressure=$(v2_swap_pressure "$swap_used_mib")

        plan_out=$(v2_plan_caps \
            "$vram_free_mib" "$ram_free_mib" \
            "$ready_micro" "$ready_small" "$ready_medium" "$ready_large" "$ready_huge" \
            "$cur_micro" "$cur_small" "$cur_medium" "$cur_large" "$cur_huge" \
            "$peak_per_worker_mib" "$running_forks" "$dlq_growth" "$inflight_huge" \
            "$ram_peak_per_worker_mib" "$ram_total_mib" "$swap_pressure")
        # Write pool advisory based on planner verdict — pool reads each iter.
        # ram_damped emitted by planner; reuse as advisory trigger.
        ram_pressure_active=$(echo "$plan_out" | awk 'match($0, /ram_damped [0-9]+/) {print substr($0, RSTART+11, RLENGTH-11)}')
        if [ "${ram_pressure_active:-0}" -eq 1 ]; then
            reason="ram_pressure"
            [ "$swap_pressure" -eq 1 ] && reason="swap_active"
            v2_write_pool_advisory 1 "$reason"
        else
            v2_write_pool_advisory 0 "none"
        fi
        diag=$(echo "$plan_out" | grep -E '^headroom_workers')
    else
        cap_out=$(read_capacity)
        total_ram_mib=$(echo "$cap_out" | awk '/^total_ram_mib/ {print $2}')
        total_vram_mib=$(echo "$cap_out" | awk '/^total_vram_mib/ {print $2}')
        ram_budget=$(awk -v t="$total_ram_mib" -v f="$RAM_BUDGET_FRACTION" 'BEGIN{printf "%d", t*f}')
        vram_budget=$(awk -v t="$total_vram_mib" -v f="$VRAM_BUDGET_FRACTION" 'BEGIN{printf "%d", t*f}')
        plan_out=$(plan_caps "$ram_budget" "$vram_budget" \
                             "$ready_micro" "$ready_small" "$ready_medium" "$ready_large")
        diag=""
    fi

    plan_micro=$(echo  "$plan_out" | awk '/^plan_micro/  {print $2}')
    plan_small=$(echo  "$plan_out" | awk '/^plan_small/  {print $2}')
    plan_medium=$(echo "$plan_out" | awk '/^plan_medium/ {print $2}')
    plan_large=$(echo  "$plan_out" | awk '/^plan_large/  {print $2}')
    # V1 doesn't emit plan_huge — treat as zero so write_caps still gets 5 args.
    plan_huge=$(echo   "$plan_out" | awk '/^plan_huge/   {print $2}')
    [ -z "$plan_huge" ] && plan_huge=0

    # Plan-change log — fires only on a transition. Operators read this to
    # see "what caused this dispatch shift".
    if [ "$cur_micro" != "$plan_micro" ] || [ "$cur_small" != "$plan_small" ] || \
       [ "$cur_medium" != "$plan_medium" ] || [ "$cur_large" != "$plan_large" ] || \
       [ "$cur_huge"  != "$plan_huge"  ]; then
        if [ -n "$diag" ]; then
            log "ready=$total_ready (μ$ready_micro s$ready_small m$ready_medium L$ready_large H$ready_huge) cur=(μ$cur_micro s$cur_small m$cur_medium L$cur_large H$cur_huge) -> plan=(μ$plan_micro s$plan_small m$plan_medium L$plan_large H$plan_huge) $diag"
        else
            log "ready=$total_ready (μ$ready_micro s$ready_small m$ready_medium L$ready_large H$ready_huge) cur=(μ$cur_micro s$cur_small m$cur_medium L$cur_large H$cur_huge) -> plan=(μ$plan_micro s$plan_small m$plan_medium L$plan_large H$plan_huge)"
        fi
        write_caps "$plan_micro" "$plan_small" "$plan_medium" "$plan_large" "$plan_huge"
    else
        # Stable-plan heartbeat — fires every AUTOSCALER_STABLE_HEARTBEAT_POLLS
        # polls (default 6 = 2 min at poll=20s) when plan does NOT change.
        # Without this, a long drain phase produces no log output, making
        # the autoscaler look frozen to operators reading autoscaler.log
        # while it actually polls silently.
        STABLE_POLLS=$((${STABLE_POLLS:-0} + 1))
        if [ "$STABLE_POLLS" -ge "${AUTOSCALER_STABLE_HEARTBEAT_POLLS:-6}" ]; then
            log "STABLE plan=(μ$plan_micro s$plan_small m$plan_medium L$plan_large H$plan_huge) $diag"
            STABLE_POLLS=0
        fi
    fi

    # Persist baseline for cold-recovery — only when plan looks healthy
    # (total >= AUTOSCALER_CPU_CEILING/4 to filter throttled states).
    # HUGE counts toward total though its cap stays solo-1.
    plan_total=$((plan_micro + plan_small + plan_medium + plan_large + plan_huge))
    if [ "$plan_total" -ge $((AUTOSCALER_CPU_CEILING / 4)) ]; then
        {
            echo "# bend-autoscaler last-known-good caps — written $(date -u +%FT%TZ)"
            echo "# Source for cold recovery: supervisor falls back to this on autoscaler stall."
            printf 'TIER_MICRO_DISPATCH=%d\nTIER_SMALL_DISPATCH=%d\nTIER_MEDIUM_DISPATCH=%d\nTIER_LARGE_DISPATCH=%d\nTIER_HUGE_DISPATCH=%d\n' \
                "$plan_micro" "$plan_small" "$plan_medium" "$plan_large" "$plan_huge"
        } > "$BASELINE_FILE"
    fi

    sleep "$POLL_SECONDS"
done
