#!/bin/bash
# install-linux.sh — Full setup for PRE (Personal Reasoning Engine) on Linux
#
# This script handles everything from Ollama to ready-to-run:
#   0. Checks system requirements (distro, GPU/VRAM, RAM, disk)
#   1. Installs/verifies Ollama
#   1b. Configures Ollama environment (FA, KV cache, keep-alive)
#   2. Pulls the base model (q8_0 ~28GB or q4_K_M ~15GB, auto-selected by VRAM)
#   2b. Pulls embedding model (nomic-embed-text for experience ledger + RAG)
#   3. Creates optimized custom model (pre-gemma4) from Modelfile
#   4. Checks for Node.js
#   5. Sets up Web GUI (npm install)
#   6. Sets up ~/.pre/ directories
#   7. Auto-sizes context window based on RAM headroom
#   8. Optional: voice tools (Whisper, FFmpeg, espeak-ng)
#   9. Optional: GNOME PIM integration (evolution-data-server)
#   10. Optional: desktop automation deps (xdotool, scrot, xclip)
#   11. Optional: auto-start (systemd user service)
#   12. Pre-warms model into GPU memory
#   13. Creates ~/bin/pre launcher script
#
# Target distros: Ubuntu/Debian (primary), Fedora/RHEL (secondary), Arch (best-effort)
#
# Usage:
#   ./install-linux.sh           Interactive install
#   ./install-linux.sh --yes     Non-interactive: accept all defaults, skip optional steps
#   ./install-linux.sh -y        Same as --yes
#   ./install-linux.sh --llamacpp  Use llama.cpp server as inference backend (instead of Ollama)

set -e

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
RESET='\033[0m'

REPO_DIR="$(cd "$(dirname "$0")" && pwd)"
ENGINE_DIR="$REPO_DIR/engine"
CUSTOM_MODEL="pre-gemma4"

# --- Flag parsing ---
AUTO_YES=false
FORCE_LLAMACPP=false
for arg in "$@"; do
    case "$arg" in
        --yes|-y) AUTO_YES=true ;;
        --llamacpp) FORCE_LLAMACPP=true ;;
    esac
done

ask_yn() {
    local prompt="$1" default="$2"
    if [ "$AUTO_YES" = true ]; then
        [ "$default" = "Y" ] && return 0 || return 1
    fi
    read -p "$prompt" -n 1 -r
    echo ""
    if [ "$default" = "Y" ]; then
        [[ ! $REPLY =~ ^[Nn]$ ]]
    else
        [[ $REPLY =~ ^[Yy]$ ]]
    fi
}

step() { echo -e "\n${BOLD}${CYAN}=== $1 ===${RESET}\n"; }
ok()   { echo -e "${GREEN}$1${RESET}"; }
warn() { echo -e "${YELLOW}$1${RESET}"; }
fail() { echo -e "${RED}$1${RESET}"; exit 1; }

trap 'echo -e "\n${RED}Installation failed at line $LINENO.${RESET}\n${DIM}Re-run the script to retry — it will pick up where it left off.${RESET}"' ERR

# ── Distro detection ──
detect_pkg_manager() {
    if command -v apt-get &>/dev/null; then
        PKG_MGR="apt"
        PKG_INSTALL="sudo apt-get install -y"
        PKG_UPDATE="sudo apt-get update"
    elif command -v dnf &>/dev/null; then
        PKG_MGR="dnf"
        PKG_INSTALL="sudo dnf install -y"
        PKG_UPDATE="sudo dnf check-update || true"
    elif command -v pacman &>/dev/null; then
        PKG_MGR="pacman"
        PKG_INSTALL="sudo pacman -S --noconfirm"
        PKG_UPDATE="sudo pacman -Sy"
    else
        PKG_MGR="unknown"
        PKG_INSTALL=""
        PKG_UPDATE=""
    fi
}

detect_pkg_manager

# ============================================================================
# Step 0: System requirements
# ============================================================================
step "Checking system requirements"

OS=$(uname -s)
if [ "$OS" != "Linux" ]; then
    fail "This installer is for Linux. Detected: $OS. Use install.sh for macOS."
fi

# Distro info
if [ -f /etc/os-release ]; then
    . /etc/os-release
    DISTRO_NAME="$NAME"
    DISTRO_VER="$VERSION_ID"
    ok "  Distro: $DISTRO_NAME $DISTRO_VER"
else
    DISTRO_NAME="Unknown"
    warn "  Unable to detect distro (no /etc/os-release)"
fi

ARCH=$(uname -m)
if [ "$ARCH" = "x86_64" ] || [ "$ARCH" = "aarch64" ]; then
    ok "  Architecture: $ARCH"
else
    fail "Unsupported architecture: $ARCH (need x86_64 or aarch64)"
fi

# ── GPU detection ──
GPU_BACKEND=""
GPU_NAME=""
GPU_VRAM_GB=0

# NVIDIA GPU detection (native — no Docker needed on Linux)
if command -v nvidia-smi &>/dev/null; then
    NVOUT=$(nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null || true)
    if [ -n "$NVOUT" ]; then
        GPU_BACKEND="cuda"
        GPU_NAME=$(echo "$NVOUT" | cut -d',' -f1 | xargs)
        VRAM_MIB=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 | xargs)
        if [ -n "$VRAM_MIB" ] && [ "$VRAM_MIB" -gt 0 ] 2>/dev/null; then
            GPU_VRAM_GB=$((VRAM_MIB / 1024))
        fi
        ok "  GPU: $GPU_NAME (NVIDIA CUDA)"
        if [ "$GPU_VRAM_GB" -gt 0 ]; then
            ok "  VRAM: ${GPU_VRAM_GB}GB"
        fi
    fi
fi

# AMD GPU detection (ROCm — future support)
if [ -z "$GPU_BACKEND" ] && command -v rocm-smi &>/dev/null; then
    AMD_NAME=$(rocm-smi --showproductname 2>/dev/null | grep -i "card" | head -1 | sed 's/.*: //' || true)
    if [ -n "$AMD_NAME" ]; then
        GPU_BACKEND="rocm"
        GPU_NAME="$AMD_NAME"
        # ROCm VRAM detection
        AMD_VRAM=$(rocm-smi --showmeminfo vram 2>/dev/null | grep "Total" | head -1 | grep -oE '[0-9]+' || true)
        if [ -n "$AMD_VRAM" ] && [ "$AMD_VRAM" -gt 0 ] 2>/dev/null; then
            GPU_VRAM_GB=$((AMD_VRAM / 1024 / 1024 / 1024))
        fi
        ok "  GPU: $GPU_NAME (AMD ROCm)"
        if [ "$GPU_VRAM_GB" -gt 0 ]; then
            ok "  VRAM: ${GPU_VRAM_GB}GB"
        fi
    fi
fi

# Fallback: lspci detection
if [ -z "$GPU_BACKEND" ]; then
    LSPCI_GPU=$(lspci 2>/dev/null | grep -i 'vga\|3d\|display' | head -1 || true)
    if [ -n "$LSPCI_GPU" ]; then
        warn "  GPU: $LSPCI_GPU (no CUDA/ROCm driver detected)"
        warn "  Install NVIDIA or AMD GPU drivers for GPU acceleration."
    else
        warn "  No GPU detected — Ollama will use CPU only (very slow for 26B model)."
    fi
fi

# RAM
RAM_KB=$(grep MemTotal /proc/meminfo | awk '{print $2}')
RAM_GB=$((RAM_KB / 1048576))

if [ "$RAM_GB" -lt 16 ]; then
    fail "PRE requires at least 16GB RAM. Detected: ${RAM_GB}GB"
fi
if [ "$RAM_GB" -lt 32 ]; then
    warn "  RAM: ${RAM_GB}GB — functional, but 32GB+ recommended."
else
    ok "  RAM: ${RAM_GB}GB"
fi

# ── Quantization selection (VRAM-aware) ──
if [ "$GPU_VRAM_GB" -ge 28 ]; then
    BASE_MODEL="gemma4:26b-a4b-it-q8_0"
    QUANT="q8_0"
    MODEL_SIZE_GB=28
    ok "  Quant: q8_0 (near-lossless, ~28GB) — fits in ${GPU_VRAM_GB}GB VRAM"
elif [ "$GPU_VRAM_GB" -gt 0 ]; then
    BASE_MODEL="gemma4:26b-a4b-it-q4_K_M"
    QUANT="q4_K_M"
    MODEL_SIZE_GB=15
    ok "  Quant: q4_K_M (~15GB) — fits in ${GPU_VRAM_GB}GB VRAM for full GPU acceleration"
elif [ "$RAM_GB" -ge 32 ]; then
    # No GPU detected — CPU mode, enough RAM for q8_0
    BASE_MODEL="gemma4:26b-a4b-it-q8_0"
    QUANT="q8_0"
    MODEL_SIZE_GB=28
    warn "  Quant: q8_0 (CPU mode — inference will be slow without GPU)"
else
    BASE_MODEL="gemma4:26b-a4b-it-q4_K_M"
    QUANT="q4_K_M"
    MODEL_SIZE_GB=15
    warn "  Quant: q4_K_M (CPU mode, limited RAM)"
fi

# Context window sizing (headroom-based, same thresholds as macOS/Windows)
HEADROOM=$((RAM_GB - MODEL_SIZE_GB))
if [ "$HEADROOM" -ge 68 ]; then
    OPTIMAL_CTX=131072
elif [ "$HEADROOM" -ge 36 ]; then
    OPTIMAL_CTX=65536
elif [ "$HEADROOM" -ge 20 ]; then
    OPTIMAL_CTX=32768
elif [ "$HEADROOM" -ge 8 ]; then
    OPTIMAL_CTX=16384
elif [ "$HEADROOM" -ge 4 ]; then
    OPTIMAL_CTX=8192
else
    OPTIMAL_CTX=4096
fi
CTX_HUMAN="$(( OPTIMAL_CTX / 1024 ))K"
ok "  Context window: ${CTX_HUMAN} (${OPTIMAL_CTX} tokens, ${HEADROOM}GB headroom after ${QUANT} model)"

# Disk space
AVAIL_KB=$(df -k "$REPO_DIR" | tail -1 | awk '{print $4}')
AVAIL_GB=$((AVAIL_KB / 1048576))
DISK_NEEDED=$((MODEL_SIZE_GB + 5))
if [ "$AVAIL_GB" -lt "$DISK_NEEDED" ]; then
    warn "  Disk: ${AVAIL_GB}GB available — need ~${MODEL_SIZE_GB}GB for ${QUANT} model download."
    if ! ask_yn "  Continue anyway? [y/N] " "N"; then exit 1; fi
else
    ok "  Disk: ${AVAIL_GB}GB available"
fi

echo ""
ok "System requirements met."

# ============================================================================
# Step 0b: Inference backend selection
# ============================================================================
step "Inference backend"

# Smart default: NVIDIA GPU (CUDA) → llama.cpp for Flash Attention + grammar;
# No GPU or AMD (ROCm) → Ollama (simpler, better ROCm compatibility).
if [ "$GPU_BACKEND" = "cuda" ]; then
    SMART_DEFAULT="llamacpp"
    SMART_REASON="NVIDIA GPU detected — llama.cpp enables Flash Attention and grammar constraints for faster, more reliable inference."
else
    SMART_DEFAULT="ollama"
    if [ "$GPU_BACKEND" = "rocm" ]; then
        SMART_REASON="AMD GPU detected — Ollama provides better ROCm compatibility."
    else
        SMART_REASON="No dedicated GPU detected — Ollama provides simpler setup for CPU inference."
    fi
fi

BACKEND="$SMART_DEFAULT"
if [ "$FORCE_LLAMACPP" = true ]; then
    BACKEND="llamacpp"
    ok "  Backend: llama.cpp (--llamacpp flag)"
elif [ "$AUTO_YES" = true ]; then
    BACKEND="$SMART_DEFAULT"
    ok "  Backend: $([ "$SMART_DEFAULT" = "llamacpp" ] && echo "llama.cpp" || echo "Ollama") (auto — $SMART_REASON)"
else
    echo ""
    echo -e "  ${DIM}$SMART_REASON${RESET}"
    echo ""
    echo -e "  ${BOLD}Choose inference backend:${RESET}"
    echo ""
    if [ "$SMART_DEFAULT" = "llamacpp" ]; then
        echo -e "  ${CYAN}1)${RESET} ${BOLD}llama.cpp${RESET} (recommended) — Direct GPU inference"
        echo -e "     Flash Attention, grammar-constrained output, lower overhead."
        echo -e "     Ollama still installed for embeddings."
        echo ""
        echo -e "  ${CYAN}2)${RESET} ${BOLD}Ollama${RESET} — Managed model server"
        echo -e "     Easy setup, automatic model management, wide compatibility."
    else
        echo -e "  ${CYAN}1)${RESET} ${BOLD}Ollama${RESET} (recommended) — Managed model server"
        echo -e "     Easy setup, automatic model management, wide compatibility."
        echo ""
        echo -e "  ${CYAN}2)${RESET} ${BOLD}llama.cpp${RESET} — Direct GPU inference"
        echo -e "     Lower overhead, grammar-constrained output, speculative decoding."
        echo -e "     Ollama still installed for embeddings."
    fi
    echo ""
    BACKEND_CHOICE=""
    read -p "  Choice [1/2, default=1]: " -n 1 -r BACKEND_CHOICE
    echo ""
    if [ "$BACKEND_CHOICE" = "2" ]; then
        if [ "$SMART_DEFAULT" = "llamacpp" ]; then
            BACKEND="ollama"
            ok "  Backend: Ollama"
        else
            BACKEND="llamacpp"
            ok "  Backend: llama.cpp"
        fi
    else
        BACKEND="$SMART_DEFAULT"
        if [ "$SMART_DEFAULT" = "llamacpp" ]; then
            ok "  Backend: llama.cpp (recommended for NVIDIA GPU)"
        else
            ok "  Backend: Ollama"
        fi
    fi
fi

# ============================================================================
# Step 1: Ollama (always installed — primary backend or embeddings-only)
# ============================================================================
if [ "$BACKEND" = "llamacpp" ]; then
    step "Checking Ollama (for embeddings)"
    echo -e "  ${DIM}llama.cpp is the primary backend — Ollama provides embedding model only.${RESET}"
else
    step "Checking Ollama"
fi

if command -v ollama &>/dev/null; then
    OLLAMA_VER=$(ollama --version 2>&1 | head -1)
    ok "  Ollama installed: $OLLAMA_VER"
else
    echo "  Ollama is not installed."
    if ask_yn "  Install Ollama via official installer? [Y/n] " "Y"; then
        echo "  Installing Ollama..."
        curl -fsSL https://ollama.com/install.sh | sh
        ok "  Ollama installed."
    else
        if [ "$BACKEND" = "ollama" ]; then
            echo ""
            echo "  Install Ollama manually:"
            echo "    curl -fsSL https://ollama.com/install.sh | sh"
            echo "    — or —"
            echo "    Download from https://ollama.com/download"
            exit 1
        else
            warn "  Skipping Ollama — embeddings will use keyword search fallback."
        fi
    fi
fi

# Ensure Ollama is running (if installed)
PORT="${PRE_PORT:-11434}"
if command -v ollama &>/dev/null; then
    if ! curl -sf "http://localhost:${PORT}/v1/models" >/dev/null 2>&1; then
        echo "  Starting Ollama..."
        if systemctl is-active ollama &>/dev/null; then
            ok "  Ollama service already active."
        elif systemctl start ollama &>/dev/null 2>&1; then
            sleep 3
            ok "  Started Ollama systemd service."
        else
            ollama serve >/dev/null 2>&1 &
            sleep 3
        fi
        if ! curl -sf "http://localhost:${PORT}/v1/models" >/dev/null 2>&1; then
            if [ "$BACKEND" = "ollama" ]; then
                warn "  Could not start Ollama automatically."
                echo "  Please run 'ollama serve' in another terminal, then re-run this script."
                exit 1
            else
                warn "  Could not start Ollama — embeddings may be limited."
            fi
        fi
    fi
    ok "  Ollama running on port $PORT."
fi

# ============================================================================
# Step 1b: Ollama environment (Ollama-primary only)
# ============================================================================
if [ "$BACKEND" = "ollama" ]; then
step "Configuring Ollama environment"

# Flash Attention: ON for CUDA (NVIDIA benefits), OFF if no GPU / ROCm (safe default)
FA_VAL=0
if [ "$GPU_BACKEND" = "cuda" ]; then
    FA_VAL=1
fi

KV_CACHE_TYPE="q8_0"
if [ "$QUANT" = "q4_K_M" ]; then
    KV_CACHE_TYPE="q4_0"
fi

# Match OLLAMA_NUM_PARALLEL to inference-queue slot count
NUM_PARALLEL=1
if [ "$HEADROOM" -ge 16 ] 2>/dev/null; then
    NUM_PARALLEL=2
fi

# Write Ollama env vars to systemd override (if Ollama runs as systemd service)
OLLAMA_SERVICE_DIR="/etc/systemd/system/ollama.service.d"
if systemctl is-enabled ollama &>/dev/null 2>&1; then
    if [ -d "/etc/systemd/system" ]; then
        sudo mkdir -p "$OLLAMA_SERVICE_DIR"
        cat <<ENVEOF | sudo tee "$OLLAMA_SERVICE_DIR/pre-environment.conf" >/dev/null
[Service]
Environment="OLLAMA_FLASH_ATTENTION=${FA_VAL}"
Environment="OLLAMA_KV_CACHE_TYPE=${KV_CACHE_TYPE}"
Environment="OLLAMA_KEEP_ALIVE=24h"
Environment="OLLAMA_NUM_PARALLEL=${NUM_PARALLEL}"
Environment="OLLAMA_MAX_LOADED_MODELS=1"
ENVEOF
        sudo systemctl daemon-reload
        sudo systemctl restart ollama
        sleep 3
        ok "  Ollama systemd environment configured (FA=${FA_VAL}, KV=${KV_CACHE_TYPE}, PARALLEL=${NUM_PARALLEL})"
    fi
fi

# Also persist in shell profile
SHELL_RC=""
if [ -f "$HOME/.bashrc" ]; then
    SHELL_RC="$HOME/.bashrc"
elif [ -f "$HOME/.zshrc" ]; then
    SHELL_RC="$HOME/.zshrc"
elif [ -f "$HOME/.profile" ]; then
    SHELL_RC="$HOME/.profile"
fi

OLLAMA_ENVS=(
    "OLLAMA_FLASH_ATTENTION:${FA_VAL}"
    "OLLAMA_KV_CACHE_TYPE:${KV_CACHE_TYPE}"
    "OLLAMA_KEEP_ALIVE:24h"
    "OLLAMA_NUM_PARALLEL:${NUM_PARALLEL}"
    "OLLAMA_MAX_LOADED_MODELS:1"
)

if [ -n "$SHELL_RC" ]; then
    CHANGED=false
    for entry in "${OLLAMA_ENVS[@]}"; do
        KEY="${entry%%:*}"
        VAL="${entry#*:}"
        if ! grep -q "^export ${KEY}=" "$SHELL_RC" 2>/dev/null; then
            echo "export ${KEY}=${VAL}" >> "$SHELL_RC"
            CHANGED=true
        fi
    done
    if [ "$CHANGED" = true ]; then
        ok "  Ollama env vars added to $SHELL_RC"
    else
        ok "  Ollama env vars already in $SHELL_RC"
    fi
fi
fi  # end BACKEND=ollama check for Step 1b

# ============================================================================
# Step 1c: llama.cpp server installation (when llamacpp backend selected)
# ============================================================================
if [ "$BACKEND" = "llamacpp" ]; then
    step "Installing llama.cpp"

    if command -v llama-server &>/dev/null; then
        LLAMA_VER=$(llama-server --version 2>&1 | head -1 || echo "installed")
        ok "  llama-server already installed: $LLAMA_VER"
    else
        echo "  Installing llama.cpp..."
        # Try pre-built binary from GitHub releases
        LLAMA_ARCH=$(uname -m)
        LLAMA_TAG=$(curl -sf "https://api.github.com/repos/ggml-org/llama.cpp/releases/latest" 2>/dev/null | grep '"tag_name"' | sed 's/.*"tag_name":[[:space:]]*"\([^"]*\)".*/\1/' || echo "")
        INSTALLED_LLAMA=false

        if [ -n "$LLAMA_TAG" ] && [ "$GPU_BACKEND" = "cuda" ]; then
            # Download CUDA pre-built binary
            LLAMA_URL="https://github.com/ggml-org/llama.cpp/releases/download/${LLAMA_TAG}/llama-${LLAMA_TAG}-bin-ubuntu-x64-cuda.tar.gz"
            echo "  Downloading llama.cpp ${LLAMA_TAG} (CUDA)..."
            if curl -fsSL "$LLAMA_URL" -o /tmp/llama-cpp.tar.gz 2>/dev/null; then
                mkdir -p /tmp/llama-cpp-extract
                tar xzf /tmp/llama-cpp.tar.gz -C /tmp/llama-cpp-extract 2>/dev/null || true
                LLAMA_BIN=$(find /tmp/llama-cpp-extract -name 'llama-server' -type f 2>/dev/null | head -1)
                if [ -n "$LLAMA_BIN" ]; then
                    sudo cp "$LLAMA_BIN" /usr/local/bin/llama-server
                    sudo chmod +x /usr/local/bin/llama-server
                    INSTALLED_LLAMA=true
                    ok "  llama-server installed to /usr/local/bin/"
                fi
                rm -rf /tmp/llama-cpp-extract /tmp/llama-cpp.tar.gz
            fi
        fi

        if [ "$INSTALLED_LLAMA" = false ]; then
            # Fallback: build from source
            echo "  Building llama.cpp from source (requires cmake, g++)..."
            if ! command -v cmake &>/dev/null; then
                if [ -n "$PKG_INSTALL" ]; then
                    $PKG_INSTALL cmake build-essential 2>&1 | tail -3
                else
                    fail "  cmake required to build llama.cpp. Install cmake and re-run."
                fi
            fi
            LLAMA_SRC="/tmp/llama-cpp-src"
            rm -rf "$LLAMA_SRC"
            git clone --depth 1 https://github.com/ggml-org/llama.cpp "$LLAMA_SRC" 2>&1 | tail -3

            CMAKE_ARGS="-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON"
            if [ "$GPU_BACKEND" = "cuda" ]; then
                CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
            fi

            cmake -B "$LLAMA_SRC/build" -S "$LLAMA_SRC" $CMAKE_ARGS 2>&1 | tail -5
            cmake --build "$LLAMA_SRC/build" --target llama-server -j$(nproc) 2>&1 | tail -5

            if [ -f "$LLAMA_SRC/build/bin/llama-server" ]; then
                sudo cp "$LLAMA_SRC/build/bin/llama-server" /usr/local/bin/
                sudo chmod +x /usr/local/bin/llama-server
                ok "  llama-server built and installed"
            else
                fail "  Failed to build llama-server"
            fi
            rm -rf "$LLAMA_SRC"
        fi
    fi

    LLAMACPP_PORT=8080
fi

# ============================================================================
# Step 2: Pull base model (via Ollama — both backends need the GGUF)
# ============================================================================
if command -v ollama &>/dev/null; then
step "Pulling base model ($BASE_MODEL)"

MODEL_TAG=$(echo "$BASE_MODEL" | sed 's/.*://')
if ollama list 2>/dev/null | grep -q "gemma4.*${MODEL_TAG}"; then
    ok "  Base model already available."
else
    echo "  Downloading ~${MODEL_SIZE_GB}GB model ($QUANT). This may take a while..."
    echo "  (Download is resumable — safe to interrupt and re-run)"
    ollama pull "$BASE_MODEL" || fail "Failed to pull $BASE_MODEL"
    ok "  Base model downloaded."
fi

# ============================================================================
# Step 2b: Pull embedding model
# ============================================================================
step "Pulling embedding model (nomic-embed-text — experience ledger + RAG)"

if ollama list 2>/dev/null | grep -q "nomic-embed-text"; then
    ok "  Embedding model already available."
else
    echo "  Downloading nomic-embed-text (~274MB)..."
    ollama pull nomic-embed-text || warn "  Failed to pull nomic-embed-text — RAG will use keyword fallback."
    if ollama list 2>/dev/null | grep -q "nomic-embed-text"; then
        ok "  Embedding model downloaded."
    fi
fi
fi  # end Ollama model pull

# ============================================================================
# Step 3: Create custom model from Modelfile (Ollama backend only)
# ============================================================================
if [ "$BACKEND" = "ollama" ] && command -v ollama &>/dev/null; then
step "Creating optimized model ($CUSTOM_MODEL)"

EFFECTIVE_MODELFILE="$ENGINE_DIR/Modelfile"
if [ ! -f "$EFFECTIVE_MODELFILE" ]; then
    fail "Modelfile not found at $EFFECTIVE_MODELFILE"
fi
if [ "$QUANT" != "q8_0" ]; then
    TEMP_MODELFILE="/tmp/pre-Modelfile-${QUANT}"
    sed "s|FROM gemma4:26b-a4b-it-q8_0|FROM $BASE_MODEL|" "$EFFECTIVE_MODELFILE" > "$TEMP_MODELFILE"
    EFFECTIVE_MODELFILE="$TEMP_MODELFILE"
    ok "  Using modified Modelfile with $QUANT base"
fi

if ollama list 2>/dev/null | grep -q "$CUSTOM_MODEL"; then
    ok "  Custom model already exists."
    if ask_yn "  Recreate from Modelfile? [y/N] " "N"; then
        ollama create "$CUSTOM_MODEL" -f "$EFFECTIVE_MODELFILE" || fail "Failed to create $CUSTOM_MODEL"
        ok "  Custom model recreated."
    fi
else
    ollama create "$CUSTOM_MODEL" -f "$EFFECTIVE_MODELFILE" || fail "Failed to create $CUSTOM_MODEL"
    ok "  Custom model created ($QUANT, dynamic context, optimized batch size)."
fi

# Fast tier model — Gemma 4 E2B for quick Q&A (~140 tok/s on Apple Silicon)
FAST_MODEL="pre-gemma4-fast"
FAST_MODELFILE="$ENGINE_DIR/Modelfile.fast"
if [ -f "$FAST_MODELFILE" ]; then
    if ! ollama list 2>/dev/null | grep -q "$FAST_MODEL"; then
        echo -e "  ${DIM}Pulling fast tier base model (gemma4:e2b, ~6GB)...${RESET}"
        ollama pull gemma4:e2b 2>/dev/null || warn "  Failed to pull gemma4:e2b — fast tier will be unavailable"
        if ollama list 2>/dev/null | grep -q "gemma4:e2b"; then
            ollama create "$FAST_MODEL" -f "$FAST_MODELFILE" 2>/dev/null && \
                ok "  Fast tier model created ($FAST_MODEL — Gemma 4 E2B)" || \
                warn "  Failed to create $FAST_MODEL — fast tier will fall back to main model"
        fi
    else
        ok "  Fast tier model already exists ($FAST_MODEL)"
    fi
fi

# ----------------------------------------------------------------------------
# Optional: Joby IT Service Desk model (private, password-gated)
# ----------------------------------------------------------------------------
# A LoRA fine-tune of pre-gemma4 specialized for Joby Aviation's IT service
# desk. Distributed via HuggingFace private repo, gated here by a simple
# install-time password. Internal Joby use only.
#
# Note: the password is an intent-gate, not a crypto-gate. Actual access
# control is the HF private-repo token. Anyone with this script can read
# the password hash; the hash just keeps casual installers from grabbing it.
JOBY_MODEL_REPO="sunkencity/joby-it-servicedesk"
JOBY_MODEL_TAG="q8_0"
JOBY_PASSWORD_HASH="18278044341fbdb26d8b0efae30087a1d38c1283c4af7448abf638580dd269cb"

if ask_yn "  Install Joby IT service desk model? (Joby internal only) [y/N] " "N"; then
    read -s -r -p "    Access password: " joby_input_password
    echo
    joby_input_hash=$(echo -n "$joby_input_password" | sha256sum | awk '{print $1}')
    unset joby_input_password
    if [ "$joby_input_hash" != "$JOBY_PASSWORD_HASH" ]; then
        warn "  Incorrect password — skipping Joby model."
    else
        ok "  Password accepted."

        # We download the GGUF locally then `ollama create`, instead of
        # `ollama pull hf.co/...`, because Ollama's HF integration trips on a
        # redirect realm-mismatch ("hf.co" vs "huggingface.co") when the repo
        # is private. The `hf` CLI handles auth correctly.

        # ── Check prerequisites: hf CLI + HuggingFace token ──
        _joby_ready=false
        while true; do
            _missing=""
            if ! command -v hf >/dev/null 2>&1; then
                _missing="${_missing}  • HuggingFace CLI ('hf') is not installed.\n"
            fi
            # Check for HF token (env var or cached file)
            _hf_token="${HF_TOKEN:-}"
            if [ -z "$_hf_token" ] && [ -f "$HOME/.cache/huggingface/token" ]; then
                _hf_token="$(cat "$HOME/.cache/huggingface/token")"
            fi
            if [ -z "$_hf_token" ]; then
                _missing="${_missing}  • HuggingFace auth token not found.\n"
            fi

            if [ -z "$_missing" ]; then
                export HF_TOKEN="$_hf_token"
                _joby_ready=true
                break
            fi

            echo ""
            warn "  Missing prerequisites:"
            printf "$_missing"
            echo ""
            echo "  To fix, open another terminal and run:"
            echo ""
            if ! command -v hf >/dev/null 2>&1; then
                echo "    pip install --upgrade huggingface_hub    # install the hf CLI"
            fi
            if [ -z "$_hf_token" ]; then
                echo "    hf auth login                            # authenticate with HuggingFace"
                echo "    (paste your token from https://huggingface.co/settings/tokens)"
            fi
            echo ""
            if ask_yn "  Press Y when ready to re-check, or N to skip the Joby model [Y/n] " "Y"; then
                echo "  Re-checking prerequisites..."
                # Re-source PATH in case user installed hf in this shell session
                export PATH="$HOME/.local/bin:$HOME/bin:$PATH"
                continue
            else
                warn "  Skipping Joby model."
                break
            fi
        done

        if [ "$_joby_ready" = true ]; then
            JOBY_GGUF_NAME="joby-it-servicedesk.${JOBY_MODEL_TAG}.gguf"
            JOBY_WORKDIR="$(mktemp -d -t pre-joby-XXXXXX)"
            trap 'rm -rf "$JOBY_WORKDIR"' EXIT

            echo "    Downloading ${JOBY_GGUF_NAME} (~30 GB)..."
            if hf download "$JOBY_MODEL_REPO" "$JOBY_GGUF_NAME" --local-dir "$JOBY_WORKDIR"; then
                cat > "$JOBY_WORKDIR/Modelfile.joby" <<MODELFILE
FROM $JOBY_WORKDIR/$JOBY_GGUF_NAME
PARAMETER num_ctx 8192
PARAMETER num_batch 512
PARAMETER temperature 1.0
PARAMETER top_k 64
PARAMETER top_p 0.95
PARAMETER min_p 0.05
PARAMETER repeat_penalty 1.1
PARAMETER repeat_last_n 256
MODELFILE
                if ollama create pre-gemma4-itsd -f "$JOBY_WORKDIR/Modelfile.joby"; then
                    ok "  Joby IT model installed as: pre-gemma4-itsd"

                    # Flip the Joby service-desk mode flag in router.json so
                    # PRE's system prompt auto-injects the anti-hallucination
                    # guard (see web/src/context.js → isJobyServiceDeskMode).
                    mkdir -p "$HOME/.pre"
                    ROUTER_JSON="$HOME/.pre/router.json"
                    if [ -f "$ROUTER_JSON" ]; then
                        python3 -c "
import json, sys
p = '$ROUTER_JSON'
try:
    with open(p) as f: cfg = json.load(f)
except Exception:
    cfg = {}
cfg['joby_servicedesk_mode'] = True
with open(p, 'w') as f: json.dump(cfg, f, indent=2)
" 2>/dev/null && ok "  Set joby_servicedesk_mode=true in router.json"
                    else
                        echo '{"joby_servicedesk_mode": true}' > "$ROUTER_JSON"
                        ok "  Wrote router.json with joby_servicedesk_mode=true"
                    fi
                else
                    warn "  ollama create failed. The GGUF is at $JOBY_WORKDIR — retry manually."
                    trap - EXIT
                fi
            else
                warn "  hf download failed. Check HF access to $JOBY_MODEL_REPO."
            fi
        fi
    fi
fi

fi

# ============================================================================
# Step 3.5: Optional escalation model (eGPU offload via Ollama)
# ============================================================================
# On Linux/Windows, Ollama already drives the GPU directly — escalation just
# registers a second, larger Ollama model alongside pre-gemma4 that the user
# can route to per-message via /escalate or the ↑GPU toggle in the chat.

if [ "$GPU_VRAM_GB" -ge 16 ]; then
    step "Optional: escalation model (eGPU offload)"

    ESC_REPO=""; ESC_FILE=""; ESC_LABEL=""; ESC_SIZE_GB=0; ESC_GATED=false
    if [ "$GPU_VRAM_GB" -ge 80 ]; then
        ESC_REPO="bartowski/Llama-3.3-70B-Instruct-GGUF"
        ESC_FILE="Llama-3.3-70B-Instruct-Q6_K.gguf"
        ESC_LABEL="Llama-3.3-70B Q6_K"
        ESC_SIZE_GB=57; ESC_GATED=true
    elif [ "$GPU_VRAM_GB" -ge 48 ]; then
        ESC_REPO="bartowski/Llama-3.3-70B-Instruct-GGUF"
        ESC_FILE="Llama-3.3-70B-Instruct-Q4_K_M.gguf"
        ESC_LABEL="Llama-3.3-70B Q4_K_M"
        ESC_SIZE_GB=42; ESC_GATED=true
    elif [ "$GPU_VRAM_GB" -ge 24 ]; then
        ESC_REPO="bartowski/THUDM_GLM-4-32B-0414-GGUF"
        ESC_FILE="THUDM_GLM-4-32B-0414-Q4_K_M.gguf"
        ESC_LABEL="GLM-4-32B Q4_K_M"
        ESC_SIZE_GB=20; ESC_GATED=false
    else
        ESC_REPO="bartowski/DeepSeek-Coder-V2-Lite-Instruct-GGUF"
        ESC_FILE="DeepSeek-Coder-V2-Lite-Instruct-Q4_K_M.gguf"
        ESC_LABEL="DeepSeek-Coder-V2-Lite Q4_K_M"
        ESC_SIZE_GB=10; ESC_GATED=false
    fi

    ESC_OLLAMA_NAME="pre-escalate"

    echo "  Detected ${GPU_VRAM_GB}GB GPU VRAM"
    echo "  Recommended escalation model: $ESC_LABEL (~${ESC_SIZE_GB}GB)"
    echo "  Runtime: Ollama on the GPU as '$ESC_OLLAMA_NAME'. Triggered by /escalate or the ↑GPU toggle."
    if [ "$ESC_GATED" = true ]; then
        echo "  Note: this model requires accepting Meta's Llama license on HuggingFace."
    fi

    if ask_yn "  Install escalation model? [Y/n] " "Y"; then
        _esc_ready=false
        while true; do
            _missing=""
            if ! command -v hf >/dev/null 2>&1; then
                _missing="${_missing}  • HuggingFace CLI ('hf') is not installed.\n"
            fi
            _hf_token="${HF_TOKEN:-}"
            if [ -z "$_hf_token" ] && [ -f "$HOME/.cache/huggingface/token" ]; then
                _hf_token="$(cat "$HOME/.cache/huggingface/token")"
            fi
            if [ -z "$_hf_token" ] && [ "$ESC_GATED" = true ]; then
                _missing="${_missing}  • HuggingFace auth token not found (required for gated model).\n"
            fi
            if [ -z "$_missing" ]; then
                [ -n "$_hf_token" ] && export HF_TOKEN="$_hf_token"
                _esc_ready=true
                break
            fi
            echo ""
            warn "  Missing prerequisites:"
            printf "$_missing"
            echo ""
            echo "  To fix, open another terminal and run:"
            if ! command -v hf >/dev/null 2>&1; then
                echo "    pip install --upgrade huggingface_hub"
            fi
            if [ -z "$_hf_token" ] && [ "$ESC_GATED" = true ]; then
                echo "    hf auth login    # paste token from https://huggingface.co/settings/tokens"
            fi
            echo ""
            if ask_yn "  Press Y to re-check, or N to skip the escalation model [Y/n] " "Y"; then
                export PATH="$HOME/.local/bin:$HOME/bin:$PATH"
                continue
            else
                warn "  Skipping escalation model."
                break
            fi
        done

        if [ "$_esc_ready" = true ]; then
            ESC_DIR="$HOME/.pre/models/escalation"
            mkdir -p "$ESC_DIR"

            _esc_dl_log="$(mktemp)"
            if ! hf download "$ESC_REPO" "$ESC_FILE" --local-dir "$ESC_DIR" 2>&1 | tee "$_esc_dl_log"; then
                if grep -qiE "403|gated|access denied|restricted" "$_esc_dl_log"; then
                    warn "  This model is gated — you need to accept the license on HuggingFace."
                    echo "    GGUF mirror: https://huggingface.co/$ESC_REPO"
                    # xdg-open works on most Linux desktops
                    xdg-open "https://huggingface.co/$ESC_REPO" 2>/dev/null || true
                    read -r -p "  Press Enter once you've accepted the license to retry, or Ctrl+C to skip: "
                    hf download "$ESC_REPO" "$ESC_FILE" --local-dir "$ESC_DIR" || \
                        warn "  Download still failed — escalation model not installed."
                fi
            fi
            rm -f "$_esc_dl_log"

            ESC_GGUF_PATH="$ESC_DIR/$ESC_FILE"
            if [ -f "$ESC_GGUF_PATH" ]; then
                # Register with Ollama
                ESC_MODELFILE="$(mktemp -t pre-escalate-modelfile.XXXXXX)"
                cat > "$ESC_MODELFILE" <<MODELFILE
FROM $ESC_GGUF_PATH
PARAMETER num_ctx 8192
PARAMETER temperature 0.7
PARAMETER top_p 0.9
PARAMETER repeat_penalty 1.1
MODELFILE
                if ollama create "$ESC_OLLAMA_NAME" -f "$ESC_MODELFILE"; then
                    ok "  Ollama escalation model created: $ESC_OLLAMA_NAME"

                    # Write _provider.escalation to connections.json (ollama runtime)
                    mkdir -p "$HOME/.pre"
                    CONN_JSON="$HOME/.pre/connections.json"
                    python3 - "$CONN_JSON" "$ESC_OLLAMA_NAME" "$ESC_LABEL" <<'PYEOF'
import json, os, sys
p, name, label = sys.argv[1:4]
data = {}
if os.path.exists(p):
    try:
        with open(p) as f: data = json.load(f)
    except Exception: data = {}
prov = data.setdefault("_provider", {})
prov["escalation"] = {
    "enabled": True,
    "runtime": "ollama",
    "model_name": name,
    "label": label,
}
os.makedirs(os.path.dirname(p), exist_ok=True)
with open(p, "w") as f: json.dump(data, f, indent=2)
print("OK")
PYEOF
                    ok "  Escalation model registered: $ESC_LABEL"
                else
                    warn "  ollama create failed. The GGUF is at $ESC_GGUF_PATH — retry manually."
                fi
                rm -f "$ESC_MODELFILE"
            else
                warn "  Escalation GGUF not found at $ESC_GGUF_PATH — skipping registration."
            fi
        fi
    fi
fi

# ============================================================================
# Step 4: Node.js
# ============================================================================
step "Checking Node.js"

if command -v node &>/dev/null; then
    NODE_VER=$(node --version)
    NODE_MAJOR=$(echo "$NODE_VER" | sed 's/v//' | cut -d. -f1)
    if [ "$NODE_MAJOR" -ge 18 ]; then
        ok "  Node.js: $NODE_VER"
    else
        warn "  Node.js $NODE_VER is too old — need v18+."
        echo "  Install a newer version via:"
        echo "    nvm install 22"
        echo "    — or —"
        echo "    https://nodejs.org/en/download"
        if [ "$PKG_MGR" = "apt" ]; then
            echo "    — or —"
            echo "    curl -fsSL https://deb.nodesource.com/setup_22.x | sudo -E bash - && sudo apt-get install -y nodejs"
        fi
        exit 1
    fi
else
    warn "  Node.js not found."
    echo "  Install Node.js (v18+ required):"
    if [ "$PKG_MGR" = "apt" ]; then
        echo "    curl -fsSL https://deb.nodesource.com/setup_22.x | sudo -E bash -"
        echo "    sudo apt-get install -y nodejs"
    elif [ "$PKG_MGR" = "dnf" ]; then
        echo "    sudo dnf install nodejs"
    elif [ "$PKG_MGR" = "pacman" ]; then
        echo "    sudo pacman -S nodejs npm"
    fi
    echo "    — or —"
    echo "    curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash"
    echo "    nvm install 22"
    exit 1
fi

# ============================================================================
# Step 5: Web GUI dependencies
# ============================================================================
step "Setting up Web GUI"

WEB_DIR="$REPO_DIR/web"
if [ -f "$WEB_DIR/package.json" ]; then
    echo "  Installing web dependencies..."
    cd "$WEB_DIR"
    npm install --silent 2>&1 | tail -3
    ok "  Web GUI dependencies installed"
    cd "$REPO_DIR"
else
    warn "  Web GUI not found at $WEB_DIR — skipping."
fi

# ============================================================================
# Step 6: Create ~/.pre/ directories
# ============================================================================
step "Setting up data directories"

PRE_DATA="$HOME/.pre"
mkdir -p "$PRE_DATA"/{sessions,memory/experience,artifacts,rag,workflows,custom_tools,training/exports,agent-checkpoints}
ok "  ~/.pre/ directory structure created"

# Write context window
echo "$OPTIMAL_CTX" > "$PRE_DATA/context"
ok "  Context window set to ${CTX_HUMAN} (${OPTIMAL_CTX} tokens)"

# Write provider configuration for llama.cpp backend
if [ "$BACKEND" = "llamacpp" ] && command -v node &>/dev/null; then
    CONN_FILE="$PRE_DATA/connections.json"
    node -e "
const fs = require('fs');
const p = '$CONN_FILE';
let cfg = {};
try { cfg = JSON.parse(fs.readFileSync(p, 'utf-8')); } catch {}
if (!cfg._provider) cfg._provider = {};
cfg._provider.active = 'llamacpp';
cfg._provider.llamacpp = {
  base_url: 'http://127.0.0.1:${LLAMACPP_PORT:-8080}',
  port: ${LLAMACPP_PORT:-8080},
  model_path: '',
  draft_model_path: '',
  grammar_enabled: true,
  flash_attn: true,
  n_gpu_layers: -1,
  context_size: $OPTIMAL_CTX,
  parallel: 1
};
fs.writeFileSync(p, JSON.stringify(cfg, null, 2) + '\n', { mode: 0o600 });
"
    ok "  Provider configured: llama.cpp → ~/.pre/connections.json"
fi

# ============================================================================
# Step 7: Optional — voice tools
# ============================================================================
step "Optional: Voice Tools"

if ask_yn "  Install voice tools (Whisper STT + espeak-ng TTS + FFmpeg)? [y/N] " "N"; then
    if [ -n "$PKG_INSTALL" ]; then
        echo "  Installing FFmpeg and espeak-ng..."
        $PKG_UPDATE 2>/dev/null || true
        $PKG_INSTALL ffmpeg espeak-ng 2>&1 | tail -3
        ok "  FFmpeg and espeak-ng installed"
    else
        warn "  Unknown package manager — install ffmpeg and espeak-ng manually."
    fi
    echo ""
    echo "  For Whisper STT, install via pip:"
    echo "    pip install openai-whisper"
    echo "  (requires Python 3.9+ and ~1GB disk for the base model)"
else
    ok "  Skipping voice tools."
fi

# ============================================================================
# Step 8: Optional — GNOME PIM integration
# ============================================================================
step "Optional: GNOME Calendar/Contacts/Reminders"

# Check if running a GNOME desktop
GNOME_DETECTED=false
if [ -n "$GNOME_DESKTOP_SESSION_ID" ] || echo "$XDG_CURRENT_DESKTOP" | grep -qi "gnome"; then
    GNOME_DETECTED=true
fi

if [ "$GNOME_DETECTED" = true ]; then
    if command -v gdbus &>/dev/null; then
        ok "  GNOME desktop detected with gdbus — Calendar/Contacts/Reminders will work via Evolution Data Server."
    else
        warn "  GNOME detected but gdbus not found."
        if [ -n "$PKG_INSTALL" ]; then
            if ask_yn "  Install evolution-data-server for calendar/contacts/reminders? [Y/n] " "Y"; then
                $PKG_INSTALL evolution-data-server 2>&1 | tail -3
                ok "  evolution-data-server installed"
            fi
        fi
    fi
else
    ok "  Non-GNOME desktop detected. Calendar/Contacts/Reminders require GNOME Evolution Data Server."
    if [ -n "$PKG_INSTALL" ]; then
        if ask_yn "  Install evolution-data-server anyway? (works if GNOME libs are available) [y/N] " "N"; then
            $PKG_INSTALL evolution-data-server 2>&1 | tail -3
            ok "  evolution-data-server installed"
        else
            ok "  Skipping — PRE will show a helpful message if you try to use these tools."
        fi
    fi
fi

# ============================================================================
# Step 9: Optional — desktop automation deps
# ============================================================================
step "Optional: Desktop Automation (Computer Use)"

if ask_yn "  Install desktop automation tools (xdotool, scrot, xclip, wmctrl)? [y/N] " "N"; then
    if [ -n "$PKG_INSTALL" ]; then
        $PKG_UPDATE 2>/dev/null || true
        if [ "$PKG_MGR" = "apt" ]; then
            $PKG_INSTALL xdotool scrot xclip wmctrl libnotify-bin 2>&1 | tail -3
        elif [ "$PKG_MGR" = "dnf" ]; then
            $PKG_INSTALL xdotool scrot xclip wmctrl libnotify 2>&1 | tail -3
        elif [ "$PKG_MGR" = "pacman" ]; then
            $PKG_INSTALL xdotool scrot xclip wmctrl libnotify 2>&1 | tail -3
        fi
        ok "  Desktop automation tools installed (X11)"
    else
        warn "  Unknown package manager. Install manually: xdotool scrot xclip wmctrl"
    fi
else
    ok "  Skipping desktop automation tools."
fi

# ============================================================================
# Step 10: Optional — auto-start (systemd user service)
# ============================================================================
step "Optional: Auto-start at Login"

if ask_yn "  Set up PRE web server to start automatically at login? [y/N] " "N"; then
    SYSTEMD_USER_DIR="$HOME/.config/systemd/user"
    mkdir -p "$SYSTEMD_USER_DIR"

    cat > "$SYSTEMD_USER_DIR/pre-server.service" <<SVCEOF
[Unit]
Description=PRE Personal Reasoning Engine — Web GUI
After=network.target

[Service]
Type=simple
WorkingDirectory=$WEB_DIR
ExecStart=$(command -v node) $WEB_DIR/server.js
Environment=PRE_WEB_PORT=7749
Restart=on-failure
RestartSec=5

[Install]
WantedBy=default.target
SVCEOF

    systemctl --user daemon-reload
    systemctl --user enable pre-server
    ok "  systemd user service created and enabled"
    ok "  PRE will auto-start at login on port 7749"
    echo -e "  ${DIM}Manage: systemctl --user start|stop|status pre-server${RESET}"
else
    ok "  Skipping auto-start."
fi

# ============================================================================
# Step 10b: tinygrad setup (optional — NVIDIA GPU inference + training)
# ============================================================================
if [ "$GPU_BACKEND" = "cuda" ]; then
    step "tinygrad Setup (GPU inference + LoRA training)"

    echo -e "  ${BOLD}tinygrad${RESET} enables Llama model inference and LoRA fine-tuning on your NVIDIA GPU."
    echo -e "  PRE auto-detects your GPU at runtime — tinygrad activates when available."
    echo ""

    # Check for Python 3.11+
    TINYGRAD_PYTHON=""
    for PY_CMD in python3.13 python3.12 python3.11 python3; do
        if command -v "$PY_CMD" &>/dev/null; then
            PY_VER=$($PY_CMD --version 2>&1 | grep -oE '[0-9]+\.[0-9]+')
            PY_MAJOR=$(echo "$PY_VER" | cut -d. -f1)
            PY_MINOR=$(echo "$PY_VER" | cut -d. -f2)
            if [ "$PY_MAJOR" -ge 3 ] && [ "$PY_MINOR" -ge 11 ]; then
                TINYGRAD_PYTHON=$(command -v "$PY_CMD")
                break
            fi
        fi
    done

    if [ -z "$TINYGRAD_PYTHON" ]; then
        warn "  Python 3.11+ not found — tinygrad requires Python 3.11 or newer."
        echo -e "  ${DIM}Install via your package manager: sudo apt install python3.12${RESET}"
    else
        ok "  Python: $($TINYGRAD_PYTHON --version 2>&1) ($TINYGRAD_PYTHON)"

        if $TINYGRAD_PYTHON -c "import tinygrad" 2>/dev/null; then
            TINYGRAD_VER=$($TINYGRAD_PYTHON -c "import tinygrad; print(tinygrad.__version__)" 2>/dev/null || echo "unknown")
            ok "  tinygrad already installed (v${TINYGRAD_VER})"
        else
            echo "  Installing tinygrad..."
            $TINYGRAD_PYTHON -m pip install tinygrad --quiet 2>&1 | tail -3
            if $TINYGRAD_PYTHON -c "import tinygrad" 2>/dev/null; then
                TINYGRAD_VER=$($TINYGRAD_PYTHON -c "import tinygrad; print(tinygrad.__version__)" 2>/dev/null || echo "unknown")
                ok "  tinygrad installed (v${TINYGRAD_VER})"
            else
                warn "  tinygrad installation failed. Install manually: pip3 install tinygrad"
            fi
        fi

        # Verify GPU access
        TINYGRAD_DEV=$(DEV=NV $TINYGRAD_PYTHON -c "from tinygrad import Device; print(Device.DEFAULT)" 2>/dev/null || true)
        if echo "$TINYGRAD_DEV" | grep -qi "NV\|CUDA"; then
            ok "  GPU access verified via tinygrad (device: $TINYGRAD_DEV)"
            echo "{\"tinygrad\":true,\"python\":\"$TINYGRAD_PYTHON\",\"device\":\"$TINYGRAD_DEV\",\"verified\":\"$(date -u +%Y-%m-%dT%H:%M:%SZ)\"}" \
                > "$HOME/.pre/tinygrad-status.json"
        else
            warn "  tinygrad installed but GPU not accessible via DEV=NV"
            echo -e "  ${DIM}Ensure NVIDIA drivers and CUDA toolkit are installed: nvidia-smi${RESET}"
        fi
    fi
fi

# ============================================================================
# Step 11: Pre-warm model (Ollama backend only)
# ============================================================================
if [ "$BACKEND" = "ollama" ] && command -v ollama &>/dev/null; then
step "Pre-warming model into GPU memory"

if ollama ps 2>/dev/null | grep -q "$CUSTOM_MODEL"; then
    ok "  Model already loaded."
else
    echo "  Loading model into GPU (context: ${OPTIMAL_CTX} tokens)..."
    curl -sf --max-time 300 "http://localhost:${PORT}/api/chat" \
        -d "{\"model\":\"${CUSTOM_MODEL}\",\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"keep_alive\":\"24h\",\"options\":{\"num_predict\":1,\"num_ctx\":${OPTIMAL_CTX}}}" \
        >/dev/null 2>&1 || true
    if ollama ps 2>/dev/null | grep -q "$CUSTOM_MODEL"; then
        ok "  Model loaded and warm."
    else
        warn "  Model may not have fully loaded. First request may be slow."
    fi
fi
fi

# ============================================================================
# Step 12: Create launcher script
# ============================================================================
step "Setting up launcher"

LAUNCHER_DIR="$HOME/.local/bin"
mkdir -p "$LAUNCHER_DIR"
LAUNCHER="$LAUNCHER_DIR/pre-web"

cat > "$LAUNCHER" <<LAUNCHEOF
#!/bin/bash
# pre-web — Launch PRE Web GUI
cd "$WEB_DIR" && exec node server.js
LAUNCHEOF
chmod +x "$LAUNCHER"
ok "  Launcher created at $LAUNCHER"

# Add ~/.local/bin to PATH if not already there
if ! echo "$PATH" | grep -q "$HOME/.local/bin"; then
    if [ -n "$SHELL_RC" ]; then
        if ! grep -q '\.local/bin' "$SHELL_RC" 2>/dev/null; then
            echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$SHELL_RC"
            ok "  Added ~/.local/bin to PATH in $SHELL_RC"
        fi
    fi
fi

# ============================================================================
# Step 13: Create XDG desktop entry
# ============================================================================
DESKTOP_DIR="$HOME/.local/share/applications"
mkdir -p "$DESKTOP_DIR"

# Check for a logo
ICON_PATH="$WEB_DIR/public/img/logo.png"
if [ ! -f "$ICON_PATH" ]; then
    ICON_PATH=""
fi

cat > "$DESKTOP_DIR/pre.desktop" <<DESKTOPEOF
[Desktop Entry]
Type=Application
Name=PRE - Personal Reasoning Engine
Comment=Local AI assistant with 70+ tools
Exec=bash -c 'cd $WEB_DIR && node server.js & sleep 2 && xdg-open http://localhost:7749'
${ICON_PATH:+Icon=$ICON_PATH}
Terminal=false
Categories=Development;AI;Science;
StartupNotify=true
DESKTOPEOF
chmod +x "$DESKTOP_DIR/pre.desktop"
ok "  Desktop entry created (find PRE in your app launcher)"

# ============================================================================
# Done!
# ============================================================================
echo ""
echo -e "${BOLD}${GREEN}═══════════════════════════════════════════════════════════${RESET}"
echo -e "${BOLD}${GREEN}  PRE is installed and ready!${RESET}"
echo -e "${BOLD}${GREEN}═══════════════════════════════════════════════════════════${RESET}"
echo ""
echo -e "  ${BOLD}Start PRE:${RESET}"
echo -e "    cd $WEB_DIR && node server.js"
echo -e "    Then open: ${CYAN}http://localhost:7749${RESET}"
echo ""
echo -e "  ${BOLD}Quick launch:${RESET}"
echo -e "    pre-web      ${DIM}(if ~/.local/bin is in PATH)${RESET}"
echo ""
echo -e "  ${BOLD}Hardware profile:${RESET}"
echo -e "    GPU: ${GPU_NAME:-CPU only}  |  VRAM: ${GPU_VRAM_GB}GB  |  RAM: ${RAM_GB}GB"
echo -e "    Model: ${QUANT}  |  Context: ${CTX_HUMAN}  |  GPU: ${GPU_BACKEND:-cpu}  |  Inference: ${BACKEND}"
echo ""
if [ -n "$GPU_BACKEND" ]; then
    echo -e "  ${DIM}Expect ~30-70+ tok/s with GPU acceleration.${RESET}"
else
    echo -e "  ${YELLOW}No GPU detected — expect ~3-6 tok/s (CPU mode).${RESET}"
    echo -e "  ${DIM}Install NVIDIA drivers + CUDA toolkit for GPU acceleration.${RESET}"
fi
echo ""
