#!/usr/bin/env bash
# =============================================================================
# Local AI Stack — Ubuntu 24.04
# Services: Ollama · Open WebUI · RAG · MCP · ChromaDB · SearXNG
#           Kiwix · Gitea · InvokeAI · Portainer
#
# Usage:
#   ./laptop_full_setup.sh           — new install or update
#   ./laptop_full_setup.sh --force   — overwrite existing config files too
#
# ZIM downloads:  run ./kiwix_download.sh separately (large files)
# GPU:            Ollama auto-detects VRAM — works with any NVIDIA GPU
# =============================================================================
set -euo pipefail

# ── colours ───────────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
info()    { echo -e "${CYAN}[..]${NC}  $*"; }
ok()      { echo -e "${GREEN}[OK]${NC}  $*"; }
warn()    { echo -e "${YELLOW}[!!]${NC}  $*"; }
die()     { echo -e "${RED}[XX]${NC}  $*" >&2; exit 1; }
section() { echo -e "\n${BOLD}━━━  $*  ━━━${NC}"; }

# ── args ──────────────────────────────────────────────────────────────────────
FORCE=false
for arg in "$@"; do [[ "$arg" == "--force" ]] && FORCE=true; done

# ── config ────────────────────────────────────────────────────────────────────
BASE="$HOME/docker/ai-stack"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOCAL_IP=$(ip route get 1.1.1.1 2>/dev/null | grep -oP 'src \K\S+' \
           || hostname -I | awk '{print $1}')
[[ -z "$LOCAL_IP" ]] && read -rp "Enter your LAN IP: " LOCAL_IP

# Models pulled on first install (adjust to your VRAM)
EMBED_MODEL="nomic-embed-text"   # always needed for RAG
CHAT_MODEL="qwen2.5:14b"
CODE_MODEL="qwen2.5-coder:7b"
FAST_MODEL="qwen2.5:7b"

# ── new vs update ─────────────────────────────────────────────────────────────
IS_UPDATE=false
[[ -f "$BASE/docker-compose.yml" ]] && IS_UPDATE=true

section "Local AI Stack — $($IS_UPDATE && echo UPDATE || echo NEW INSTALL)"
info "Base   : $BASE"
info "Host   : $LOCAL_IP"
$IS_UPDATE && warn "Existing install found. Config files kept unless --force is passed."

# ── helper: write only if missing (or --force) ────────────────────────────────
# Usage:  write_if_new /path/to/file << 'EOF' ... EOF
write_if_new() {
    local dest="$1"
    local content; content=$(cat)
    if [[ ! -f "$dest" ]] || $FORCE; then
        printf '%s\n' "$content" > "$dest"
        ok "Wrote $(basename "$dest")"
    else
        info "Kept   $(basename "$dest")  (--force to overwrite)"
    fi
}

# =============================================================================
section "Prerequisites"
# =============================================================================
if ! $IS_UPDATE; then

    # Docker
    if ! command -v docker &>/dev/null; then
        info "Installing Docker..."
        curl -fsSL https://get.docker.com | sh
        sudo usermod -aG docker "$USER"
        ok "Docker installed — log out/in or run: newgrp docker"
    else
        ok "Docker: $(docker --version)"
    fi

    # NVIDIA Container Toolkit
    if command -v nvidia-smi &>/dev/null && \
       ! dpkg -l 2>/dev/null | grep -q nvidia-container-toolkit; then
        info "Installing NVIDIA Container Toolkit..."
        distribution=$(. /etc/os-release && echo "$ID$VERSION_ID")
        curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
          | sudo gpg --dearmor \
            -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
        curl -fsSL \
          "https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list" \
          | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
          | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
        sudo apt-get update -qq
        sudo apt-get install -y nvidia-container-toolkit
        sudo nvidia-ctk runtime configure --runtime=docker
        sudo systemctl restart docker
        ok "NVIDIA Container Toolkit installed"
    elif command -v nvidia-smi &>/dev/null; then
        ok "NVIDIA Container Toolkit already present"
    else
        warn "No NVIDIA GPU detected — CPU-only (Ollama will use CPU)"
    fi

    # ripgrep (used by MCP search_code tool)
    if ! command -v rg &>/dev/null; then
        sudo apt-get install -y ripgrep
        ok "ripgrep installed"
    fi

else
    ok "Skipping prereq checks (update mode)"
fi

# =============================================================================
section "Directories"
# =============================================================================
for d in papers repos workspace index searxng invokeai-data invokeai-outputs \
          kiwix gitea portainer-data logs; do
    mkdir -p "$BASE/$d"
done
ok "Directories ready under $BASE"

# =============================================================================
section "Server Files"
# =============================================================================

# Copy Python servers from repo (always update — they are version-controlled)
cp "$SCRIPT_DIR/server.py"     "$BASE/server.py"     && ok "server.py"
cp "$SCRIPT_DIR/mcp_server.py" "$BASE/mcp_server.py" && ok "mcp_server.py"

# RAG dependencies
cat > "$BASE/requirements.txt" << 'REQ'
fastapi
uvicorn[standard]
httpx
pydantic
chromadb
pypdf
watchdog
python-multipart
REQ
ok "requirements.txt"

# MCP server dependencies
cat > "$BASE/mcp_requirements.txt" << 'REQ'
mcp[cli]
fastapi
uvicorn[standard]
httpx
REQ
ok "mcp_requirements.txt"

# =============================================================================
section "SearXNG Config"
# =============================================================================
write_if_new "$BASE/searxng/settings.yml" << SEARXNG
use_default_settings: true
general:
  instance_name: "Local Search"
server:
  secret_key: "$(openssl rand -hex 32)"
  limiter: false
search:
  safe_search: 0
  default_lang: "en"
  formats: [html, json]
SEARXNG

# =============================================================================
section ".env File"
# =============================================================================
# Never overwrite — this is where users store their tokens
if [[ ! -f "$BASE/.env" ]]; then
    cat > "$BASE/.env" << ENV
# Local AI Stack — API Tokens
# Edit this file, then restart: bash $BASE/start.sh

# Gitea — generate at http://$LOCAL_IP:3001/user/settings/applications
GITEA_TOKEN=your-gitea-token-here

# GitHub — optional, for GitHub API access via MCP
GITHUB_TOKEN=your-github-token-here
ENV
    ok "Created .env — add your tokens before using MCP Gitea/GitHub tools"
else
    info "Kept   .env (never overwritten)"
fi

# =============================================================================
section "Docker Compose"
# =============================================================================
# Always written — it's the stack definition and safe to update
cat > "$BASE/docker-compose.yml" << COMPOSE
# Local AI Stack — generated $(date '+%Y-%m-%d')
# Edit \$BASE/.env to add API tokens.
# GPU: OLLAMA_NUM_GPU=999 means "use all available VRAM" — auto-adapts to any GPU.

services:

  # ── Ollama — LLM inference ──────────────────────────────────────────────────
  ollama:
    image: ollama/ollama:latest
    container_name: ollama
    restart: unless-stopped
    ports:
      - "0.0.0.0:11434:11434"
    volumes:
      - ollama-models:/root/.ollama
    environment:
      - OLLAMA_NUM_GPU=999       # use all available VRAM (auto-detects GPU size)
      - OLLAMA_NUM_CTX=8192      # lower to 4096 if you hit OOM
      - OLLAMA_KEEP_ALIVE=24h
      - OLLAMA_MAX_LOADED_MODELS=1
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    healthcheck:
      test: ["CMD", "ollama", "list"]
      interval: 30s
      timeout: 10s
      retries: 5

  # ── Open WebUI — Chat interface ─────────────────────────────────────────────
  open-webui:
    image: ghcr.io/open-webui/open-webui:main
    container_name: open-webui
    restart: unless-stopped
    ports:
      - "0.0.0.0:3000:8080"
    volumes:
      - open-webui-data:/app/backend/data
    environment:
      - OLLAMA_BASE_URL=http://ollama:11434
      - OPENAI_API_BASE_URL=http://rag-server:8001/v1
      - OPENAI_API_KEY=local-rag-key
      - ENABLE_OPENAI_API=true
      - ENABLE_RAG_WEB_SEARCH=true
      - RAG_WEB_SEARCH_ENGINE=searxng
      - SEARXNG_QUERY_URL=http://searxng:8080/search?q=<query>&format=json
      - WEBUI_AUTH=false
    depends_on:
      ollama:
        condition: service_healthy

  # ── ChromaDB — Vector store ─────────────────────────────────────────────────
  chromadb:
    image: chromadb/chroma:latest
    container_name: chromadb
    restart: unless-stopped
    ports:
      - "0.0.0.0:8000:8000"
    volumes:
      - $BASE/index:/chroma/chroma
    environment:
      - IS_PERSISTENT=TRUE
      - ANONYMIZED_TELEMETRY=FALSE
    healthcheck:
      test: ["CMD-SHELL", "wget -qO- http://localhost:8000/api/v2/heartbeat || exit 1"]
      interval: 15s
      timeout: 5s
      retries: 5

  # ── RAG Server — code-aware retrieval ──────────────────────────────────────
  rag-server:
    image: python:3.11-slim
    container_name: rag-server
    restart: unless-stopped
    ports:
      - "0.0.0.0:8001:8001"
    volumes:
      - $BASE/papers:/papers
      - $BASE/repos:/repos
      - $BASE/index:/index
      - $BASE/server.py:/app/server.py
      - $BASE/requirements.txt:/app/requirements.txt
    working_dir: /app
    environment:
      - OLLAMA_URL=http://ollama:11434
      - CHROMA_URL=http://chromadb:8000
      - EMBED_MODEL=nomic-embed-text
      - CHAT_MODEL=qwen2.5:14b
      - PAPERS_DIR=/papers
      - REPOS_DIR=/repos
    command: >
      bash -c "apt-get update -qq &&
               apt-get install -y --no-install-recommends git &&
               pip install --no-cache-dir -r requirements.txt &&
               uvicorn server:app --host 0.0.0.0 --port 8001"
    depends_on:
      chromadb:
        condition: service_healthy
      ollama:
        condition: service_healthy

  # ── MCP Server — Claude Code-equivalent tools ───────────────────────────────
  # Connect via: http://$LOCAL_IP:8002/sse
  # Add to Claude Code:  claude mcp add local http://$LOCAL_IP:8002/sse
  mcp-server:
    image: python:3.11-slim
    container_name: mcp-server
    restart: unless-stopped
    ports:
      - "0.0.0.0:8002:8002"
    volumes:
      - $BASE/workspace:/workspace
      - $BASE/repos:/repos
      - $BASE/mcp_server.py:/app/mcp_server.py
      - $BASE/mcp_requirements.txt:/app/mcp_requirements.txt
    working_dir: /app
    env_file: $BASE/.env
    environment:
      - WORKSPACE_DIR=/workspace
      - REPOS_DIR=/repos
      - GITEA_URL=http://gitea:3000
      - RAG_URL=http://rag-server:8001
    command: >
      bash -c "apt-get update -qq &&
               apt-get install -y --no-install-recommends git ripgrep &&
               pip install --no-cache-dir -r mcp_requirements.txt &&
               python mcp_server.py"
    depends_on:
      - rag-server

  # ── SearXNG — Private web search ───────────────────────────────────────────
  searxng:
    image: searxng/searxng:latest
    container_name: searxng
    restart: unless-stopped
    ports:
      - "0.0.0.0:8888:8080"
    volumes:
      - $BASE/searxng:/etc/searxng
    cap_drop: [ALL]
    cap_add: [CHOWN, SETGID, SETUID]

  # ── Kiwix — Offline Wikipedia/docs ─────────────────────────────────────────
  # Download ZIMs first: ./kiwix_download.sh
  kiwix:
    image: ghcr.io/kiwix/kiwix-serve:latest
    container_name: kiwix
    restart: unless-stopped
    ports:
      - "0.0.0.0:8181:8080"
    volumes:
      - $BASE/kiwix:/data
    command: "*.zim"

  # ── Gitea — Self-hosted Git ─────────────────────────────────────────────────
  gitea:
    image: gitea/gitea:latest
    container_name: gitea
    restart: unless-stopped
    ports:
      - "0.0.0.0:3001:3000"
      - "0.0.0.0:2222:22"
    volumes:
      - $BASE/gitea:/data
      - /etc/timezone:/etc/timezone:ro
      - /etc/localtime:/etc/localtime:ro
    environment:
      - USER_UID=1000
      - USER_GID=1000
      - GITEA__database__DB_TYPE=sqlite3
      - GITEA__database__PATH=/data/gitea/gitea.db
      - GITEA__webhook__ALLOWED_HOST_LIST=rag-server,mcp-server

  # ── InvokeAI — Image generation ────────────────────────────────────────────
  invokeai:
    image: ghcr.io/invoke-ai/invokeai:latest
    container_name: invokeai
    restart: unless-stopped
    ports:
      - "0.0.0.0:9090:9090"
    volumes:
      - invokeai-models:/invokeai/models
      - $BASE/invokeai-outputs:/invokeai/outputs
      - $BASE/invokeai-data:/invokeai/databases
    environment:
      - INVOKEAI_HOST=0.0.0.0
      - INVOKEAI_PORT=9090
      - INVOKEAI_PRECISION=float16
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

  # ── Portainer — Docker management UI ───────────────────────────────────────
  portainer:
    image: portainer/portainer-ce:latest
    container_name: portainer
    restart: unless-stopped
    ports:
      - "0.0.0.0:9000:9000"
      - "0.0.0.0:9443:9443"
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock
      - $BASE/portainer-data:/data

volumes:
  ollama-models:
  open-webui-data:
  invokeai-models:
COMPOSE
ok "docker-compose.yml written"

# =============================================================================
section "Firewall (UFW)"
# =============================================================================
if command -v ufw &>/dev/null; then
    if [[ ! -f "$BASE/.ufw-done" ]] || $FORCE; then
        read -rp "  LAN subnet for firewall [192.168.1.0/24]: " LAN_SUBNET
        LAN_SUBNET="${LAN_SUBNET:-192.168.1.0/24}"
        [[ "$LAN_SUBNET" =~ /[0-9]+$ ]] || LAN_SUBNET="${LAN_SUBNET}/24"

        for port_comment in \
            "3000:Open WebUI" "11434:Ollama" "8001:RAG Server" \
            "8002:MCP Server" "8000:ChromaDB" "8888:SearXNG" \
            "8181:Kiwix" "3001:Gitea" "2222:Gitea SSH" \
            "9090:InvokeAI" "9000:Portainer" "9443:Portainer S"; do
            port="${port_comment%%:*}"
            comment="${port_comment##*:}"
            sudo ufw allow from "$LAN_SUBNET" to any port "$port" proto tcp \
                comment "$comment" > /dev/null
        done
        sudo ufw reload > /dev/null
        ok "UFW rules set for $LAN_SUBNET"
        touch "$BASE/.ufw-done"
    else
        info "UFW rules already set (--force to redo)"
    fi
else
    warn "ufw not found — skipping firewall config"
fi

# =============================================================================
section "Helper Scripts"
# =============================================================================

cat > "$BASE/start.sh" << STARTSH
#!/bin/bash
cd "$BASE"
echo "Pulling latest images..."
docker compose pull --quiet
docker compose up -d
echo ""
echo "  Open WebUI  →  http://$LOCAL_IP:3000"
echo "  InvokeAI    →  http://$LOCAL_IP:9090"
echo "  SearXNG     →  http://$LOCAL_IP:8888"
echo "  Kiwix       →  http://$LOCAL_IP:8181"
echo "  Gitea       →  http://$LOCAL_IP:3001"
echo "  RAG Health  →  http://$LOCAL_IP:8001/health"
echo "  MCP SSE     →  http://$LOCAL_IP:8002/sse"
echo "  Portainer   →  https://$LOCAL_IP:9443"
echo ""
echo "  Workspace   →  $BASE/workspace/"
echo "  Drop PDFs   →  $BASE/papers/"
echo "  Images out  →  $BASE/invokeai-outputs/"
echo ""
echo "  Claude Code MCP:"
echo "    claude mcp add local http://$LOCAL_IP:8002/sse"
STARTSH
chmod +x "$BASE/start.sh"
ok "start.sh"

cat > "$BASE/stop.sh" << STOPSH
#!/bin/bash
cd "$BASE"
docker compose down
echo "All services stopped."
STOPSH
chmod +x "$BASE/stop.sh"
ok "stop.sh"

cat > "$BASE/status.sh" << 'STATUSSH'
#!/bin/bash
echo "=== GPU ==="
nvidia-smi --query-gpu=name,temperature.gpu,utilization.gpu,memory.used,memory.total \
    --format=csv,noheader 2>/dev/null || echo "(no GPU)"

echo ""
echo "=== Containers ==="
docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}"

echo ""
echo "=== Ollama models ==="
docker exec ollama ollama ps 2>/dev/null || echo "(not running)"

echo ""
echo "=== RAG ==="
curl -s http://localhost:8001/health | python3 -m json.tool 2>/dev/null \
    || echo "(not running)"

echo ""
echo "=== MCP ==="
curl -s http://localhost:8002/health 2>/dev/null || echo "(not running)"

echo ""
echo "=== Disk ==="
du -sh ~/docker/ai-stack/*/  2>/dev/null
STATUSSH
chmod +x "$BASE/status.sh"
ok "status.sh"

cat > "$BASE/pull-models.sh" << PULLSH
#!/bin/bash
echo "Waiting for Ollama..."
until docker exec ollama ollama list &>/dev/null; do sleep 3; done
echo "Ollama ready."

echo ""
echo "Pulling embed model (needed for RAG)..."
docker exec ollama ollama pull $EMBED_MODEL

echo "Pulling fast chat model (~5GB)..."
docker exec ollama ollama pull $FAST_MODEL

echo "Pulling smart chat model (~9GB)..."
docker exec ollama ollama pull $CHAT_MODEL

echo "Pulling code model (~5GB)..."
docker exec ollama ollama pull $CODE_MODEL

echo ""
echo "Done. Models:"
docker exec ollama ollama list
PULLSH
chmod +x "$BASE/pull-models.sh"
ok "pull-models.sh"

write_if_new "$BASE/Caddyfile.example" << CADDY
# Caddy2 reverse proxy — copy to your proxy machine
# Replace yourdomain.com with your actual domain

webui.yourdomain.com    { reverse_proxy $(hostname).local:3000 }
invokeai.yourdomain.com { reverse_proxy $(hostname).local:9090 }
search.yourdomain.com   { reverse_proxy $(hostname).local:8888 }
git.yourdomain.com      { reverse_proxy $(hostname).local:3001 }
kiwix.yourdomain.com    { reverse_proxy $(hostname).local:8181 }
rag.yourdomain.com      { reverse_proxy $(hostname).local:8001 }
mcp.yourdomain.com      { reverse_proxy $(hostname).local:8002 }
portainer.yourdomain.com {
    reverse_proxy https://$(hostname).local:9443 {
        transport http { tls_insecure_skip_verify }
    }
}
CADDY

# =============================================================================
section "Systemd Auto-Start"
# =============================================================================
sudo tee /etc/systemd/system/local-ai.service > /dev/null << SYSD
[Unit]
Description=Local AI Stack
After=docker.service network-online.target
Requires=docker.service

[Service]
Type=oneshot
RemainAfterExit=yes
User=$USER
WorkingDirectory=$BASE
ExecStart=/bin/bash $BASE/start.sh
ExecStop=/bin/bash $BASE/stop.sh
TimeoutStartSec=300

[Install]
WantedBy=multi-user.target
SYSD
sudo systemctl daemon-reload
sudo systemctl enable local-ai.service
ok "Systemd service enabled (local-ai.service)"

# =============================================================================
section "Starting Stack"
# =============================================================================
cd "$BASE"
info "Pulling images (this takes a few minutes on first run)..."
docker compose pull --quiet
docker compose up -d
ok "Stack started"

if ! $IS_UPDATE; then
    echo ""
    read -rp "Pull Ollama models now? (~20GB total, takes 10-40 min) [Y/n]: " DO_PULL
    if [[ "${DO_PULL,,}" != "n" ]]; then
        bash "$BASE/pull-models.sh"
    else
        info "Run later: bash $BASE/pull-models.sh"
    fi
fi

# =============================================================================
echo ""
echo -e "${GREEN}${BOLD}━━━  Done!  ━━━${NC}"
echo ""
echo -e "  ${CYAN}Open WebUI${NC}   →  http://$LOCAL_IP:3000"
echo -e "  ${CYAN}InvokeAI${NC}     →  http://$LOCAL_IP:9090"
echo -e "  ${CYAN}SearXNG${NC}      →  http://$LOCAL_IP:8888"
echo -e "  ${CYAN}Kiwix${NC}        →  http://$LOCAL_IP:8181  (needs ZIMs — run kiwix_download.sh)"
echo -e "  ${CYAN}Gitea${NC}        →  http://$LOCAL_IP:3001"
echo -e "  ${CYAN}RAG Health${NC}   →  http://$LOCAL_IP:8001/health"
echo -e "  ${CYAN}MCP SSE${NC}      →  http://$LOCAL_IP:8002/sse"
echo -e "  ${CYAN}Portainer${NC}    →  https://$LOCAL_IP:9443"
echo ""
echo -e "  ${YELLOW}Add MCP to Claude Code:${NC}"
echo    "    claude mcp add local http://$LOCAL_IP:8002/sse"
echo ""
echo -e "  ${YELLOW}Add API tokens to:${NC}  $BASE/.env"
echo -e "  ${YELLOW}Drop PDFs into:${NC}     $BASE/papers/"
echo -e "  ${YELLOW}Your workspace:${NC}     $BASE/workspace/"
echo -e "  ${YELLOW}ZIM downloads:${NC}      ./kiwix_download.sh"
echo ""
$IS_UPDATE && echo -e "  ${GREEN}Update complete.${NC}" \
           || echo -e "  ${GREEN}Fresh install complete.${NC}"
echo ""