deploy.sh

#!/usr/bin/env bash
# FLAI v8.8 — Single-Server Deployment Script

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# ── Colors ──
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
info()  { echo -e "${GREEN}[INFO]${NC} $1"; }
warn()  { echo -e "${YELLOW}[WARN]${NC} $1"; }
error() { echo -e "${RED}[ERROR]${NC} $1"; exit 1; }

# ── Prerequisites ──
check_prereqs() {
    info "Checking prerequisites..."
    command -v docker &>/dev/null     || error "Docker is not installed. Install Docker first."
    command -v docker compose &>/dev/null || error "Docker Compose plugin is not installed."
    command -v git &>/dev/null        || error "Git is not installed."
    if ! command -v nvidia-smi &>/dev/null; then
        warn "No NVIDIA GPU detected — GPU acceleration will not be available."
    fi
    if docker info 2>/dev/null | grep -q "Runtimes.*nvidia"; then
        info "NVIDIA Docker runtime detected."
    elif command -v nvidia-smi &>/dev/null; then
        warn "nvidia-smi found, but nvidia-container-runtime not detected in Docker."
    fi
    info "Prerequisites OK."
}

# ── Configuration ──
setup_env() {
    if [[ -f .env ]]; then
        warn ".env already exists — skipping configuration."
        return
    fi
    info "Creating .env from .env.example..."
    cp .env.example .env
    SECRET=$(python3 -c "import secrets; print(secrets.token_hex(32))" 2>/dev/null || openssl rand -hex 32)
    sed -i "s/SECRET_KEY=.*/SECRET_KEY=$SECRET/" .env
    info ".env created. Edit it to set models, URLs, and preferences."
}

# ── Data directories ──
setup_dirs() {
    info "Creating data directories..."
    mkdir -p data/uploads data/documents data/db_backups data/slm
    chown -R 1000:1000 data 2>/dev/null || warn "Could not chown data/ — run with sudo if needed."
}

# ── Validate .env ──
validate_env() {
    if [[ ! -f .env ]]; then
        warn ".env not found — run deploy.sh to create it first."
        return
    fi
    local errors=0
    grep -qP '^SECRET_KEY=\S+' .env 2>/dev/null || { warn "SECRET_KEY is not set in .env"; errors=1; }
    grep -qP '^DATABASE_URL=\S+' .env 2>/dev/null || { warn "DATABASE_URL is not set in .env"; errors=1; }
    [[ "$errors" -eq 1 ]] && warn "Edit .env with your settings before proceeding."
}

# ── Generate minimal llama-swap config ──
generate_llama_swap_config() {
    local CONFIG_DIR="llama-swap-config"
    mkdir -p "$CONFIG_DIR"

    # Don't overwrite if already generated by the web app
    if [[ -f "$CONFIG_DIR/llama-swap.yaml" ]]; then
        return
    fi

    info "Generating minimal llama-swap config..."
    cat > "$CONFIG_DIR/llama-swap.yaml" << 'YAML'
# Auto-generated by deploy.sh — replaced by web app on startup
default_model: chat
models:
  - model: /models/Qwen3-4B-Instruct-2507-Q4_K_M.gguf
    name: chat
    ctx_size: 8192
YAML
}

# ── Model download helpers ──
HF_DOWNLOAD() {
    local repo="$1" file="$2" dest="$3"
    local url="https://huggingface.co/$repo/resolve/main/$file"
    info "Downloading $file from HuggingFace..."
    local dir="$(dirname "$dest")"
    mkdir -p "$dir"
    
    # Retry download up to 3 times
    for attempt in 1 2 3; do
        curl -L --progress-bar -o "$dest" "$url"
        
        # Check if file is valid (more than 1KB)
        local size=$(stat -c%s "$dest" 2>/dev/null || echo "0")
        if [[ "$size" -gt 1024 ]]; then
            info "Successfully downloaded: $file ($size bytes)"
            return 0
        else
            warn "Attempt $attempt failed, retrying..."
            rm -f "$dest"
            sleep 2
        fi
    done
    error "Failed to download $file after 3 attempts"
}

# ── llama.cpp models ──
download_llamacpp_models() {
    info "Downloading llama.cpp models..."
    local MODEL_DIR="services/llamacpp/models"
    local VRAM_MB=0

    # Detect GPU VRAM for model selection
    if command -v nvidia-smi &>/dev/null; then
        VRAM_MB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -d ' ')
    fi

    # Chat/router model — Qwen3-4B-Instruct MXFP4 (~2 GB), fast for routing
    if [[ ! -f "$MODEL_DIR/Qwen3-4B-Instruct-2507-MXFP4_MOE.gguf" ]]; then
        info "Downloading Qwen3-4B-Instruct-2507-MXFP4_MOE.gguf (chat/router)..."
        HF_DOWNLOAD "unsloth/Qwen3-4B-Instruct-2507-GGUF" \
            "Qwen3-4B-Instruct-2507-MXFP4_MOE.gguf" \
            "$MODEL_DIR/Qwen3-4B-Instruct-2507-MXFP4_MOE.gguf"
    else
        warn "Qwen3-4B-Instruct-2507-MXFP4_MOE.gguf already exists — skipping."
    fi

    # Reasoning model — select based on VRAM
    if [[ "$VRAM_MB" -ge 16000 ]]; then
        # 16GB+ — full model
        if [[ ! -f "$MODEL_DIR/gpt-oss-20b-Q4_K_M.gguf" ]]; then
            info "Downloading gpt-oss-20b-Q4_K_M.gguf (reasoning, 16GB+ tier)..."
            HF_DOWNLOAD "unsloth/gpt-oss-20b-GGUF" \
                "gpt-oss-20b-Q4_K_M.gguf" \
                "$MODEL_DIR/gpt-oss-20b-Q4_K_M.gguf"
        else
            warn "gpt-oss-20b-Q4_K_M.gguf already exists — skipping."
        fi
    elif [[ "$VRAM_MB" -ge 12000 ]]; then
        # 12GB — Qwen3-8B-Thinking (~5GB, fits entirely with all layers on GPU)
        info "VRAM ${VRAM_MB}MB: downloading Qwen3-8B-Thinking (reasoning)..."
        if [[ ! -f "$MODEL_DIR/Qwen3-8B-Thinking-2507-Q4_K_M.gguf" ]]; then
            info "Downloading Qwen3-8B-Thinking-2507-Q4_K_M.gguf (reasoning, 12GB tier)..."
            HF_DOWNLOAD "unsloth/Qwen3-8B-Thinking-2507-GGUF" \
                "Qwen3-8B-Thinking-2507-Q4_K_M.gguf" \
                "$MODEL_DIR/Qwen3-8B-Thinking-2507-Q4_K_M.gguf" || \
            warn "8B thinking model download failed. Falling back to 4B."
            if [[ ! -f "$MODEL_DIR/Qwen3-8B-Thinking-2507-Q4_K_M.gguf" ]]; then
                HF_DOWNLOAD "unsloth/Qwen3-4B-Thinking-2507-GGUF" \
                    "Qwen3-4B-Thinking-2507-Q4_K_M.gguf" \
                    "$MODEL_DIR/Qwen3-4B-Thinking-2507-Q4_K_M.gguf"
            fi
        fi
    elif [[ "$VRAM_MB" -ge 8000 ]]; then
        # 8GB — Qwen3-4B-Thinking (~2.5GB)
        info "VRAM ${VRAM_MB}MB: downloading light reasoning model (Qwen3-4B-Thinking)..."
        if [[ ! -f "$MODEL_DIR/Qwen3-4B-Thinking-2507-Q4_K_M.gguf" ]]; then
            info "Downloading Qwen3-4B-Thinking-2507-Q4_K_M.gguf (light reasoning)..."
            HF_DOWNLOAD "unsloth/Qwen3-4B-Thinking-2507-GGUF" \
                "Qwen3-4B-Thinking-2507-Q4_K_M.gguf" \
                "$MODEL_DIR/Qwen3-4B-Thinking-2507-Q4_K_M.gguf" || \
            warn "Light reasoning model download failed."
        fi
    else
        # Unknown VRAM — download both
        info "VRAM unknown: downloading standard reasoning model..."
        if [[ ! -f "$MODEL_DIR/gpt-oss-20b-Q4_K_M.gguf" ]]; then
            HF_DOWNLOAD "unsloth/gpt-oss-20b-GGUF" \
                "gpt-oss-20b-Q4_K_M.gguf" \
                "$MODEL_DIR/gpt-oss-20b-Q4_K_M.gguf"
        fi
    fi

    # Multimodal model — select based on VRAM
    if [[ "$VRAM_MB" -ge 12000 ]]; then
        # 12GB+ — full multimodal (Qwen3VL-8B)
        if [[ ! -d "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M" ]]; then
            info "Downloading Qwen3VL-8B-Instruct-Q4_K_M (multimodal, 12GB+ tier)..."
            mkdir -p "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M"
            HF_DOWNLOAD "Qwen/Qwen3-VL-8B-Instruct-GGUF" \
                "Qwen3VL-8B-Instruct-Q4_K_M.gguf" \
                "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M/Qwen3VL-8B-Instruct-Q4_K_M.gguf"
            HF_DOWNLOAD "Qwen/Qwen3-VL-8B-Instruct-GGUF" \
                "mmproj-Qwen3VL-8B-Instruct-F16.gguf" \
                "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M/mmproj-F16.gguf"
        else
            warn "Qwen3VL-8B-Instruct-Q4_K_M already exists — skipping."
        fi
    elif [[ "$VRAM_MB" -ge 8000 ]]; then
        # 8GB — light multimodal (Qwen3VL-4B, ~2.5GB)
        info "VRAM ${VRAM_MB}MB: downloading light multimodal (Qwen3VL-4B)..."
        if [[ ! -d "$MODEL_DIR/Qwen3VL-4B-Instruct-Q4_K_M" ]]; then
            mkdir -p "$MODEL_DIR/Qwen3VL-4B-Instruct-Q4_K_M"
            HF_DOWNLOAD "Qwen/Qwen3-VL-4B-Instruct-GGUF" \
                "Qwen3VL-4B-Instruct-Q4_K_M.gguf" \
                "$MODEL_DIR/Qwen3VL-4B-Instruct-Q4_K_M/Qwen3VL-4B-Instruct-Q4_K_M.gguf" || \
            warn "Light multimodal download failed. Falling back to 8B version."
            HF_DOWNLOAD "Qwen/Qwen3-VL-4B-Instruct-GGUF" \
                "mmproj-Qwen3VL-4B-Instruct-F16.gguf" \
                "$MODEL_DIR/Qwen3VL-4B-Instruct-Q4_K_M/mmproj-F16.gguf" || true
            # Fallback to 8B if 4B not available
            if [[ ! -f "$MODEL_DIR/Qwen3VL-4B-Instruct-Q4_K_M/Qwen3VL-4B-Instruct-Q4_K_M.gguf" ]]; then
                warn "Qwen3VL-4B not available on HuggingFace, downloading 8B instead..."
                rm -rf "$MODEL_DIR/Qwen3VL-4B-Instruct-Q4_K_M"
                mkdir -p "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M"
                HF_DOWNLOAD "Qwen/Qwen3-VL-8B-Instruct-GGUF" \
                    "Qwen3VL-8B-Instruct-Q4_K_M.gguf" \
                    "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M/Qwen3VL-8B-Instruct-Q4_K_M.gguf"
                HF_DOWNLOAD "Qwen/Qwen3-VL-8B-Instruct-GGUF" \
                    "mmproj-Qwen3VL-8B-Instruct-F16.gguf" \
                    "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M/mmproj-F16.gguf"
            fi
        fi
    else
        # Unknown — download 8B
        if [[ ! -d "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M" ]]; then
            info "VRAM unknown: downloading Qwen3VL-8B-Instruct-Q4_K_M..."
            mkdir -p "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M"
            HF_DOWNLOAD "Qwen/Qwen3-VL-8B-Instruct-GGUF" \
                "Qwen3VL-8B-Instruct-Q4_K_M.gguf" \
                "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M/Qwen3VL-8B-Instruct-Q4_K_M.gguf"
            HF_DOWNLOAD "Qwen/Qwen3-VL-8B-Instruct-GGUF" \
                "mmproj-Qwen3VL-8B-Instruct-F16.gguf" \
                "$MODEL_DIR/Qwen3VL-8B-Instruct-Q4_K_M/mmproj-F16.gguf"
        fi
    fi

    # Embedding model — fits all tiers
    if [[ ! -f "$MODEL_DIR/bge-m3-Q8_0.gguf" ]]; then
        info "Downloading bge-m3-Q8_0.gguf (embeddings)..."
        HF_DOWNLOAD "gpustack/bge-m3-GGUF" \
            "bge-m3-Q8_0.gguf" \
            "$MODEL_DIR/bge-m3-Q8_0.gguf"
    else
        warn "bge-m3-Q8_0.gguf already exists — skipping."
    fi
}

# ── Stable Diffusion models ──
download_sd_cpp_models() {
    info "Downloading stable-diffusion.cpp models..."
    local DIFF_DIR="services/sd_cpp/models/diffusion_models"
    local VAE_DIR="services/sd_cpp/models/vae"
    local TXT_DIR="services/sd_cpp/models/text_encoders"

    # Z-Image Turbo (image generation)
    if [[ ! -f "$DIFF_DIR/z_image_turbo-Q8_0.gguf" ]]; then
        info "Downloading z_image_turbo-Q8_0.gguf..."
        HF_DOWNLOAD "leejet/Z-Image-Turbo-GGUF" \
            "z_image_turbo-Q8_0.gguf" "${DIFF_DIR}/z_image_turbo-Q8_0.gguf"
    else
        warn "z_image_turbo-Q8_0.gguf already exists — skipping."
    fi

    # Flux.2 Klein 4B (image editing)
    if [[ ! -f "$DIFF_DIR/flux-2-klein-4b-Q8_0.gguf" ]]; then
        info "Downloading flux-2-klein-4b-Q8_0.gguf (editing)..."
        HF_DOWNLOAD "leejet/FLUX.2-klein-4B-GGUF" \
            "flux-2-klein-4b-Q8_0.gguf" "${DIFF_DIR}/flux-2-klein-4b-Q8_0.gguf"
    else
        warn "flux-2-klein-4b-Q8_0.gguf already exists — skipping."
    fi

    # VAE (for generation)
    if [[ ! -f "$VAE_DIR/ae.safetensors" ]]; then
        info "Downloading ae.safetensors (VAE for generation)..."
        HF_DOWNLOAD "Comfy-Org/z_image_turbo" \
            "split_files/vae/ae.safetensors" "${VAE_DIR}/ae.safetensors"
    else
        warn "ae.safetensors already exists — skipping."
    fi

    # VAE (for editing)
    if [[ ! -f "$VAE_DIR/flux2_ae.safetensors" ]]; then
        info "Downloading flux2_ae.safetensors (VAE for editing)..."
        HF_DOWNLOAD "Comfy-Org/flux2-dev" \
            "split_files/vae/flux2-vae.safetensors" \
            "${VAE_DIR}/flux2_ae.safetensors"
    else
        warn "flux2_ae.safetensors already exists — skipping."
    fi

    # Text encoder (shared with chat model)
    if [[ ! -f "$TXT_DIR/Qwen3-4B-Instruct-2507-Q4_K_M.gguf" ]]; then
        info "Downloading Qwen3-4B-Instruct-2507-Q4_K_M.gguf (text encoder)..."
        # Symlink from chat model if already downloaded
        if [[ -f "services/llamacpp/models/Qwen3-4B-Instruct-2507-Q4_K_M.gguf" ]]; then
            info "Chat model already downloaded — creating symlink..."
            ln -sf "../../llamacpp/models/Qwen3-4B-Instruct-2507-Q4_K_M.gguf" \
                "$TXT_DIR/Qwen3-4B-Instruct-2507-Q4_K_M.gguf"
        else
            HF_DOWNLOAD "unsloth/Qwen3-4B-Instruct-2507-GGUF" \
                "Qwen3-4B-Instruct-2507-Q4_K_M.gguf" \
                "${TXT_DIR}/Qwen3-4B-Instruct-2507-Q4_K_M.gguf"
        fi
    else
        warn "Qwen3-4B-Instruct-2507-Q4_K_M.gguf already exists — skipping."
    fi
}

# ── LTX-Video models ──
download_ltx_video_models() {
    info "Downloading LTX-Video models..."
    local VIDEO_DIR="services/ltx_video/models"

    # LTX-Video distilled model (2B params)
    if [[ ! -f "$VIDEO_DIR/ltxv-2b-0.9.8-distilled.safetensors" ]]; then
        info "Downloading ltxv-2b-0.9.8-distilled.safetensors..."
        HF_DOWNLOAD "Lightricks/LTX-Video" \
            "ltxv-2b-0.9.8-distilled.safetensors" \
            "$VIDEO_DIR/ltxv-2b-0.9.8-distilled.safetensors"
    else
        warn "ltxv-2b-0.9.8-distilled.safetensors already exists — skipping."
    fi

    # T5 text encoder (PixArt T5 ≈ 8.9 GB)
    if [[ ! -d "$VIDEO_DIR/t5_encoder/text_encoder" ]]; then
        info "Downloading T5 text encoder (PixArt T5-XXL, ~18 GB on disk)…"
        bash "$SCRIPT_DIR/services/ltx_video/download-t5-encoder.sh"
    else
        warn "T5 text encoder already exists — skipping."
    fi
}

# ── TTS (Speech Synthesis) models ──
download_tts_models() {
    info "Downloading TTS (Piper) models..."
    local TTS_DIR="services/piper/piper_models"
    mkdir -p "$TTS_DIR"
    # English (male)
    if [[ ! -f "$TTS_DIR/en_US-ryan-medium.onnx" ]]; then
        info "Downloading en_US-ryan-medium.onnx..."
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "en/en_US/ryan/medium/en_US-ryan-medium.onnx" \
            "$TTS_DIR/en_US-ryan-medium.onnx"
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "en/en_US/ryan/medium/en_US-ryan-medium.onnx.json" \
            "$TTS_DIR/en_US-ryan-medium.onnx.json"
    fi
    # English (female)
    if [[ ! -f "$TTS_DIR/en_US-ljspeech-medium.onnx" ]]; then
        info "Downloading en_US-ljspeech-medium.onnx..."
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "en/en_US/ljspeech/medium/en_US-ljspeech-medium.onnx" \
            "$TTS_DIR/en_US-ljspeech-medium.onnx"
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "en/en_US/ljspeech/medium/en_US-ljspeech-medium.onnx.json" \
            "$TTS_DIR/en_US-ljspeech-medium.onnx.json"
    fi
    # Russian (male)
    if [[ ! -f "$TTS_DIR/ru_RU-dmitri-medium.onnx" ]]; then
        info "Downloading ru_RU-dmitri-medium.onnx..."
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "ru/ru_RU/dmitri/medium/ru_RU-dmitri-medium.onnx" \
            "$TTS_DIR/ru_RU-dmitri-medium.onnx"
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "ru/ru_RU/dmitri/medium/ru_RU-dmitri-medium.onnx.json" \
            "$TTS_DIR/ru_RU-dmitri-medium.onnx.json"
    fi
    # Russian (female)
    if [[ ! -f "$TTS_DIR/ru_RU-irina-medium.onnx" ]]; then
        info "Downloading ru_RU-irina-medium.onnx..."
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "ru/ru_RU/irina/medium/ru_RU-irina-medium.onnx" \
            "$TTS_DIR/ru_RU-irina-medium.onnx"
        HF_DOWNLOAD "rhasspy/piper-voices" \
            "ru/ru_RU/irina/medium/ru_RU-irina-medium.onnx.json" \
            "$TTS_DIR/ru_RU-irina-medium.onnx.json"
    fi
}

# ── Whisper ASR Models ──
download_whisper_models() {
    info "Downloading Whisper ASR model (Systran/faster-whisper-medium)..."
    local CACHE_DIR="data/hf-cache"
    mkdir -p "$CACHE_DIR/hub"

    local MODEL_DIR="$CACHE_DIR/hub/models--Systran--faster-whisper-medium"
    if [[ -d "$MODEL_DIR" && -f "$MODEL_DIR/snapshots/"*"/model.bin" ]]; then
        info "Whisper model already downloaded."
        return 0
    fi

    info "Downloading ~1.5 GB — this may take several minutes..."
    python3 -c "
from huggingface_hub import snapshot_download
import os, sys

os.environ['HF_HUB_DOWNLOAD_TIMEOUT'] = '600'
try:
    path = snapshot_download(
        'Systran/faster-whisper-medium',
        cache_dir='$CACHE_DIR',
        ignore_patterns=['*.h5', '*.ot', '*.msgpack']
    )
    print(f'OK: model downloaded to {path}')
except Exception as e:
    print(f'ERROR: {e}')
    sys.exit(1)
" && info "Whisper model downloaded successfully." || warn "Failed to download Whisper model. ASR will be unavailable."
}

# ── Build & Launch ──
build_and_launch() {
    local PROFILE=""
    [[ "$WITH_IMAGE_GEN" == "true" ]] && PROFILE="$PROFILE --profile with-image-gen"
    [[ "$WITH_VOICE" == "true" ]] && PROFILE="$PROFILE --profile with-voice"
    [[ "$WITH_RAG" == "true" ]]    && PROFILE="$PROFILE --profile with-rag"
    [[ "$WITH_VIDEO" == "true" ]]  && PROFILE="$PROFILE --profile with-video"
    [[ "$WITH_SLM" == "true" ]]    && PROFILE="$PROFILE --profile with-slm"

    COMPOSE_FILE="docker-compose.gpu.yml"
    info "GPU mode — using GPU compose file."

    # Detect GPU VRAM tier
    local VRAM_MB=0
    if command -v nvidia-smi &>/dev/null; then
        VRAM_MB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -d ' ')
    fi
    if [[ "$VRAM_MB" -ge 16000 ]]; then
        info "GPU VRAM: ${VRAM_MB}MB (tier: 16GB+) — full performance"
    elif [[ "$VRAM_MB" -ge 12000 ]]; then
        info "GPU VRAM: ${VRAM_MB}MB (tier: 12GB) — moderate, consider light reasoning model"
    elif [[ "$VRAM_MB" -ge 8000 ]]; then
        info "GPU VRAM: ${VRAM_MB}MB (tier: 8GB) — limited, consider light reasoning model"
    else
        warn "GPU VRAM: ${VRAM_MB}MB — below 8GB minimum. Performance will be severely limited."
    fi

    # Clean up old containers to avoid runtime conflicts
    info "Stopping old containers (if any)..."
    docker compose -f "$COMPOSE_FILE" down --remove-orphans 2>/dev/null || true

    info "Building Docker images..."
    docker compose -f "$COMPOSE_FILE" $PROFILE build

    info "Starting services..."
    docker compose -f "$COMPOSE_FILE" $PROFILE up -d

    info "Waiting for services to start..."
    sleep 10

    local STATUS
    STATUS=$(curl -s http://localhost:5000/health 2>/dev/null | python3 -c "import sys,json; print(json.load(sys.stdin).get('status','unknown'))" 2>/dev/null || echo "unreachable")
    if [[ "$STATUS" == "ok" ]]; then
        info "FLAI v8.8 is running! Open http://localhost:5000 in your browser."
    else
        warn "Health check returned: $STATUS — check 'docker compose logs' for details."
    fi
}

# ── Tests ──
run_tests() {
    info "Running unit tests..."
    pip install -e ".[test]" -q 2>/dev/null || pip install -r requirements.txt -q 2>/dev/null
    pytest --tb=short -q 2>&1 || warn "Some tests failed — check output above."
}

# ── Usage ──
usage() {
    cat <<'USAGE'
FLAI v8.8 — Deployment Script

Usage: ./deploy.sh [OPTIONS]

Options:
  --with-voice        Deploy Whisper ASR + Piper TTS
  --with-rag          Deploy Qdrant for RAG (document search)
  --with-image-gen    Deploy stable-diffusion.cpp for image generation/editing
                      (default: disabled, use this flag to enable)
  --with-video        Deploy LTX-Video for video generation
  --with-slm          Deploy SuperLocalMemory for long-term memory
                      (default: disabled, use this flag to enable)
  --download-models   Download GGUF/safetensors models from HuggingFace
  --run-tests         Run unit tests after deployment
  --help, -h          Show this help message

Model Download Sizes (approximate):
  llama.cpp (16GB+ tier — full models):
    Qwen3-4B-Instruct (chat)         ~2.5 GB
    gpt-oss-20b (reasoning)          ~12 GB
    Qwen3VL-8B (multimodal)          ~5.5 GB
    bge-m3 (embeddings)              ~1.5 GB
  llama.cpp (8GB tier — light models):
    Qwen3-4B-Instruct (chat)         ~2.5 GB
    Qwen3-4B-Thinking (reasoning)    ~2.5 GB
    Qwen3VL-4B (multimodal)          ~2.5 GB
    bge-m3 (embeddings)              ~1.5 GB
  llama.cpp (12GB tier — mid models):
    Qwen3-4B-Instruct (chat)         ~2.5 GB
    Qwen3-8B-Thinking (reasoning)    ~5 GB
    Qwen3VL-8B (multimodal)          ~5.5 GB
    bge-m3 (embeddings)              ~1.5 GB
  Image generation (Z-Image Turbo)   ~6.5 GB
  Image editing (Flux.2 Klein 4B)    ~5 GB
  Video generation (LTX-Video 2B)    ~5.9 GB
  T5 text encoder (PixArt T5-XXL)    ~18 GB (on disk, float32)
  TTS (Piper)                        ~0.2 GB
  SLM embedding model                ~0.5 GB (pre-downloaded during Docker build)
USAGE
}

# ── Parse arguments ──
WITH_VOICE=false
WITH_RAG=false
WITH_IMAGE_GEN=false
WITH_VIDEO=false
WITH_SLM=false
DOWNLOAD_MODELS=false
RUN_TESTS=false

for arg in "$@"; do
    case "$arg" in
        --with-voice)     WITH_VOICE=true ;;
        --with-rag)       WITH_RAG=true ;;
        --with-image-gen) WITH_IMAGE_GEN=true ;;
        --with-video)     WITH_VIDEO=true ;;
        --with-slm)       WITH_SLM=true ;;
        --download-models) DOWNLOAD_MODELS=true ;;
        --run-tests)      RUN_TESTS=true ;;
        --help|-h)        usage; exit 0 ;;
    esac
done

# ── Main ──
main() {
    echo "============================================"
    echo "  FLAI v8.8 — Deployment Script"
    echo "============================================"
    echo ""

    check_prereqs
    setup_dirs
    setup_env
    validate_env
    generate_llama_swap_config

    if [[ "$DOWNLOAD_MODELS" == "true" ]]; then
        download_llamacpp_models
        download_sd_cpp_models
        download_tts_models
        [[ "$WITH_VIDEO" == "true" ]] && download_ltx_video_models
        [[ "$WITH_VOICE" == "true" ]] && download_whisper_models
    fi

    build_and_launch

    if [[ "$RUN_TESTS" == "true" ]]; then
        run_tests
    fi

    echo ""
    echo "============================================"
    info "Deployment complete!"
    echo "============================================"
    echo "  Web UI:   http://localhost:5000"
    echo "  Health:   http://localhost:5000/health"
    echo "  Logs:     docker compose -f $COMPOSE_FILE logs -f"
    echo "  Stop:     docker compose -f $COMPOSE_FILE down --remove-orphans"
    echo "============================================"
}

main "$@"