#!/usr/bin/env bash
# llm-cache-watchdog.sh — Detect poisoned KV cache and auto-restart primary LLM.
#
# Probes the primary model directly on a fixed interval. An empty response
# indicates oMLX has cached a bad state and is replaying it — the only fix is
# a process restart (no cache flush API available).
#
# Trigger threshold: TRIGGER_COUNT consecutive empty-response probes within a
# PROBE_INTERVAL cycle before a restart is issued.
#
# Run in the foreground or via a long-lived background process (e.g. screen/tmux).
# Cron is also fine for infrequent checks:
#   */2 * * * * /path/to/scripts/llm-cache-watchdog.sh --once

source "$(dirname "$0")/llm.conf"

PROBE_INTERVAL=30     # seconds between health probes
RESTART_COOLDOWN=120  # seconds before allowing another restart
TRIGGER_COUNT=2       # consecutive failures required before restart

last_restart=0
fail_count=0

log() { echo "$(date '+%Y-%m-%d %H:%M:%S') [cache-watchdog] $*" | tee -a "$LOG_DIR/cache-watchdog.log"; }

_probe_primary() {
  # Returns: 0 = healthy, 1 = bad/empty response, 2 = server not up
  local model_id
  model_id=$(curl -sf --max-time 5 "http://localhost:$PRIMARY_PORT/v1/models" | \
    python3 -c "import sys,json; print(json.load(sys.stdin)['data'][0]['id'])" 2>/dev/null)
  [[ -z "$model_id" ]] && return 2

  local content
  content=$(curl -sf --max-time 30 "http://localhost:$PRIMARY_PORT/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d "{\"model\":\"$model_id\",\"messages\":[{\"role\":\"user\",\"content\":\"Say OK\"}],\"max_tokens\":5,\"temperature\":0,\"stream\":false}" | \
    python3 -c "import sys,json; print(json.load(sys.stdin)['choices'][0]['message'].get('content',''))" 2>/dev/null)

  [[ -z "$content" ]] && return 1
  return 0
}

restart_primary() {
  log "CACHE POISON DETECTED — restarting primary LLM to clear KV cache"
  local pid
  pid=$(cat "$PRIMARY_PID" 2>/dev/null)
  if [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null; then
    kill "$pid"
    for i in $(seq 1 15); do kill -0 "$pid" 2>/dev/null || break; sleep 1; done
    kill -0 "$pid" 2>/dev/null && kill -9 "$pid"
  fi
  rm -f "$PRIMARY_PID"

  nohup "$RMLX_BIN" serve "$PRIMARY_MODEL" \
    --port "$PRIMARY_PORT" \
    --enable-auto-tool-choice \
    --tool-call-parser "$PRIMARY_TOOL_PARSER" \
    --reasoning-parser "$PRIMARY_REASONING_PARSER" \
    --kv-cache-quantization \
    --kv-cache-quantization-bits 4 \
    --kv-cache-quantization-group-size 64 \
    --enable-prefix-cache \
    --pin-system-prompt \
    --gpu-memory-utilization 0.60 \
    --cache-memory-percent 0.20 \
    --max-tokens 65536 \
    --chunked-prefill-tokens 4096 \
    --prefill-step-size 4096 \
    --max-num-seqs 4 \
    --gc-control \
    --default-temperature 0.6 \
    --default-top-p 0.95 \
    --default-top-k 64 \
    >> "$PRIMARY_LOG" 2>&1 &
  echo $! > "$PRIMARY_PID"
  log "Primary LLM restarted (PID $(cat $PRIMARY_PID))"
  last_restart=$(date +%s)
  fail_count=0
}

# ── --once mode: single probe, exit 0 if healthy, 1 if restarted ─────────────
if [[ "${1:-}" == "--once" ]]; then
  _probe_primary
  result=$?
  if [[ $result -eq 2 ]]; then log "Server not up — skipping"; exit 0; fi
  if [[ $result -eq 0 ]]; then log "Healthy"; exit 0; fi
  log "Bad response on single probe — restarting"
  restart_primary
  exit 0
fi

# ── Continuous loop ───────────────────────────────────────────────────────────
log "Watchdog started — probing :$PRIMARY_PORT every ${PROBE_INTERVAL}s"

while true; do
  sleep "$PROBE_INTERVAL"
  now=$(date +%s)

  if (( now - last_restart < RESTART_COOLDOWN )); then
    continue
  fi

  _probe_primary
  probe_result=$?

  if [[ $probe_result -eq 2 ]]; then
    log "Server not responding on :$PRIMARY_PORT — skipping probe"
    fail_count=0
    continue
  fi

  if [[ $probe_result -eq 0 ]]; then
    fail_count=0
    continue
  fi

  fail_count=$((fail_count + 1))
  log "Empty/bad response from probe ($fail_count/$TRIGGER_COUNT)"

  if (( fail_count >= TRIGGER_COUNT )); then
    restart_primary
  fi
done