#!/usr/bin/env zsh
# setup.sh — Install and configure the local LLM stack.
# Idempotent: safe to run multiple times.
# No sudo required after initial Homebrew setup.
set -euo pipefail

SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
# shellcheck source=llm.conf
source "$SCRIPT_DIR/llm.conf"

_ok()   { printf "  \033[32m✅\033[0m  %s\n" "$*"; }
_info() { printf "  →  %s\n" "$*"; }
_warn() { printf "  \033[33m⚠️\033[0m   %s\n" "$*"; }
_fail() { printf "  \033[31m❌\033[0m  %s\n" "$*" >&2; exit 1; }

echo ""
echo "╔══════════════════════════════════════════════╗"
echo "║  Local LLM Stack — Setup                    ║"
echo "╚══════════════════════════════════════════════╝"
echo ""

# ── Homebrew ──────────────────────────────────────────────────────────────────
command -v brew >/dev/null 2>&1 || _fail "Homebrew not found. Install from https://brew.sh"
_ok "Homebrew: $(brew --version | head -1)"

# ── omlx — version-pinned install ────────────────────────────────────────────
# All patches (model_discovery.py, config.json vision_config) are specific to
# oMLX 0.4.1. brew upgrade omlx WILL silently break VLM loading.
REQUIRED_OMLX_VERSION="0.4.1"

INSTALLED_OMLX_VERSION=$(brew list --versions omlx 2>/dev/null | awk '{print $2}')
if [[ -z "$INSTALLED_OMLX_VERSION" ]]; then
  _info "Installing omlx ${REQUIRED_OMLX_VERSION}..."
  brew tap jundot/omlx 2>/dev/null
  brew install jundot/omlx/omlx
  INSTALLED_OMLX_VERSION=$(brew list --versions omlx 2>/dev/null | awk '{print $2}')
fi

if [[ "$INSTALLED_OMLX_VERSION" != "$REQUIRED_OMLX_VERSION" ]]; then
  _fail "omlx version mismatch: installed=${INSTALLED_OMLX_VERSION} required=${REQUIRED_OMLX_VERSION}
  The VLM patches (model_discovery.py, config.json vision_config) are specific to ${REQUIRED_OMLX_VERSION}.
  To downgrade: brew uninstall omlx && brew install jundot/omlx/omlx@${REQUIRED_OMLX_VERSION}
  To upgrade: re-validate all three capability tests after running setup.sh against the new version."
fi

brew pin omlx 2>/dev/null && _ok "omlx ${INSTALLED_OMLX_VERSION} installed and pinned (brew upgrade omlx is blocked)"

# Verify rapid-mlx binary (version-pinned in llm.conf — stale after brew upgrade omlx)
if [[ ! -x "$RMLX_BIN" ]]; then
  FOUND=$(find /opt/homebrew/Cellar/omlx -name rapid-mlx -type f 2>/dev/null | sort | tail -1)
  if [[ -n "$FOUND" ]]; then
    _warn "RMLX_BIN in llm.conf appears stale."
    _warn "  Configured: $RMLX_BIN"
    _warn "  Found:      $FOUND"
    _warn "  Update RMLX_BIN in $SCRIPT_DIR/llm.conf"
  else
    _fail "rapid-mlx binary not found at $RMLX_BIN (omlx installed but binary missing?)"
  fi
else
  _ok "rapid-mlx binary OK: $RMLX_BIN"
fi

# ── Runtime dirs ─────────────────────────────────────────────────────────────
for dir in "$LLM_BASE" "$MODEL_DIR" "$LOG_DIR" "$RUN_DIR" "$CACHE_DIR"; do
  mkdir -p "$dir"
done
_ok "Runtime dirs: $LLM_BASE/{models,logs,run,cache}"

# ── Proxy Python venv ─────────────────────────────────────────────────────────
VENV_DIR=$LLM_BASE/proxy-venv

if [[ -x "$VENV_DIR/bin/python3" ]]; then
  _ok "proxy-venv already exists: $VENV_DIR"
else
  _info "Creating proxy venv at $VENV_DIR..."
  # Use Homebrew Python explicitly to avoid PIP_REQUIRE_VIRTUALENV conflicts
  /opt/homebrew/bin/python3 -m venv "$VENV_DIR"
  _ok "proxy-venv created"
fi

if "$VENV_DIR/bin/python3" -c "import aiohttp" 2>/dev/null; then
  VER=$("$VENV_DIR/bin/python3" -c "import aiohttp; print(aiohttp.__version__)")
  _ok "aiohttp $VER already installed in proxy-venv"
else
  _info "Installing aiohttp into proxy-venv..."
  "$VENV_DIR/bin/pip" install --quiet aiohttp
  VER=$("$VENV_DIR/bin/python3" -c "import aiohttp; print(aiohttp.__version__)")
  _ok "aiohttp $VER installed"
fi

# ── opencode integration ───────────────────────────────────────────────────────
echo ""
echo "── opencode ─────────────────────────────────────────────────────────────"

OPENCODE_CFG_DIR="$HOME/.config/opencode"
OPENCODE_CFG="$OPENCODE_CFG_DIR/opencode.json"

mkdir -p "$OPENCODE_CFG_DIR"

# Install @ai-sdk/openai-compatible into the opencode config dir
if [[ -f "$OPENCODE_CFG_DIR/node_modules/@ai-sdk/openai-compatible/package.json" ]]; then
  _ok "@ai-sdk/openai-compatible already installed"
else
  _info "Installing @ai-sdk/openai-compatible (opencode local provider adapter)..."
  if command -v npm >/dev/null 2>&1; then
    (cd "$OPENCODE_CFG_DIR" && npm install @ai-sdk/openai-compatible --save --silent 2>/dev/null) \
      && _ok "@ai-sdk/openai-compatible installed" \
      || _warn "npm install failed — install manually: cd ~/.config/opencode && npm install @ai-sdk/openai-compatible"
  else
    _warn "npm not found. Install Node.js then run: cd ~/.config/opencode && npm install @ai-sdk/openai-compatible"
  fi
fi

# Install MCP servers
for pkg in "@modelcontextprotocol/server-github" "tavily-mcp" "@playwright/mcp"; do
  pkg_dir="$OPENCODE_CFG_DIR/node_modules/${pkg}"
  if [[ -d "$pkg_dir" ]]; then
    _ok "$pkg already installed"
  else
    _info "Installing $pkg..."
    if command -v npm >/dev/null 2>&1; then
      (cd "$OPENCODE_CFG_DIR" && npm install "$pkg" --save --silent 2>/dev/null) \
        && _ok "$pkg installed" \
        || _warn "npm install failed for $pkg — install manually: cd ~/.config/opencode && npm install $pkg"
    else
      _warn "npm not found — skipping $pkg"
    fi
  fi
done

if [[ -f "$OPENCODE_CFG" ]]; then
  _ok "opencode.json already exists: $OPENCODE_CFG"
  echo "     (To regenerate: rm $OPENCODE_CFG && re-run setup.sh)"
else
  _info "Writing opencode.json..."
  OPENCODE_CFG_TMP=$(mktemp)
  cat > "$OPENCODE_CFG_TMP" << 'OPENCODE_CONF'
{
  "$schema": "https://opencode.ai/config.json",
  "model": "omlx/gemma-4-26b-a4b-it-4bit",
  "provider": {
    "omlx": {
      "npm": "@ai-sdk/openai-compatible",
      "name": "oMLX (local)",
      "options": {
        "baseURL": "http://localhost:8090/v1",
        "apiKey": "omlx-local",
        "headerTimeout": false,
        "timeout": 600000,
        "chunkTimeout": 600000
      },
      "models": {
        "gemma-4-26b-a4b-it-4bit": {
          "name": "Gemma 4 26B-A4B (local)",
          "tool_call": true,
          "modalities": {
            "input": ["text", "image"],
            "output": ["text"]
          },
          "limit": {
            "context": 262144,
            "output": 16384
          }
        }
      }
    }
  },
  "compaction": {
    "auto": true
  },
  "permission": {
    "*": "allow",
    "doom_loop": "deny"
  },
  "agent": {
    "build": {
      "temperature": 0.3,
      "steps": 30
    },
    "plan": {
      "temperature": 0.1
    },
    "title": {
      "temperature": 0.1
    }
  },
  "experimental": {
    "mcp_timeout": 180000
  },
  "mcp": {
    "playwright": {
      "type": "local",
      "command": ["node", "$OPENCODE_CFG_DIR/node_modules/@playwright/mcp/cli.js"],
      "enabled": true
    },
    "tavily": {
      "type": "local",
      "command": ["node", "$OPENCODE_CFG_DIR/node_modules/tavily-mcp/build/index.js"],
      "enabled": true
    },
    "github": {
      "type": "local",
      "command": ["node", "$OPENCODE_CFG_DIR/node_modules/@modelcontextprotocol/server-github/dist/index.js"],
      "enabled": true,
      "env": {
        "GITHUB_PERSONAL_ACCESS_TOKEN": ""
      }
    }
  }
}
OPENCODE_CONF
  mv "$OPENCODE_CFG_TMP" "$OPENCODE_CFG"
  _ok "opencode.json written to $OPENCODE_CFG"
fi

# ── Claude Code integration ────────────────────────────────────────────────────
echo ""
echo "── Claude Code ──────────────────────────────────────────────────────────"

CLAUDE_ENV_FILE="$HOME/.config/local-llm/claude-code-env.sh"
mkdir -p "$(dirname "$CLAUDE_ENV_FILE")"

if [[ -f "$CLAUDE_ENV_FILE" ]]; then
  _ok "Claude Code env file already exists: $CLAUDE_ENV_FILE"
else
  _info "Writing Claude Code env file..."
  cat > "$CLAUDE_ENV_FILE" << 'CLAUDE_ENV'
# Source this file to route Claude Code through the local oMLX stack.
# Add to ~/.zshrc or ~/.bashrc, or run: source ~/.config/local-llm/claude-code-env.sh
#
# oMLX exposes a native Anthropic Messages API (/v1/messages) — no proxy
# needed for the Anthropic wire format. Claude Code sends requests to this
# endpoint and oMLX translates them to mlx-vlm internally.

export ANTHROPIC_BASE_URL="http://localhost:8080"
export ANTHROPIC_API_KEY="omlx-local"

# Tell Claude Code which model to use (maps to gemma-4-26b-a4b-it-4bit via
# model_alias in ~/.omlx/model_settings.json)
export CLAUDE_CODE_MODEL="gemma-4-26b-a4b-it-4bit"

echo "Claude Code → local oMLX stack (http://localhost:8080)"
CLAUDE_ENV
  _ok "Claude Code env file written to $CLAUDE_ENV_FILE"
  echo "     To activate: source $CLAUDE_ENV_FILE"
  echo "     To make permanent: echo 'source $CLAUDE_ENV_FILE' >> ~/.zshrc"
fi

OMLX_SETTINGS="$HOME/.omlx/settings.json"
PREFIX_CACHE_DIR="$LLM_BASE/prefix-cache"
mkdir -p "$PREFIX_CACHE_DIR"
if [[ -f "$OMLX_SETTINGS" ]]; then
  _info "Applying performance settings to ~/.omlx/settings.json..."
  python3 -c "
import json
path = '$OMLX_SETTINGS'
s = json.load(open(path))

# Performance: disable HF cache scanning (prevents LRU fragmentation)
s.setdefault('huggingface', {})['hf_cache_enabled'] = False

# Performance: 34 GB ceiling — model 15.3 GB + 6 GB hot KV + ~9.3 GB dynamic KV headroom
# soft_threshold=0.90 → eviction/throttle at 30.6 GB (default 0.85 = 28.9 GB was too low;
# concurrent sessions peaked around 30.3 GB, causing adaptive_prefill_throttle below this threshold)
# prefill_min_chunk_tokens=512 → 16× fewer memory checks per large prefill (default 32)
s.setdefault('memory', {})['memory_guard_tier'] = 'custom'
s['memory']['memory_guard_custom_ceiling_gb'] = 34.0
s['memory']['soft_threshold'] = 0.90
s['memory']['hard_threshold'] = 0.95
s['memory']['prefill_safe_zone_ratio'] = 0.90
s['memory']['prefill_min_chunk_tokens'] = 512

# Performance: hot KV cache — 6 GB (reduced from 8 GB to free 2 GB headroom for concurrent large-context sessions); hot_cache_only=false
# allows KV spill to SSD on a shared system (don't lock all available RAM)
s.setdefault('cache', {})['hot_cache_max_size'] = '6GB'
s['cache']['hot_cache_only'] = False
s['cache']['initial_cache_blocks'] = 1024

# Performance: prefix cache on SSD — oMLX reuses computed prefill states for repeated prompts
# Opencode sends the same ~3K system prompt every turn; this skips that prefill entirely
s['cache']['ssd_cache_dir'] = '$PREFIX_CACHE_DIR'
s['cache']['ssd_cache_max_size'] = '20GB'

# Scheduler: chunked_prefill interleaves long prefill with decode for responsiveness
s.setdefault('scheduler', {})['chunked_prefill'] = True
s['scheduler'].setdefault('max_concurrent_requests', 2)

open(path,'w').write(json.dumps(s, indent=2))
print('ok')
" && _ok "settings.json: ceiling=34GB soft=0.90 hot_cache=6GB(SSD-spill-ok) chunk_tokens=512 prefix_cache=$PREFIX_CACHE_DIR hf_cache=false"
fi

# Set model_alias in omlx model_settings.json so gemma4 is addressable by name
OMLX_MODEL_SETTINGS="$HOME/.omlx/model_settings.json"
if [[ -f "$OMLX_MODEL_SETTINGS" ]]; then
  HAS_ALIAS=$(python3 -c "
import json
ms = json.load(open('$OMLX_MODEL_SETTINGS'))
g = ms.get('models',{}).get('gemma-4-26b-a4b-it-4bit',{})
print('yes' if g.get('model_alias') else 'no')
" 2>/dev/null)
  if [[ "$HAS_ALIAS" == "yes" ]]; then
    _ok "model_alias already set in model_settings.json"
  else
    _info "Setting model_alias in ~/.omlx/model_settings.json..."
    python3 -c "
import json
path = '$OMLX_MODEL_SETTINGS'
ms = json.load(open(path))
ms.setdefault('models',{}).setdefault('gemma-4-26b-a4b-it-4bit',{})['model_alias'] = 'gemma4'
open(path,'w').write(json.dumps(ms, indent=2))
print('done')
" && _ok "model_alias=gemma4 set (Claude Code can use model: gemma4)"
  fi
fi

# ── Summary ───────────────────────────────────────────────────────────────────
echo ""
echo "╔══════════════════════════════════════════════╗"
echo "║  Setup complete                             ║"
echo "╠══════════════════════════════════════════════╣"
echo "║  Start stack:   ./llm-start.sh              ║"
echo "║  CLI ask:       ./llm-ask.sh \"your prompt\" ║"
echo "║  Web UI:        http://localhost:8080/admin/chat ║"
echo "║  Opencode:      opencode (auto-configured)  ║"
echo "║  Claude Code:   source ~/.config/local-llm/claude-code-env.sh ║"
echo "║  Stop stack:    ./llm-stop.sh               ║"
echo "╚══════════════════════════════════════════════╝"
echo ""