#!/bin/bash
# =============================================================================
# AI Stack Update Script — Safe to run on existing installation
# Adds: Kiwix (offline Wikipedia) + Gitea (self-hosted git)
# Never overwrites: docker-compose.yml, searxng config, server.py
# Usage: bash update-ai-stack.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}"; }

BASE="$HOME/docker/ai-stack"
COMPOSE="$BASE/docker-compose.yml"
LOCAL_IP=$(ip route get 1.1.1.1 2>/dev/null | grep -oP 'src \K\S+' || hostname -I | awk '{print $1}')

# Verify this is an existing installation
if [[ ! -f "$COMPOSE" ]]; then
    echo "Error: $COMPOSE not found. Run the main setup script first."
    exit 1
fi

echo ""
echo -e "  ${C}AI Stack Update${N}"
echo "  Base: $BASE"
echo "  IP:   $LOCAL_IP"
echo ""
read -rp "  Proceed? (Y/n): " CONFIRM
[[ "${CONFIRM,,}" == "n" ]] && exit 0

# =============================================================================
sec "STEP 1 — Create new directories"
# =============================================================================
mkdir -p "$BASE"/{kiwix,gitea}
ok "Directories ready"

# =============================================================================
sec "STEP 2 — Add Kiwix to docker-compose.yml"
# =============================================================================
if grep -q "kiwix" "$COMPOSE"; then
    ok "Kiwix already in docker-compose.yml — skipping"
else
    inf "Adding Kiwix service..."
    # Insert before the volumes: section
    python3 - << PYEOF
import re

with open('$COMPOSE', 'r') as f:
    content = f.read()

kiwix_service = '''
  # ── Kiwix — Offline Wikipedia ─────────────────────────────────────────────
  # ZIM file goes in: $BASE/kiwix/
  # Download: wget -c 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"

'''

# Insert before volumes: section
content = re.sub(r'(\nvolumes:)', kiwix_service + r'\1', content, count=1)

with open('$COMPOSE', 'w') as f:
    f.write(content)
PYEOF
    ok "Kiwix added to docker-compose.yml"
fi

# =============================================================================
sec "STEP 3 — Add Gitea to docker-compose.yml"
# =============================================================================
if grep -q "gitea" "$COMPOSE"; then
    ok "Gitea already in docker-compose.yml — skipping"
else
    inf "Adding Gitea service..."
    read -rp "  Your domain for Gitea (e.g. git.yourdomain.com): " GITEA_DOMAIN
    GITEA_DOMAIN="${GITEA_DOMAIN:-git.yourdomain.com}"

    python3 - << PYEOF
import re

with open('$COMPOSE', 'r') as f:
    content = f.read()

gitea_service = '''
  # ── Gitea — Self-hosted Git ────────────────────────────────────────────────
  # First run: go to http://$LOCAL_IP:3001 to complete setup
  # Mirror GitHub repos: Gitea → Explore → Migrate → GitHub
  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://$GITEA_DOMAIN
      - GITEA__server__DOMAIN=$GITEA_DOMAIN
      - GITEA__server__SSH_DOMAIN=$GITEA_DOMAIN
      - GITEA__webhook__ALLOWED_HOST_LIST=rag-server

'''

content = re.sub(r'(\nvolumes:)', gitea_service + r'\1', content, count=1)

with open('$COMPOSE', 'w') as f:
    f.write(content)
PYEOF
    ok "Gitea added to docker-compose.yml"
fi

# =============================================================================
sec "STEP 4 — Fix SearXNG port in docker-compose.yml"
# =============================================================================
if grep -q "8080:8080" "$COMPOSE"; then
    inf "Fixing SearXNG port 8080 → 8888..."
    sed -i 's/0\.0\.0\.0:8080:8080/0.0.0.0:8888:8080/g' "$COMPOSE"
    ok "SearXNG port fixed"
else
    ok "SearXNG port already correct — skipping"
fi

# =============================================================================
sec "STEP 5 — Fix InvokeAI precision fp16 → float16"
# =============================================================================
if grep -q "INVOKEAI_PRECISION=fp16" "$COMPOSE"; then
    inf "Fixing InvokeAI precision..."
    sed -i 's/INVOKEAI_PRECISION=fp16/INVOKEAI_PRECISION=float16/g' "$COMPOSE"
    ok "InvokeAI precision fixed"
else
    ok "InvokeAI precision already correct — skipping"
fi

# =============================================================================
sec "STEP 6 — UFW rules for new services"
# =============================================================================
read -rp "  Your LAN subnet [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"

# Only add rules that don't exist yet
add_ufw_rule() {
    local port=$1 proto=$2 comment=$3
    if ! sudo ufw status | grep -q "$port/$proto"; then
        sudo ufw allow from "$LAN_SUBNET" to any port "$port" proto "$proto" comment "$comment" > /dev/null
        ok "UFW: opened port $port ($comment)"
    else
        ok "UFW: port $port already open — skipping"
    fi
}

add_ufw_rule 8181 tcp "Kiwix"
add_ufw_rule 3001 tcp "Gitea"
add_ufw_rule 2222 tcp "Gitea SSH"
add_ufw_rule 8888 tcp "SearXNG"
sudo ufw reload > /dev/null
ok "UFW updated"

# =============================================================================
sec "STEP 7 — Start new services"
# =============================================================================
cd "$BASE"
inf "Starting Kiwix and Gitea..."
docker compose up -d kiwix gitea

inf "Restarting fixed services (InvokeAI, SearXNG)..."
docker compose up -d invokeai searxng

ok "Services started"

# =============================================================================
sec "STEP 8 — Status check"
# =============================================================================
echo ""
echo "Waiting 10 seconds for services to initialize..."
sleep 10
docker compose ps

# =============================================================================
echo ""
echo -e "${G}════════════════════════════════════════════════════${N}"
echo -e "${G}  Update complete!${N}"
echo -e "${G}════════════════════════════════════════════════════${N}"
echo ""
echo -e "  ${C}NEW SERVICES${N}"
echo "  ──────────────────────────────────────────────────"
echo "  Kiwix       →  http://$LOCAL_IP:8181"
echo "  Gitea       →  http://$LOCAL_IP:3001"
echo ""
echo -e "  ${C}FIXED SERVICES${N}"
echo "  ──────────────────────────────────────────────────"
echo "  InvokeAI    →  http://$LOCAL_IP:9090  (precision fix)"
echo "  SearXNG     →  http://$LOCAL_IP:8888  (port fix)"
echo ""
echo -e "  ${C}NEXT STEPS${N}"
echo "  ──────────────────────────────────────────────────"
echo "  1. Set up Gitea at http://$LOCAL_IP:3001"
echo "     → Complete install wizard"
echo "     → Create admin account"
echo "     → Mirror your GitHub repos"
echo ""
echo "  2. Download Wikipedia ZIM (22GB — use torrent for best results):"
echo "     cd $BASE/kiwix"
echo "     # Via wget:"
echo "     wget -c https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic_2026-01.zim"
echo "     # Via torrent (faster + resumes):"
echo "     sudo apt install transmission-cli"
echo "     transmission-cli https://download.kiwix.org/zim/wikipedia/wikipedia_en_all_nopic_2026-01.zim.torrent --download-dir $BASE/kiwix"
echo ""
echo "  3. Update Caddy config — add to your Caddyfile:"
echo "     wiki.yourdomain.com { reverse_proxy $LOCAL_IP:8181 }"
echo "     git.yourdomain.com  { reverse_proxy $LOCAL_IP:3001 }"
echo ""
echo "  4. Update Open WebUI SearXNG URL to:"
echo "     http://searxng:8080/search?q=<query>&format=json"
echo ""
echo "  5. Set up Gitea webhook for RAG auto-indexing:"
echo "     Repo → Settings → Webhooks → Add"
echo "     URL: http://rag-server:8001/webhook/gitea"
echo ""
echo -e "  ${Y}NOTE: Wikipedia download is 22GB — start it and walk away${N}"