#!/bin/bash
# =============================================================================
# Laptop AI Stack — Ubuntu 24.04 + RTX 3000 Mobile (6GB VRAM / 32GB RAM)
# All files live under ~/docker/ai-stack/
# Usage: bash laptop-ai-setup.sh
# =============================================================================
set -euo pipefail

G='\033[0;32m'; Y='\033[1;33m'; C='\033[0;36m'; N='\033[0m'
ok()  { echo -e "${G}[OK]${N}  $1"; }
inf() { echo -e "${C}[..]${N}  $1"; }
wrn() { echo -e "${Y}[!!]${N}  $1"; }
sec() { echo -e "\n${C}══════════════════════════════════════════════════${N}";
        echo -e "${C}  $1${N}";
        echo -e "${C}══════════════════════════════════════════════════${N}"; }

# All paths relative to this base — no exceptions
BASE="$HOME/docker/ai-stack"
# Get the primary non-Docker, non-loopback IP
LOCAL_IP=$(ip route get 1.1.1.1 2>/dev/null | grep -oP 'src \K\S+' || hostname -I | awk '{print $1}')
if [[ -z "$LOCAL_IP" ]]; then
    read -rp "  Could not detect IP. Enter your LAN IP manually: " LOCAL_IP
fi

# =============================================================================
sec "STEP 1 — Configuration"
# =============================================================================
echo ""
echo "  Your hardware: RTX 3000 Mobile (6GB VRAM) + 32GB RAM"
echo "  All files will be stored under: $BASE"
echo ""
echo "  Model options:"
echo "    [1] Smart  — 7B + 14B only  (fast, recommended)"
echo "        7B:  30-50 tok/s — instant responses"
echo "        14B: 10-20 tok/s — ~20 sec typical response"
echo ""
echo "    [2] Full   — 7B + 14B + 32B  (adds slow but powerful model)"
echo "        32B:  3-5 tok/s  — 1-4 min typical response"
echo "        Good for: deep research — ask and walk away"
echo ""
read -rp "  Model choice [1]: " MODEL_CHOICE; MODEL_CHOICE="${MODEL_CHOICE:-1}"
USE_32B=false
[[ "$MODEL_CHOICE" == "2" ]] && USE_32B=true

echo ""
echo ""
echo "  LAN subnet for firewall access — common options:"
echo "    192.168.1.0/24  — only your 192.168.1.x devices (most secure)"
echo "    192.168.0.0/16  — any 192.168.x.x device (covers subnet changes)"
echo "    10.0.0.0/8      — if your network uses 10.x.x.x addresses"
read -rp "  Your LAN subnet [192.168.1.0/24]: " LAN_SUBNET
LAN_SUBNET="${LAN_SUBNET:-192.168.1.0/24}"
# Validate format — must end in /number
if [[ ! "$LAN_SUBNET" =~ /[0-9]+$ ]]; then
    wrn "Subnet missing CIDR mask — appending /24"
    LAN_SUBNET="${LAN_SUBNET}/24"
fi

echo ""
echo -e "  ${C}─── Directory Layout ─────────────────────────────${N}"
echo "  $BASE/"
echo "  ├── docker-compose.yml"
echo "  ├── server.py"
echo "  ├── requirements.txt"
echo "  ├── papers/          ← drop PDFs here to auto-index"
echo "  ├── repos/           ← git repos for code RAG"
echo "  ├── index/           ← ChromaDB index files"
echo "  ├── searxng/         ← SearXNG config"
echo "  ├── invokeai-outputs/ ← generated images"
echo "  └── logs/            ← service logs"
echo ""
read -rp "  Proceed? (Y/n): " CONFIRM
[[ "${CONFIRM,,}" == "n" ]] && exit 0

# =============================================================================
sec "STEP 2 — Prerequisites"
# =============================================================================
inf "Checking Docker..."
if ! command -v docker &>/dev/null; then
    inf "Installing Docker..."
    sudo apt-get update -qq
    sudo apt-get install -y curl
    curl -fsSL https://get.docker.com | sudo bash
    sudo usermod -aG docker "$USER"
    ok "Docker installed — you may need to log out/in for group to take effect"
    wrn "If docker commands fail, run: newgrp docker"
else
    ok "Docker present: $(docker --version)"
fi

inf "Checking NVIDIA Container Toolkit..."
if ! dpkg -l | grep -q nvidia-container-toolkit 2>/dev/null; then
    inf "Installing NVIDIA Container Toolkit..."
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
        | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
    curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.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"
else
    ok "NVIDIA Container Toolkit present"
fi

# =============================================================================
sec "STEP 3 — Directory Structure"
# =============================================================================
mkdir -p "$BASE"/{papers,repos,index,searxng,invokeai-outputs,logs,kiwix,gitea}
ok "Directories created:"
ok "  $BASE/papers"
ok "  $BASE/repos"
ok "  $BASE/index"
ok "  $BASE/searxng"
ok "  $BASE/invokeai-outputs"
ok "  $BASE/logs"

# =============================================================================
sec "STEP 4 — RAG Server Files"
# =============================================================================
inf "Writing $BASE/server.py..."
cat > "$BASE/server.py" << 'PYTHON_SERVER_EOF'
#!/usr/bin/env python3
"""
Local RAG Server — Complete Standalone Script
Drop PDFs into the papers folder and they auto-index.
OpenAI-compatible /v1/chat/completions API.
"""
import os, sys, json, time, hashlib, logging, asyncio, threading
from pathlib import Path
from typing import Optional

try:
    import chromadb
    import httpx
    from fastapi import FastAPI, BackgroundTasks, HTTPException
    from fastapi.responses import StreamingResponse
    from fastapi.middleware.cors import CORSMiddleware
    from pydantic import BaseModel
    from watchdog.observers import Observer
    from watchdog.events import FileSystemEventHandler
    from langchain_community.document_loaders import PyPDFLoader, TextLoader
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    from langchain_ollama import OllamaEmbeddings
    from langchain_chroma import Chroma
except ImportError as e:
    print(f"[ERROR] Missing dependency: {e}")
    sys.exit(1)

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger("rag")

OLLAMA_URL  = os.getenv("OLLAMA_URL",  "http://ollama:11434")
CHROMA_URL  = os.getenv("CHROMA_URL",  "http://chromadb:8000")
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
CHAT_MODEL  = os.getenv("CHAT_MODEL",  "qwen2.5:14b")
PAPERS_DIR  = Path(os.getenv("PAPERS_DIR", "/papers"))
REPOS_DIR   = Path(os.getenv("REPOS_DIR",  "/repos"))
INDEX_DIR   = Path(os.getenv("INDEX_DIR",  "/index"))
INDEX_FILE  = INDEX_DIR / "indexed.json"
PORT        = int(os.getenv("PORT", "8001"))

PAPER_EXTENSIONS = {".pdf",".txt",".md",".docx",".rst"}
CODE_EXTENSIONS  = {".py",".js",".ts",".html",".css",".json",".sh",".yaml",".sql"}
SKIP_DIRS        = {"node_modules",".git","__pycache__",".venv","dist","build"}

for d in [PAPERS_DIR, REPOS_DIR, INDEX_DIR]:
    d.mkdir(parents=True, exist_ok=True)

chroma_client = papers_db = code_db = embeddings = None
_index_lock = threading.Lock()

def init_chroma():
    host = CHROMA_URL.replace("http://","").replace("https://","").split(":")[0]
    port = int(CHROMA_URL.split(":")[-1])
    for attempt in range(15):
        try:
            c = chromadb.HttpClient(host=host, port=port)
            c.heartbeat()
            log.info("Connected to ChromaDB")
            return c
        except Exception as e:
            log.warning(f"ChromaDB not ready ({attempt+1}/15): {e}")
            time.sleep(3)
    sys.exit(1)

def setup_databases():
    global chroma_client, papers_db, code_db, embeddings
    chroma_client = init_chroma()
    embeddings = OllamaEmbeddings(model=EMBED_MODEL, base_url=OLLAMA_URL)
    papers_db = Chroma(client=chroma_client, collection_name="papers", embedding_function=embeddings)
    code_db   = Chroma(client=chroma_client, collection_name="code",   embedding_function=embeddings)
    log.info("Vector stores ready")

def load_index():
    with _index_lock:
        if INDEX_FILE.exists():
            try: return json.loads(INDEX_FILE.read_text())
            except: return {}
        return {}

def save_index(idx):
    with _index_lock:
        INDEX_FILE.write_text(json.dumps(idx, indent=2))

def file_hash(path):
    try: return hashlib.md5(path.read_bytes()).hexdigest()
    except: return ""

def ingest_paper(path):
    idx = load_index(); h = file_hash(path); key = str(path)
    if idx.get(key) == h: return
    log.info(f"[papers] Ingesting: {path.name}")
    try:
        loader = PyPDFLoader(str(path)) if path.suffix == ".pdf" \
            else TextLoader(str(path), encoding="utf-8")
        pages = loader.load()
    except Exception as e:
        log.warning(f"Could not load {path.name}: {e}"); return
    s_spl = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
    d_spl = RecursiveCharacterTextSplitter(chunk_size=400,  chunk_overlap=80)
    s_docs = s_spl.split_documents(pages)
    d_docs = d_spl.split_documents(pages)
    for i,doc in enumerate(s_docs): doc.metadata.update({"layer":"summary","source_file":path.name})
    for i,doc in enumerate(d_docs): doc.metadata.update({"layer":"detail", "source_file":path.name})
    papers_db.add_documents(s_docs + d_docs)
    log.info(f"[papers] {path.name}: {len(s_docs)} summary + {len(d_docs)} detail chunks")
    idx[key] = h; save_index(idx)

def ingest_papers_dir():
    for ext in PAPER_EXTENSIONS:
        for f in PAPERS_DIR.rglob(f"*{ext}"): ingest_paper(f)

def retrieve(query, k=6):
    results = []
    if papers_db:
        try: results.extend(papers_db.similarity_search(query, k=4))
        except: pass
    if code_db:
        try: results.extend(code_db.similarity_search(query, k=3))
        except: pass
    seen = set(); unique = []
    for doc in results:
        h = hashlib.md5(doc.page_content.encode()).hexdigest()
        if h not in seen: seen.add(h); unique.append(doc)
    return unique[:k*2]

app = FastAPI(title="Local RAG Server")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

class Message(BaseModel):
    role: str; content: str

class ChatRequest(BaseModel):
    model: Optional[str] = None
    messages: list[Message]
    stream: Optional[bool] = False
    top_k: Optional[int] = 6

@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    user_msg = next((m.content for m in reversed(req.messages) if m.role == "user"), "")
    if not user_msg.strip(): raise HTTPException(400, "No user message")
    chunks = retrieve(user_msg, req.top_k or 6)
    parts = []; sources = []
    for chunk in chunks:
        fname = chunk.metadata.get("source_file","unknown")
        parts.append(f"[{fname}]\n{chunk.page_content.strip()}")
        if fname not in sources: sources.append(fname)
    context = "\n\n---\n\n".join(parts)
    source_list = "\n".join(f"  • {s}" for s in sources)
    sys_content = f"You are a research assistant.\n\n=== CONTEXT ===\n{context}\n\n=== SOURCES ===\n{source_list}" \
        if context.strip() else "You are a helpful assistant. No relevant documents found."
    ollama_messages = [{"role":"system","content":sys_content}] + \
                      [{"role":m.role,"content":m.content} for m in req.messages]
    model = req.model or CHAT_MODEL
    if req.stream:
        async def stream_gen():
            async with httpx.AsyncClient(timeout=600) as client:
                async with client.stream("POST", f"{OLLAMA_URL}/api/chat",
                        json={"model":model,"messages":ollama_messages,"stream":True}) as resp:
                    async for line in resp.aiter_lines():
                        if not line.strip(): continue
                        try:
                            token = json.loads(line).get("message",{}).get("content","")
                            if token:
                                yield f"data: {json.dumps({'object':'chat.completion.chunk','choices':[{'index':0,'delta':{'content':token}}]})}\n\n"
                        except: pass
                    yield "data: [DONE]\n\n"
        return StreamingResponse(stream_gen(), media_type="text/event-stream")
    async with httpx.AsyncClient(timeout=600) as client:
        resp = await client.post(f"{OLLAMA_URL}/api/chat",
            json={"model":model,"messages":ollama_messages,"stream":False})
    resp.raise_for_status()
    reply = resp.json().get("message",{}).get("content","")
    return {"object":"chat.completion","model":model,
            "choices":[{"index":0,"message":{"role":"assistant","content":reply}}]}

@app.get("/health")
async def health():
    pc = papers_db._collection.count() if papers_db else 0
    cc = code_db._collection.count()   if code_db   else 0
    return {"status":"ok","papers_chunks":pc,"code_chunks":cc,
            "chat_model":CHAT_MODEL,"embed_model":EMBED_MODEL}

@app.post("/ingest/papers")
async def trigger_ingest(background: BackgroundTasks):
    background.add_task(ingest_papers_dir)
    return {"status":"paper ingestion queued"}

class PaperDropHandler(FileSystemEventHandler):
    def _handle(self, path_str):
        path = Path(path_str)
        if path.suffix.lower() in PAPER_EXTENSIONS and path.is_file():
            time.sleep(0.5); ingest_paper(path)
    def on_created(self, event):
        if not event.is_directory: self._handle(event.src_path)

@app.on_event("startup")
async def on_startup():
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, setup_databases)
    observer = Observer()
    observer.schedule(PaperDropHandler(), str(PAPERS_DIR), recursive=True)
    observer.start()
    def initial_index():
        time.sleep(2); ingest_papers_dir()
        log.info("Initial scan complete")
    threading.Thread(target=initial_index, daemon=True).start()
    log.info(f"RAG server ready on port {PORT}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("server:app", host="0.0.0.0", port=PORT, reload=False)
PYTHON_SERVER_EOF
ok "server.py written"

cat > "$BASE/requirements.txt" << 'REQ'
fastapi
uvicorn[standard]
httpx
pydantic
chromadb
langchain
langchain-community
langchain-chroma
langchain-ollama
langchain-text-splitters
pypdf
unstructured
watchdog
python-multipart
REQ
ok "requirements.txt written"

# =============================================================================
sec "STEP 5 — SearXNG Config"
# =============================================================================
cat > "$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
ok "SearXNG config written"

# =============================================================================
sec "STEP 6 — Docker Compose"
# =============================================================================
# Write .env file so BASE path is available in compose
cat > "$BASE/.env" << ENV
BASE=$BASE
ENV

cat > "$BASE/docker-compose.yml" << COMPOSE
# =============================================================================
# Laptop AI Stack — RTX 3000 Mobile (6GB) + 32GB RAM
# Base directory: $BASE
# Generated: $(date)
#
# Directory layout on host:
#   $BASE/papers/           → RAG documents (drop PDFs here)
#   $BASE/repos/            → Git repos for code RAG
#   $BASE/index/            → ChromaDB index
#   $BASE/searxng/          → SearXNG config
#   $BASE/invokeai-outputs/ → Generated images
#   $BASE/logs/             → Service logs
# =============================================================================

services:

  # ── Ollama — Local 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_CTX=16384
      - OLLAMA_KEEP_ALIVE=24h
      - OLLAMA_MAX_LOADED_MODELS=1
      - OLLAMA_NUM_GPU=999
    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
      - OPENAI_API_KEY=local-rag-key
      - 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

  # ── 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
      - invokeai-db:/invokeai/databases
    environment:
      - INVOKEAI_HOST=0.0.0.0
      - INVOKEAI_PORT=9090
      - INVOKEAI_PRECISION=float16
      - INVOKEAI_VRAM_CACHE_SIZE=3
      - INVOKEAI_RAM_CACHE_SIZE=12
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

  # ── ChromaDB — Vector database ────────────────────────────────────────────
  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 ────────────────────────────────────────────────────────────
  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
      - INDEX_DIR=/index
      - PORT=8001
    command: >
      bash -c "apt-get update -qq &&
               apt-get install -y --no-install-recommends libmagic1 poppler-utils &&
               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

  # ── 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 ─────────────────────────────────────────────
  # Download ZIM file first:
  # cd ~/docker/ai-stack/kiwix
  # wget https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic_2024-10.zim
  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__server__ROOT_URL=https://git.yourdomain.com
      - GITEA__server__SSH_DOMAIN=git.yourdomain.com
      - GITEA__server__DOMAIN=git.yourdomain.com
      - GITEA__webhook__ALLOWED_HOST_LIST=rag-server

  # ── 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
      - portainer-data:/data

# Named volumes for data that doesn't need to be directly browsable on host
# (models, webui state, invokeai db) — everything else maps to $BASE subfolders
volumes:
  ollama-models:
  open-webui-data:
  invokeai-models:
  invokeai-db:
  portainer-data:
COMPOSE
ok "docker-compose.yml written → $BASE/docker-compose.yml"

# =============================================================================
sec "STEP 7 — UFW Rules"
# =============================================================================
sudo ufw allow from "$LAN_SUBNET" to any port 9090 proto tcp comment "InvokeAI"    > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 3000 proto tcp comment "Open WebUI"  > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 8001 proto tcp comment "RAG Server"  > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 8888 proto tcp comment "SearXNG"     > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 11434 proto tcp comment "Ollama"     > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 8000 proto tcp comment "ChromaDB"    > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 9000 proto tcp comment "Portainer"   > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 9443 proto tcp comment "Portainer S" > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 8181 proto tcp comment "Kiwix"       > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 3001 proto tcp comment "Gitea"       > /dev/null
sudo ufw allow from "$LAN_SUBNET" to any port 2222 proto tcp comment "Gitea SSH"   > /dev/null
sudo ufw reload > /dev/null
ok "UFW rules set for $LAN_SUBNET"

# Fix Docker bypassing UFW on Ubuntu 24.04
DAEMON_FILE="/etc/docker/daemon.json"
if ! grep -q "iptables" "$DAEMON_FILE" 2>/dev/null; then
    inf "Applying Docker UFW fix..."
    if [[ -f "$DAEMON_FILE" ]]; then
        sudo python3 -c "
import json
with open('$DAEMON_FILE') as f: d = json.load(f)
d['iptables'] = True
with open('$DAEMON_FILE','w') as f: json.dump(d, f, indent=2)
"
    else
        echo '{"iptables": true}' | sudo tee "$DAEMON_FILE" > /dev/null
    fi
    sudo systemctl restart docker
    ok "Docker UFW fix applied"
fi

# =============================================================================
sec "STEP 8 — Helper Scripts"
# =============================================================================

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

echo "1/4 — Embedding model (required for RAG, ~274MB)..."
docker exec ollama ollama pull nomic-embed-text

echo "2/4 — Qwen 7B: fast everyday model, fits 100% in 6GB GPU (~5GB)..."
docker exec ollama ollama pull qwen2.5:7b

echo "3/4 — Qwen 14B: smarter, partial GPU offload (~9GB)..."
docker exec ollama ollama pull qwen2.5:14b

echo "4/4 — Qwen Coder 7B: code + scripts, fits in GPU (~5GB)..."
docker exec ollama ollama pull qwen2.5-coder:7b

$([ "$USE_32B" = true ] && cat << 'PULL32B'
echo ""
echo "32B model selected — 19GB, slow on 6GB GPU (1-4 min responses)"
read -rp "Pull 32B now? (y/N): " DO32B
if [[ "\${DO32B,,}" == "y" ]]; then
    echo "Pulling Qwen 32B (~19GB — takes 20-40 min)..."
    docker exec ollama ollama pull qwen2.5:32b
    echo "32B pulled."
fi
PULL32B
)

echo ""
echo "Models ready. Speed guide on your RTX 3000 Mobile:"
echo "  nomic-embed-text  — automatic (RAG embeddings only)"
echo "  qwen2.5:7b        — 30-50 tok/s  — chat, image prompts, quick Q&A"
echo "  qwen2.5:14b       — 10-20 tok/s  — RAG, analysis, writing"
echo "  qwen2.5-coder:7b  — 30-50 tok/s  — code, scripts"
$([ "$USE_32B" = true ] && echo "echo \"  qwen2.5:32b       — 3-5 tok/s    — deep research (slow)\"")
PULL
chmod +x "$BASE/pull-models.sh"

# ── start.sh ──────────────────────────────────────────────────────────────────
cat > "$BASE/start.sh" << STARTSH
#!/bin/bash
cd "$BASE"
docker compose up -d
echo ""
echo "  InvokeAI    →  http://$LOCAL_IP:9090"
echo "  Open WebUI  →  http://$LOCAL_IP:3000"
echo "  SearXNG     →  http://$LOCAL_IP:8080"
echo "  RAG Health  →  http://$LOCAL_IP:8001/health"
echo "  Portainer   →  https://$LOCAL_IP:9443"
echo ""
echo "  Generated images → $BASE/invokeai-outputs/"
echo "  Drop PDFs here   → $BASE/papers/"
echo ""
echo "  First run? Pull models: bash $BASE/pull-models.sh"
echo ""
echo "  InvokeAI tip: upload jesus.safetensors via"
echo "  Model Manager → LoRA → Import"
STARTSH
chmod +x "$BASE/start.sh"

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

# ── status.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 ""
echo "=== Containers ==="
docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}"

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

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

echo ""
echo "=== Disk usage ==="
du -sh "$BASE"/*/  2>/dev/null
STATUSSH
chmod +x "$BASE/status.sh"

# ── Caddyfile ─────────────────────────────────────────────────────────────────
cat > "$BASE/Caddyfile.example" << CADDY
# =============================================================
# Caddy2 reverse proxy config
#
# RECOMMENDED: Use hostname instead of IP so it never breaks
# when your laptop's DHCP address changes.
# Your laptop hostname is: $(hostname)
# With mDNS it's reachable as: $(hostname).local
#
# OR set a static DHCP reservation in your router:
# tie your MAC address to a fixed IP like $LOCAL_IP
# =============================================================

# Option A — use hostname (works on most home networks, never breaks)
invokeai.yourdomain.com  { reverse_proxy $(hostname).local:9090 }
webui.yourdomain.com     { reverse_proxy $(hostname).local:3000 }
search.yourdomain.com    { reverse_proxy $(hostname).local:8080 }
rag.yourdomain.com       { reverse_proxy $(hostname).local:8001 }
portainer.yourdomain.com {
    reverse_proxy https://$(hostname).local:9443 {
        transport http { tls_insecure_skip_verify }
    }
}

# Option B — use IP (simpler but breaks if DHCP reassigns)
# invokeai.yourdomain.com  { reverse_proxy $LOCAL_IP:9090 }
# webui.yourdomain.com     { reverse_proxy $LOCAL_IP:3000 }
# search.yourdomain.com    { reverse_proxy $LOCAL_IP:8080 }
# rag.yourdomain.com       { reverse_proxy $LOCAL_IP:8001 }
# portainer.yourdomain.com {
#     reverse_proxy https://$LOCAL_IP:9443 {
#         transport http { tls_insecure_skip_verify }
#     }
# }
CADDY

# =============================================================================
sec "STEP 9 — Systemd Auto-Start"
# =============================================================================
sudo bash -c "cat > /etc/systemd/system/laptop-ai.service << SYSD
[Unit]
Description=Laptop AI Stack (InvokeAI + Ollama + Open WebUI + RAG)
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 laptop-ai.service
ok "Systemd auto-start enabled (laptop-ai.service)"

# =============================================================================
echo ""
echo -e "${G}════════════════════════════════════════════════════${N}"
echo -e "${G}  Setup complete!${N}"
echo -e "${G}════════════════════════════════════════════════════${N}"
echo ""
echo -e "  ${C}Everything lives under:${N} $BASE"
echo ""
echo -e "  ${C}NEXT STEPS${N}"
echo "  ──────────────────────────────────────────────────"
echo "  1. Start the stack:"
echo "       bash $BASE/start.sh"
echo ""
echo "  2. Pull AI models (10-40 min):"
echo "       bash $BASE/pull-models.sh"
echo ""
echo "  3. Upload jesus.safetensors to InvokeAI:"
echo "       Open http://$LOCAL_IP:9090"
echo "       → Model Manager → LoRA → Import"
echo ""
echo "  4. Drop PDFs into papers folder to auto-index:"
echo "       $BASE/papers/"
echo ""
echo "  5. Generated images saved to:"
echo "       $BASE/invokeai-outputs/"
echo ""
echo "  6. Add to Caddy2 on your proxy machine:"
echo "       cat $BASE/Caddyfile.example"
echo ""
echo -e "  ${C}SERVICE URLS${N}"
echo "  ──────────────────────────────────────────────────"
echo "  InvokeAI    →  http://$LOCAL_IP:9090"
echo "  Open WebUI  →  http://$LOCAL_IP:3000"
echo "  SearXNG     →  http://$LOCAL_IP:8080"
echo "  RAG Health  →  http://$LOCAL_IP:8001/health"
echo "  Portainer   →  https://$LOCAL_IP:9443"
echo ""
echo -e "  ${Y}NOTE: If docker commands fail, run: newgrp docker${N}"
echo -e "  ${Y}      or log out and back in.${N}"