#!/usr/bin/env zsh
# llm-start.sh — Bring up the LLM inference stack.
# Manual start only. No autostart. "Lightswitch ON."
#
# Components started (in order):
#   1. oMLX 0.4.1     — multi-model server :8080 (all models in MODEL_DIR)
#   2. oMLX proxy     — :8090 (reasoning_content routing for Gemma4 thinking)
#
# oMLX 0.4.1 is a multi-model LRU server. Gemma4 26B-A4B loads as
# VLMBatchedEngine and handles vision, tools, and thinking from :8080.
# No separate vision sidecar needed.

source "$(dirname "$0")/llm.conf"
mkdir -p "$LOG_DIR" "$RUN_DIR" "$CACHE_DIR"

# Version guard — patches are specific to OMLX_VERSION in llm.conf
RUNNING_OMLX_VERSION=$(brew list --versions omlx 2>/dev/null | awk '{print $2}')
if [[ -n "$OMLX_VERSION" && "$RUNNING_OMLX_VERSION" != "$OMLX_VERSION" ]]; then
  echo "❌ oMLX version mismatch: installed=${RUNNING_OMLX_VERSION} required=${OMLX_VERSION}"
  echo "   VLM patches (model_discovery.py, vision_config) are version-specific."
  echo "   Run setup.sh to re-apply patches or update OMLX_VERSION in llm.conf."
  exit 1
fi

_is_running() {
  local pid_file=$1
  [[ -f "$pid_file" ]] && kill -0 "$(cat "$pid_file")" 2>/dev/null
}

_port_in_use() {
  lsof -ti tcp:$1 >/dev/null 2>&1
}

_wait_ready() {
  local port=$1 label=$2 max_wait=${3:-60}
  printf "  Waiting for %s on :%s " "$label" "$port"
  for i in $(seq 1 $((max_wait / 2))); do
    if curl -sf "http://localhost:$port/v1/models" >/dev/null 2>&1; then
      echo " ready"
      return 0
    fi
    printf "."
    sleep 2
  done
  echo " TIMEOUT — check log"
  return 1
}

echo "╔══════════════════════════════════════════════╗"
echo "║  LLM Stack — Starting                       ║"
echo "╚══════════════════════════════════════════════╝"
echo ""

# ── 1. oMLX 0.4.1 multi-model server ─────────────────────────────────────────
#
# Configuration rationale (M4 Pro 48 GB, gemma-4-26b-a4b-it-4bit @ 15.3 GB):
#   --memory-guard-gb 34         → 34 GB ceiling (soft=30.6 GB, hard=32.3 GB)
#                                  Raised from 30 GB to eliminate adaptive_prefill_throttle on
#                                  two concurrent 16K-token sessions (peak observed: 30.3 GB).
#                                  Leaves 14 GB for OS + other processes on 48 GB system.
#   --sse-keepalive-mode off     → no SSE keepalive chunks (avoids id mismatch)
#   --max-concurrent-requests 2  → single user, saves memory for larger context
#   --hot-cache-max-size 20GB    → KV blocks stay hot in RAM on 48 GB system
#   --initial-cache-blocks 2048  → pre-allocate cache blocks, reduces latency spikes
#   --paged-ssd-cache-dir        → prefix cache: skip prefill for repeated system prompts
#                                  (opencode sends same ~3K prompt every turn — huge win)
#   --no-hf-cache                → belt-and-suspenders with hf_cache_enabled=false in settings
#
# Per-model settings in ~/.omlx/model_settings.json:
#   max_context_window: 262144   (256K native context)
#   turboquant_kv_enabled: true  (4-bit KV — ~3 GB for 256K, vs ~12 GB fp16)
#   turboquant_kv_bits: 4.0
#   turboquant_skip_last: true   (avoid corruption on last KV layer)
#   is_pinned: true              (keep weights resident in unified memory)
#
# hot_cache_only: true + ssd_cache_dir both set in settings.json:
#   hot_cache_only keeps KV blocks in RAM (no SSD spill for KV)
#   ssd_cache_dir enables PREFIX cache on SSD (processed prefill state reuse)
#   These are orthogonal — hot_cache_only = KV blocks; paged SSD = prefix states
#
if _is_running "$PRIMARY_PID" || _port_in_use "$PRIMARY_PORT"; then
  echo "omlx:    already running on :$PRIMARY_PORT"
else
  echo "omlx:    starting multi-model server on :$PRIMARY_PORT (models: $MODEL_DIR)"
  nohup "$RMLX_BIN" serve \
    --model-dir "$MODEL_DIR" \
    --port "$PRIMARY_PORT" \
    --memory-guard-gb 34 \
    --sse-keepalive-mode off \
    --max-concurrent-requests 2 \
    --hot-cache-max-size 6GB \
    --initial-cache-blocks 1024 \
    --paged-ssd-cache-dir "$HOME/.llm/prefix-cache" \
    --no-hf-cache \
    >> "$PRIMARY_LOG" 2>&1 &
  echo $! > "$PRIMARY_PID"
  echo "  PID $(cat "$PRIMARY_PID") — log: $PRIMARY_LOG"
  _wait_ready "$PRIMARY_PORT" "omlx" 120
fi

# ── 2. oMLX proxy — :8090 → :8080 ──────────────────────────────────────────
# Routes reasoning_content → content in streaming deltas (Gemma4 thinking mode).
# Routes reasoning_content → content in stream deltas for all models.
# oMLX 0.4.1 handles streaming+tools natively.
# Requires: $HOME/.llm/proxy-venv (pip install aiohttp)
if _is_running "$PROXY_PID" || _port_in_use "$PROXY_PORT"; then
  echo "proxy:     already running on :$PROXY_PORT"
else
  echo "proxy:     starting on :$PROXY_PORT"
  nohup "$PROXY_PYTHON" "$(dirname "$0")/omlx-proxy.py" \
    >> "$PROXY_LOG" 2>&1 &
  echo $! > "$PROXY_PID"
  echo "  PID $(cat "$PROXY_PID") — log: $PROXY_LOG"
  printf "  Waiting for proxy on :$PROXY_PORT "
  for i in $(seq 1 10); do
    curl -sf "http://localhost:$PROXY_PORT/v1/models" >/dev/null 2>&1 && echo " ready" && break
    printf "."
    sleep 1
  done
fi

# ── Summary ───────────────────────────────────────────────────────────────────
echo ""
echo "╔══════════════════════════════════════════════╗"
echo "║  Stack UP                                    ║"
echo "╠══════════════════════════════════════════════╣"
printf "║  omlx      http://localhost:%-5s/v1        ║\n" "$PRIMARY_PORT"
printf "║  proxy     http://localhost:%-5s/v1        ║\n" "$PROXY_PORT"
echo "╚══════════════════════════════════════════════╝"