#!/bin/bash
# pre-server.sh — Headless PRE server launcher (for LaunchAgent auto-start)
#
# Ensures Ollama is running, pre-warms the model into GPU memory,
# then exec's the web server so launchd can manage the process directly.
#
# Usage:
#   ./pre-server.sh              Start server (foreground, for launchd)
#   ./pre-server.sh --status     Check if server is running
#   ./pre-server.sh --stop       Stop the server

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PORT="${PRE_PORT:-11434}"
WEB_PORT="${PRE_WEB_PORT:-7749}"

# Read configured context window from ~/.pre/context
MODEL_CTX=131072
CTX_FILE="$HOME/.pre/context"
if [ -f "$CTX_FILE" ]; then
    read -r _ctx < "$CTX_FILE" 2>/dev/null
    _ctx="${_ctx//[!0-9]/}"
    if [ -n "$_ctx" ] && [ "$_ctx" -ge 2048 ] 2>/dev/null && [ "$_ctx" -le 262144 ] 2>/dev/null; then
        MODEL_CTX="$_ctx"
    fi
fi

# Status check
if [ "${1:-}" = "--status" ]; then
    if curl -sf "http://localhost:${WEB_PORT}/api/status" >/dev/null 2>&1; then
        _status_json=$(curl -sf "http://localhost:${WEB_PORT}/api/status" 2>/dev/null || true)
        _status_provider=$(echo "$_status_json" | grep -o '"provider"[[:space:]]*:[[:space:]]*"[^"]*"' | sed 's/.*"\([^"]*\)"$/\1/' 2>/dev/null || echo "ollama")
        echo "PRE server is running on port $WEB_PORT (provider: ${_status_provider})"
        if [ "$_status_provider" = "ollama" ] && ollama ps 2>/dev/null | grep -q "pre-gemma4"; then
            MEM=$(ollama ps 2>/dev/null | grep "pre-gemma4" | awk '{print $3, $4}')
            echo "Model loaded (${MEM}, context: ${MODEL_CTX})"
        fi
        exit 0
    else
        echo "PRE server is not running"
        exit 1
    fi
fi

# Stop
if [ "${1:-}" = "--stop" ]; then
    IS_LINUX=false
    [ "$(uname -s)" = "Linux" ] && IS_LINUX=true

    if [ "$IS_LINUX" = true ]; then
        # Linux: try systemd user service first, then direct kill
        if systemctl --user is-active pre-server &>/dev/null 2>&1; then
            systemctl --user stop pre-server
            echo "PRE server stopped (systemd user service)"
            echo "To re-enable: systemctl --user start pre-server"
        else
            PID=$(lsof -ti :"$WEB_PORT" -sTCP:LISTEN 2>/dev/null | head -1 || ss -tlnp "sport = :$WEB_PORT" 2>/dev/null | grep -oP 'pid=\K\d+' | head -1)
            if [ -n "$PID" ]; then
                kill "$PID" 2>/dev/null
                echo "PRE server stopped (PID $PID)"
            else
                echo "PRE server is not running"
            fi
        fi
    else
        # macOS: LaunchAgent or direct kill
        PLIST="$HOME/Library/LaunchAgents/com.pre.server.plist"
        if [ -f "$PLIST" ]; then
            launchctl unload "$PLIST" 2>/dev/null || true
            echo "PRE server stopped (LaunchAgent unloaded)"
            echo "To re-enable: launchctl load $PLIST"
        else
            PID=$(lsof -ti :"$WEB_PORT" -sTCP:LISTEN 2>/dev/null | head -1)
            if [ -n "$PID" ]; then
                kill "$PID" 2>/dev/null
                echo "PRE server stopped (PID $PID)"
            else
                echo "PRE server is not running"
            fi
        fi
    fi
    exit 0
fi

# --- Startup sequence ---

# Detect configured inference provider from ~/.pre/connections.json
PROVIDER="ollama"
CONN_FILE="$HOME/.pre/connections.json"
if [ -f "$CONN_FILE" ]; then
    # Extract _provider.active using lightweight JSON parsing (no jq dependency)
    _prov=$(grep -o '"active"[[:space:]]*:[[:space:]]*"[^"]*"' "$CONN_FILE" 2>/dev/null | head -1 | sed 's/.*"active"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/')
    [ -n "$_prov" ] && PROVIDER="$_prov"
fi

# Compute dynamic OLLAMA_NUM_PARALLEL: match inference-queue slot count.
# 2 slots when headroom >= 16GB (same heuristic as inference-queue.js).
TOTAL_MEM_GB=0
if [ "$(uname -s)" = "Linux" ]; then
    TOTAL_MEM_GB=$(($(grep MemTotal /proc/meminfo | awk '{print $2}') / 1048576))
else
    TOTAL_MEM_GB=$(($(sysctl -n hw.memsize 2>/dev/null || echo 0) / 1073741824))
fi
# Estimate model size: q8_0 ~28GB on 48GB+ systems, q4_K_M ~15GB otherwise
MODEL_EST_GB=15
[ "$TOTAL_MEM_GB" -ge 48 ] 2>/dev/null && MODEL_EST_GB=28
HEADROOM_GB=$((TOTAL_MEM_GB - MODEL_EST_GB - 8))
if [ "$HEADROOM_GB" -ge 16 ] 2>/dev/null; then
    OLLAMA_NUM_PARALLEL_VAL=2
else
    OLLAMA_NUM_PARALLEL_VAL=1
fi

# Set Ollama env vars. On macOS, `open -g -a Ollama` launches Ollama.app which
# does NOT inherit shell exports — it inherits from the user's launchd session.
# Use `launchctl setenv` to inject env vars into launchd so Ollama.app sees them.
export OLLAMA_KEEP_ALIVE=24h
export OLLAMA_MAX_LOADED_MODELS=2
export OLLAMA_NUM_PARALLEL="$OLLAMA_NUM_PARALLEL_VAL"

if [ "$(uname -s)" != "Linux" ]; then
    launchctl setenv OLLAMA_KEEP_ALIVE "24h" 2>/dev/null || true
    launchctl setenv OLLAMA_MAX_LOADED_MODELS "2" 2>/dev/null || true
    launchctl setenv OLLAMA_NUM_PARALLEL "$OLLAMA_NUM_PARALLEL_VAL" 2>/dev/null || true
fi

# Detect GPU backend and set Flash Attention / KV cache accordingly
# macOS: NVIDIA eGPU via TinyGPU (Docker nvidia-smi) or Apple Silicon (Metal)
# Linux: direct nvidia-smi (no Docker needed)
GPU_BACKEND=""
DETECTED_VRAM_GB=0

if [ "$(uname -s)" = "Linux" ]; then
    # Linux: direct NVIDIA detection
    if command -v nvidia-smi &>/dev/null; then
        NVRAM_MIB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 || true)
        if [ -n "$NVRAM_MIB" ] && [ "$NVRAM_MIB" -gt 0 ] 2>/dev/null; then
            GPU_BACKEND="cuda"
            DETECTED_VRAM_GB=$((NVRAM_MIB / 1024))
        fi
    fi
else
    # macOS: Apple Silicon (Metal) or NVIDIA eGPU (CUDA via TinyGPU)
    GPU_BACKEND="metal"
    if systemextensionsctl list 2>/dev/null | grep -qi "tinygpu"; then
        if command -v docker &>/dev/null && docker info &>/dev/null 2>&1; then
            EGPU_VRAM_MIB=$(docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu24.04 \
                nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 || true)
            if [ -n "$EGPU_VRAM_MIB" ] && [ "$EGPU_VRAM_MIB" -gt 0 ] 2>/dev/null; then
                GPU_BACKEND="cuda"
                DETECTED_VRAM_GB=$((EGPU_VRAM_MIB / 1024))
            fi
        fi
    fi
fi

if [ "$GPU_BACKEND" = "cuda" ]; then
    export OLLAMA_FLASH_ATTENTION=1
    if [ "$DETECTED_VRAM_GB" -ge 28 ]; then
        export OLLAMA_KV_CACHE_TYPE=q8_0
    else
        export OLLAMA_KV_CACHE_TYPE=q4_0
    fi
    export OLLAMA_GPU_OVERHEAD=256000000
elif [ "$GPU_BACKEND" = "metal" ]; then
    export OLLAMA_FLASH_ATTENTION=0
fi

# Inject FA and KV cache settings into the launchd session so Ollama.app
# picks them up when launched via `open -g -a Ollama`.  Ollama.app is spawned
# by launchd (not as a child of this script), so it only inherits env vars
# that have been pushed into the launchd session via launchctl setenv.
if [ "$(uname -s)" != "Linux" ]; then
    [ -n "${OLLAMA_FLASH_ATTENTION+x}" ] && launchctl setenv OLLAMA_FLASH_ATTENTION "$OLLAMA_FLASH_ATTENTION" 2>/dev/null || true
    [ -n "${OLLAMA_KV_CACHE_TYPE+x}" ]   && launchctl setenv OLLAMA_KV_CACHE_TYPE   "$OLLAMA_KV_CACHE_TYPE"   2>/dev/null || true
    [ -n "${OLLAMA_GPU_OVERHEAD+x}" ]    && launchctl setenv OLLAMA_GPU_OVERHEAD    "$OLLAMA_GPU_OVERHEAD"    2>/dev/null || true
fi

# Ensure Ollama is running (skip for remote-only providers like openai/azure/anthropic)
NEED_OLLAMA=false
if [ "$PROVIDER" = "ollama" ]; then
    NEED_OLLAMA=true
elif [ "$PROVIDER" = "llamacpp" ]; then
    # llamacpp still benefits from Ollama for embeddings (nomic-embed-text)
    NEED_OLLAMA=true
fi

if [ "$NEED_OLLAMA" = true ]; then
    if ! curl -sf "http://localhost:${PORT}/v1/models" >/dev/null 2>&1; then
        if [ "$(uname -s)" = "Linux" ]; then
            # Linux: try systemd service first, then manual
            if systemctl start ollama &>/dev/null 2>&1; then
                : # started via systemd
            else
                ollama serve >/dev/null 2>&1 &
            fi
        elif [ -d "/Applications/Ollama.app" ]; then
            # macOS: prefer Ollama.app (handles updates, UI icon)
            open -g -a Ollama 2>/dev/null || true
        else
            ollama serve >/dev/null 2>&1 &
        fi
        # Wait for Ollama API
        for i in $(seq 1 30); do
            if curl -sf "http://localhost:${PORT}/v1/models" >/dev/null 2>&1; then break; fi
            sleep 1
        done
        if ! curl -sf "http://localhost:${PORT}/v1/models" >/dev/null 2>&1; then
            if [ "$PROVIDER" = "ollama" ]; then
                echo "ERROR: Ollama failed to start within 30 seconds" >&2
                exit 1
            else
                echo "WARNING: Ollama not available (embeddings may be limited)" >&2
            fi
        fi
    fi

    # Pre-warm model into GPU memory (only for Ollama-primary provider)
    if [ "$PROVIDER" = "ollama" ]; then
        if ! ollama ps 2>/dev/null | grep -q "pre-gemma4"; then
            curl -sf --max-time 300 "http://localhost:${PORT}/api/chat" \
                -d "{\"model\":\"pre-gemma4\",\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"keep_alive\":\"24h\",\"options\":{\"num_predict\":1,\"num_ctx\":${MODEL_CTX}}}" \
                >/dev/null 2>&1 || true
        fi
    fi
fi

# Start web server — exec replaces this process so launchd manages it directly
cd "$SCRIPT_DIR"
exec node server.js
