#!/usr/bin/env bash
# Start doctate-server with the chosen ASR backend (whisper or canary).
#
# Coordinates the GPU host (minerva.lan) so the chosen ASR container is
# the only heavy GPU resident:
#   - Whisper (~1.5-3 GB) coexists with Ollama (~3-4 GB) on the RTX 3060.
#   - Canary (~9.7 GB) does NOT — Ollama must be unloaded first, else
#     Canary OOMs on its last few MB of model weights.
#   See docs/canary.md §2.6 for the VRAM accounting.
#
# Flow:
#   1. cargo build -p doctate-server  (fail fast, BEFORE touching minerva)
#   2. ssh minerva.lan: docker compose stop the OTHER ASR stack
#   3. (canary only) unload all Ollama models — must happen BEFORE the
#      canary container starts, otherwise canary loads its model into a
#      still-occupied GPU and OOMs in the last few MB. Order matters:
#      „start canary, then free VRAM" is too late — canary's model load
#      runs in the container background as soon as `compose up -d` returns.
#   4. ssh minerva.lan: docker compose up -d the chosen stack — for
#      canary, the GPU is now guaranteed empty.
#   5. poll the chosen backend until it answers, with progress pulses
#      and periodic container-log snapshots so a stuck start is visible
#   6. exec cargo run -p doctate-server with ASR_BACKEND= exported
#
# Usage:
#   scripts/run-server.sh whisper
#   scripts/run-server.sh canary
#
# Env overrides (all optional):
#   MINERVA_HOST           default minerva.lan
#   WHISPER_STACK          default /opt/stacks/doctate-whisper
#   CANARY_STACK           default /opt/stacks/doctate-canary
#   WHISPER_HOST_PORT      default 9001
#   CANARY_HOST_PORT       default 9002
#   OLLAMA_HOST_PORT       default 11434
#   HEALTH_TIMEOUT_SECS    default 600  (cold model-download room)
#   HEALTH_POLL_SECS       default 3    (interval between probes)
#   LOG_SNAPSHOT_SECS      default 30   (interval between log dumps)

set -euo pipefail

MINERVA_HOST="${MINERVA_HOST:-minerva.lan}"
WHISPER_STACK="${WHISPER_STACK:-/opt/stacks/doctate-whisper}"
CANARY_STACK="${CANARY_STACK:-/opt/stacks/doctate-canary}"
WHISPER_HOST_PORT="${WHISPER_HOST_PORT:-9001}"
CANARY_HOST_PORT="${CANARY_HOST_PORT:-9002}"
OLLAMA_HOST_PORT="${OLLAMA_HOST_PORT:-11434}"
HEALTH_TIMEOUT_SECS="${HEALTH_TIMEOUT_SECS:-600}"
HEALTH_POLL_SECS="${HEALTH_POLL_SECS:-3}"
LOG_SNAPSHOT_SECS="${LOG_SNAPSHOT_SECS:-30}"

usage() {
  cat <<EOF >&2
Usage: $0 <whisper|canary>

Coordinates the GPU host (\$MINERVA_HOST) so only the matching ASR
container runs (and, for canary, Ollama is unloaded), then starts
doctate-server locally with the matching ASR_BACKEND env var.
See header comment for env overrides.
EOF
  exit 64
}

[[ $# -eq 1 ]] || usage
case "$1" in
  whisper|canary) BACKEND="$1" ;;
  *) usage ;;
esac

REPO_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "$REPO_ROOT/server"

# ---- Step 1: build before touching minerva. -------------------------------
echo ">>> [1/6] cargo build -p doctate-server"
cargo build -p doctate-server

# ---- Step 2: pick stop/start targets. -------------------------------------
if [[ "$BACKEND" == "whisper" ]]; then
  STOP_STACK="$CANARY_STACK"
  START_STACK="$WHISPER_STACK"
  HEALTH_KIND="tcp"
  HEALTH_PORT="$WHISPER_HOST_PORT"
  HEALTH_DESC="TCP $MINERVA_HOST:$WHISPER_HOST_PORT"
else
  STOP_STACK="$WHISPER_STACK"
  START_STACK="$CANARY_STACK"
  HEALTH_KIND="http"
  HEALTH_URL="http://$MINERVA_HOST:$CANARY_HOST_PORT/health"
  HEALTH_DESC="HTTP $HEALTH_URL"
fi

echo ">>> [2/6] $MINERVA_HOST: stop $STOP_STACK"
ssh "$MINERVA_HOST" "cd '$STOP_STACK' && docker compose stop"

# ---- Step 3 (canary only): unload Ollama models. --------------------------
# Ollama keeps the active model warm for `keep_alive` seconds (300 in
# settings.toml at the time of writing). When the operator switches to
# canary in the middle of that window, the model is still resident and
# canary OOMs on the last few MB. We trigger a synchronous unload via
# the Ollama API: posting `keep_alive: 0` with an empty prompt drops
# the model immediately and frees VRAM. (See ollama/ollama issue #2546.)
unload_ollama_models() {
  local base="http://$MINERVA_HOST:$OLLAMA_HOST_PORT"
  local ps_file
  ps_file=$(mktemp)
  trap 'rm -f "$ps_file"' RETURN

  if ! curl -fsS --max-time 5 "$base/api/ps" -o "$ps_file"; then
    echo "    Ollama API unreachable at $base/api/ps — skipping unload" >&2
    return 0
  fi

  # Parse the JSON model list with python3 (avoid jq as a hard dep).
  local models
  models=$(python3 -c "
import json, sys
try:
    d = json.load(open('$ps_file'))
    for m in d.get('models', []):
        print(m.get('name', ''))
except Exception:
    sys.exit(0)
")
  if [[ -z "$models" ]]; then
    echo "    Ollama: no models loaded — VRAM already free"
    return 0
  fi

  while IFS= read -r name; do
    [[ -z "$name" ]] && continue
    echo "    unloading: $name"
    curl -fsS --max-time 10 "$base/api/generate" \
      -H 'Content-Type: application/json' \
      -d "{\"model\":\"$name\",\"prompt\":\"\",\"keep_alive\":0}" \
      -o /dev/null \
      || echo "      (warn: unload request for $name returned non-2xx)" >&2
  done <<< "$models"

  # Give the GPU driver a moment to actually reclaim the pages before
  # canary starts allocating its own.
  sleep 2
}

if [[ "$BACKEND" == "canary" ]]; then
  echo ">>> [3/6] freeing VRAM: unloading Ollama models BEFORE canary starts"
  unload_ollama_models
else
  echo ">>> [3/6] (whisper) Ollama coexistence is fine — skipping unload"
fi

# ---- Step 4: NOW start the chosen stack. ----------------------------------
# For canary this is the critical moment — the GPU has just been freed
# by step 3, and canary's container will allocate into an empty GPU
# instead of fighting Ollama for the last 20 MB.
echo ">>> [4/6] $MINERVA_HOST: start $START_STACK"
ssh "$MINERVA_HOST" "cd '$START_STACK' && docker compose up -d"

# ---- Step 5: poll until the chosen backend answers. -----------------------
echo ">>> [5/6] waiting for $HEALTH_DESC (up to ${HEALTH_TIMEOUT_SECS}s)"

probe_once() {
  if [[ "$HEALTH_KIND" == "tcp" ]]; then
    if (exec 3<>"/dev/tcp/$MINERVA_HOST/$HEALTH_PORT") 2>/dev/null; then
      exec 3<&- 3>&-
      return 0
    fi
    return 1
  else
    curl -fsS --max-time 5 "$HEALTH_URL" >/dev/null 2>&1
  fi
}

start=$(date +%s)
deadline=$(( start + HEALTH_TIMEOUT_SECS ))
last_log_snapshot=0
attempt=0
while :; do
  if probe_once; then
    echo ""
    echo "    backend is up after $(( $(date +%s) - start ))s"
    break
  fi

  attempt=$(( attempt + 1 ))
  printf '.'

  now=$(date +%s)
  if (( now >= deadline )); then
    echo ""
    echo "timeout waiting for $HEALTH_DESC after ${HEALTH_TIMEOUT_SECS}s" >&2
    echo "--- container logs (tail 40) ---" >&2
    ssh "$MINERVA_HOST" "cd '$START_STACK' && docker compose logs --tail 40" >&2 || true
    exit 1
  fi

  # Periodic log snapshot: every LOG_SNAPSHOT_SECS, print elapsed time
  # and the last few log lines so the operator can SEE if the container
  # is loading, crashing, or stalled.
  if (( now - last_log_snapshot >= LOG_SNAPSHOT_SECS )); then
    echo ""
    elapsed=$(( now - start ))
    echo "    [${elapsed}s elapsed] still waiting; recent container log:"
    ssh "$MINERVA_HOST" "cd '$START_STACK' && docker compose logs --tail 5 --no-log-prefix 2>&1" \
      | sed 's/^/      | /' \
      || echo "      (could not read logs)" >&2
    last_log_snapshot=$now
  fi

  sleep "$HEALTH_POLL_SECS"
done

# ---- Step 6: start the server. --------------------------------------------
# dotenvy::dotenv() in server/src/main.rs does NOT override env vars that
# are already set, so this export wins over any ASR_BACKEND= in .env.
echo ">>> [6/6] starting doctate-server (ASR_BACKEND=$BACKEND)"
export ASR_BACKEND="$BACKEND"
exec cargo run -p doctate-server