#!/bin/bash
# update.sh - Update PRE (Personal Reasoning Engine) on macOS
#
# Detects whether PRE was installed via git clone or zip download,
# compares local version against the latest release on GitHub,
# and updates accordingly. User data (~/.pre/) is never touched.
#
# Usage:
#   bash update.sh          # Interactive
#   bash update.sh --yes    # Non-interactive (accept all defaults)

set -e

REPO_URL="https://github.com/sunkencity999/pre"
RAW_URL="https://raw.githubusercontent.com/sunkencity999/pre/main"
REPO_DIR="$(cd "$(dirname "$0")" && pwd)"
WEB_DIR="$REPO_DIR/web"
ENGINE_DIR="$REPO_DIR/engine"
PRE_PORT="${PRE_PORT:-7749}"
AUTO_YES=false

if [[ "$1" == "--yes" || "$1" == "-y" ]]; then
    AUTO_YES=true
fi

# -- Helpers ----------------------------------------------------------------

step()  { printf "\n\033[1;36m=== %s ===\033[0m\n\n" "$1"; }
ok()    { printf "  \033[32m%s\033[0m\n" "$1"; }
warn()  { printf "  \033[33m%s\033[0m\n" "$1"; }
fail()  { printf "\n  \033[31m%s\033[0m\n" "$1"; exit 1; }

ask_yn() {
    local prompt="$1" default="$2"
    if $AUTO_YES; 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
}

# -- Detect install type ----------------------------------------------------

step "Checking installation"

IS_GIT=false
if [ -d "$REPO_DIR/.git" ]; then
    IS_GIT=true
    ok "Install type: git repository"
    # Check for local modifications
    if [ -n "$(git -C "$REPO_DIR" status --porcelain 2>/dev/null)" ]; then
        warn "Local modifications detected."
        git -C "$REPO_DIR" status --short
        echo ""
    fi
else
    ok "Install type: zip download (no git repository)"
fi

# -- Get local version ------------------------------------------------------

LOCAL_VERSION="unknown"
VERSION_FILE="$REPO_DIR/VERSION"
if [ -f "$VERSION_FILE" ]; then
    LOCAL_VERSION="$(cat "$VERSION_FILE" | tr -d '[:space:]')"
fi
ok "Local version: $LOCAL_VERSION"

# -- Get remote version -----------------------------------------------------

step "Checking for updates"

REMOTE_VERSION=""
if command -v curl >/dev/null 2>&1; then
    REMOTE_VERSION="$(curl -sf "${RAW_URL}/VERSION" 2>/dev/null | tr -d '[:space:]')" || true
fi

if [ -z "$REMOTE_VERSION" ]; then
    warn "Could not fetch remote version. Check your internet connection."
    if ! ask_yn "  Continue anyway? [y/N]" "N"; then
        exit 0
    fi
    REMOTE_VERSION="unknown"
fi

ok "Remote version: $REMOTE_VERSION"

# Compare versions
if [ "$LOCAL_VERSION" = "$REMOTE_VERSION" ] && [ "$REMOTE_VERSION" != "unknown" ]; then
    ok "You're already on the latest version ($LOCAL_VERSION)."
    if ! ask_yn "  Force update anyway? [y/N]" "N"; then
        echo ""
        echo "  No update needed."
        exit 0
    fi
else
    if [ "$REMOTE_VERSION" != "unknown" ]; then
        ok "Update available: $LOCAL_VERSION -> $REMOTE_VERSION"
    fi
fi

# -- Check if server is running ---------------------------------------------

SERVER_WAS_RUNNING=false
if curl -sf "http://localhost:${PRE_PORT}/api/status" >/dev/null 2>&1; then
    SERVER_WAS_RUNNING=true
    warn "PRE server is currently running."
    if ask_yn "  Stop it for the update? [Y/n]" "Y"; then
        echo "  Stopping server..."
        # Try graceful shutdown via the API, fall back to kill
        curl -sf "http://localhost:${PRE_PORT}/api/shutdown" >/dev/null 2>&1 || true
        sleep 1
        # Kill by port if still running
        lsof -ti :${PRE_PORT} -sTCP:LISTEN 2>/dev/null | xargs kill 2>/dev/null || true
        sleep 1
        ok "Server stopped."
    else
        warn "Updating while server is running - restart manually after."
    fi
fi

# -- Perform update ---------------------------------------------------------

step "Updating PRE"

if $IS_GIT; then
    # ---- Git-based update ----
    echo "  Pulling latest changes from origin..."

    # Stash local changes if any
    STASHED=false
    if [ -n "$(git -C "$REPO_DIR" status --porcelain 2>/dev/null)" ]; then
        warn "Stashing local changes..."
        git -C "$REPO_DIR" stash push -m "pre-update-$(date +%Y%m%d-%H%M%S)" >/dev/null 2>&1
        STASHED=true
    fi

    # Pull
    PULL_OUTPUT="$(git -C "$REPO_DIR" pull origin main 2>&1)" || {
        echo "$PULL_OUTPUT"
        if $STASHED; then
            warn "Restoring stashed changes..."
            git -C "$REPO_DIR" stash pop >/dev/null 2>&1 || true
        fi
        fail "Git pull failed. Resolve conflicts manually."
    }
    echo "  $PULL_OUTPUT"

    # Restore stashed changes
    if $STASHED; then
        echo ""
        warn "Restoring stashed changes..."
        if git -C "$REPO_DIR" stash pop >/dev/null 2>&1; then
            ok "Local changes restored."
        else
            warn "Could not auto-restore local changes. Run 'git stash pop' manually."
        fi
    fi

    ok "Git update complete."

else
    # ---- Zip-based update ----
    echo "  Downloading latest version from GitHub..."

    TEMP_DIR="$(mktemp -d)"
    ZIP_FILE="$TEMP_DIR/pre-latest.zip"
    EXTRACT_DIR="$TEMP_DIR/pre-main"

    # Download
    curl -fSL "${REPO_URL}/archive/refs/heads/main.zip" -o "$ZIP_FILE" || {
        rm -rf "$TEMP_DIR"
        fail "Download failed. Check your internet connection."
    }
    ok "Downloaded."

    # Extract
    echo "  Extracting..."
    unzip -q "$ZIP_FILE" -d "$TEMP_DIR" || {
        rm -rf "$TEMP_DIR"
        fail "Extraction failed."
    }

    # The zip extracts to a directory like "pre-main/"
    if [ ! -d "$EXTRACT_DIR" ]; then
        # Try to find the extracted directory
        EXTRACT_DIR="$(find "$TEMP_DIR" -maxdepth 1 -type d -name 'pre-*' | head -1)"
    fi

    if [ ! -d "$EXTRACT_DIR" ]; then
        rm -rf "$TEMP_DIR"
        fail "Could not find extracted files."
    fi

    # Copy files, preserving user data
    echo "  Updating files..."

    # Update engine/ (source code, Modelfile, scripts)
    if [ -d "$EXTRACT_DIR/engine" ]; then
        rsync -a --exclude='*.o' --exclude='pre' --exclude='telegram' \
            "$EXTRACT_DIR/engine/" "$ENGINE_DIR/"
    fi

    # Update web/ (server, tools, frontend)
    if [ -d "$EXTRACT_DIR/web" ]; then
        rsync -a --exclude='node_modules/' \
            "$EXTRACT_DIR/web/" "$WEB_DIR/"
    fi

    # Update root files (scripts, README, VERSION, etc.)
    for f in VERSION README.md CLAUDE.md install.sh install-linux.sh install.ps1 install.cmd \
             "Launch PRE.command" "Launch PRE.cmd" "PRE Tray.cmd" \
             "Install PRE.command" "Update PRE.command" "Update PRE.cmd" \
             update.sh update.ps1 system.md benchmark.sh; do
        if [ -f "$EXTRACT_DIR/$f" ]; then
            cp "$EXTRACT_DIR/$f" "$REPO_DIR/$f"
        fi
    done

    # Make scripts executable
    chmod +x "$REPO_DIR"/*.sh "$REPO_DIR"/*.command 2>/dev/null || true

    # Clean up
    rm -rf "$TEMP_DIR"
    ok "Files updated."
fi

# -- Update dependencies ----------------------------------------------------

step "Updating dependencies"

if [ -f "$WEB_DIR/package.json" ]; then
    echo "  Running npm install..."
    cd "$WEB_DIR"
    npm install --silent 2>&1 | tail -3
    ok "Dependencies updated."
else
    warn "web/package.json not found - skipping npm install."
fi

# -- Rebuild binaries (macOS only) ------------------------------------------

if [ -f "$ENGINE_DIR/Makefile" ] && command -v make >/dev/null 2>&1; then
    step "Rebuilding binaries"

    # Desktop app + menu bar (works on both Apple Silicon and Intel)
    if command -v swiftc >/dev/null 2>&1; then
        echo "  Building desktop app and menu bar..."
        cd "$ENGINE_DIR"
        make menubar desktop 2>&1 | tail -3
        if [ -x "$ENGINE_DIR/PRE Desktop.app/Contents/MacOS/pre-desktop" ]; then
            ok "Built: PRE Desktop.app (native desktop wrapper)"
            # Generate icon if missing
            if [ ! -f "$ENGINE_DIR/PRE Desktop.app/Contents/Resources/AppIcon.icns" ] && \
               [ -x "$ENGINE_DIR/gen-icon.sh" ]; then
                "$ENGINE_DIR/gen-icon.sh" 2>/dev/null && ok "Generated app icon" || true
            fi
            # Install to ~/Applications
            mkdir -p "$HOME/Applications"
            rm -rf "$HOME/Applications/PRE Desktop.app"
            cp -R "$ENGINE_DIR/PRE Desktop.app" "$HOME/Applications/PRE Desktop.app"
            ok "Installed PRE Desktop.app → ~/Applications/"
        fi
        if [ -x "$ENGINE_DIR/PRE.app/Contents/MacOS/pre-menubar" ]; then
            ok "Built: PRE.app (menu bar helper)"
            mkdir -p "$HOME/Applications"
            rm -rf "$HOME/Applications/PRE.app"
            cp -R "$ENGINE_DIR/PRE.app" "$HOME/Applications/PRE.app"
            ok "Installed PRE.app → ~/Applications/"
        fi
    fi

    # CLI binary (Apple Silicon only)
    IS_ARM64=false
    if [ "$(uname -m)" = "arm64" ]; then
        IS_ARM64=true
    fi
    if $IS_ARM64; then
        if ask_yn "  Rebuild the CLI binary? [y/N]" "N"; then
            cd "$ENGINE_DIR"
            make pre 2>&1 | tail -5
            ok "CLI rebuilt."
        else
            echo "  Skipped. Run 'cd engine && make pre' to rebuild later."
        fi
    fi
fi

# -- Check if Modelfile changed ---------------------------------------------

# If the Modelfile changed, the custom model may need recreation
if $IS_GIT; then
    MODELFILE_CHANGED="$(git -C "$REPO_DIR" diff HEAD~1 --name-only 2>/dev/null | grep -c 'Modelfile' || true)"
    if [ "$MODELFILE_CHANGED" -gt 0 ]; then
        warn "Modelfile has changed. You may want to recreate the custom model:"
        echo "    ollama create pre-gemma4 -f engine/Modelfile"
    fi
fi

# -- Post-update migration (1.0 → 2.0+) ------------------------------------
# Ensures features added after the user's initial install are properly set up.
# Each step is idempotent — safe to run on every update.

step "Post-update migration"

PRE_DIR="$HOME/.pre"
MIGRATION_COUNT=0

# ── Create missing data directories ──────────────────────────────────────────
for dir in "$PRE_DIR/agent-checkpoints" "$PRE_DIR/models" \
           "$PRE_DIR/training" "$PRE_DIR/training/exports" \
           "$PRE_DIR/training/adapters" \
           "$PRE_DIR/custom_tools" "$PRE_DIR/workflows" \
           "$PRE_DIR/rag"; do
    if [ ! -d "$dir" ]; then
        mkdir -p "$dir"
        MIGRATION_COUNT=$((MIGRATION_COUNT + 1))
    fi
done
if [ "$MIGRATION_COUNT" -gt 0 ]; then
    ok "Created $MIGRATION_COUNT missing data directories"
    MIGRATION_COUNT=0
fi

# ── Fast tier model (pre-gemma4-fast) ────────────────────────────────────────
FAST_MODEL="pre-gemma4-fast"
FAST_MODELFILE="$ENGINE_DIR/Modelfile.fast"
if command -v ollama >/dev/null 2>&1 && [ -f "$FAST_MODELFILE" ]; then
    if ! ollama list 2>/dev/null | grep -q "$FAST_MODEL"; then
        echo ""
        echo "  The multi-model router's fast tier uses $FAST_MODEL (Gemma 4 E2B, ~3 GB)."
        echo "  This model provides ~140 tok/s responses for simple queries."
        if ask_yn "  Create the fast tier model? [Y/n]" "Y"; then
            echo "  Pulling base model (gemma4:e2b)..."
            if ollama pull gemma4:e2b 2>/dev/null; then
                if ollama create "$FAST_MODEL" -f "$FAST_MODELFILE" 2>/dev/null; then
                    ok "Fast tier model created ($FAST_MODEL)"
                else
                    warn "Failed to create $FAST_MODEL — fast tier will fall back to main model"
                fi
            else
                warn "Failed to pull gemma4:e2b — fast tier will be unavailable"
            fi
        else
            echo "  Skipped. Create later with:"
            echo "    ollama pull gemma4:e2b && ollama create $FAST_MODEL -f engine/Modelfile.fast"
        fi
    else
        ok "Fast tier model already exists ($FAST_MODEL)"
    fi
fi

# ── Ollama environment variables (NUM_PARALLEL, MAX_LOADED_MODELS) ───────────
# 2.0 uses NUM_PARALLEL=2 for parallel inference (user + background agents)
# and MAX_LOADED_MODELS=2 for main model + embedding model simultaneously.
if [ "$(uname -s)" = "Darwin" ]; then
    PLIST="$HOME/Library/LaunchAgents/com.pre.server.plist"
    if [ -f "$PLIST" ]; then
        NEEDS_PLIST_UPDATE=false

        # Check OLLAMA_NUM_PARALLEL
        CURRENT_PARALLEL="$(/usr/libexec/PlistBuddy -c "Print :EnvironmentVariables:OLLAMA_NUM_PARALLEL" "$PLIST" 2>/dev/null || echo "not_set")"
        if [ "$CURRENT_PARALLEL" = "1" ] || [ "$CURRENT_PARALLEL" = "not_set" ]; then
            NEEDS_PLIST_UPDATE=true
        fi

        # Check OLLAMA_MAX_LOADED_MODELS
        CURRENT_LOADED="$(/usr/libexec/PlistBuddy -c "Print :EnvironmentVariables:OLLAMA_MAX_LOADED_MODELS" "$PLIST" 2>/dev/null || echo "not_set")"
        if [ "$CURRENT_LOADED" = "1" ] || [ "$CURRENT_LOADED" = "not_set" ]; then
            NEEDS_PLIST_UPDATE=true
        fi

        if $NEEDS_PLIST_UPDATE; then
            echo ""
            echo "  PRE 2.0 uses OLLAMA_NUM_PARALLEL=2 (parallel user + background inference)"
            echo "  and OLLAMA_MAX_LOADED_MODELS=2 (main model + embedding model)."
            if ask_yn "  Update LaunchAgent with new Ollama settings? [Y/n]" "Y"; then
                /usr/libexec/PlistBuddy -c "Set :EnvironmentVariables:OLLAMA_NUM_PARALLEL 2" "$PLIST" 2>/dev/null || \
                    /usr/libexec/PlistBuddy -c "Add :EnvironmentVariables:OLLAMA_NUM_PARALLEL string 2" "$PLIST" 2>/dev/null || true
                /usr/libexec/PlistBuddy -c "Set :EnvironmentVariables:OLLAMA_MAX_LOADED_MODELS 2" "$PLIST" 2>/dev/null || \
                    /usr/libexec/PlistBuddy -c "Add :EnvironmentVariables:OLLAMA_MAX_LOADED_MODELS string 2" "$PLIST" 2>/dev/null || true
                ok "LaunchAgent updated (NUM_PARALLEL=2, MAX_LOADED_MODELS=2)"
                # Also inject into current launchd session so Ollama picks them up
                launchctl setenv OLLAMA_NUM_PARALLEL "2" 2>/dev/null || true
                launchctl setenv OLLAMA_MAX_LOADED_MODELS "2" 2>/dev/null || true
            fi
        else
            ok "LaunchAgent Ollama settings are current"
        fi
    fi
fi

# Linux: update systemd unit if present
if [ "$(uname -s)" = "Linux" ]; then
    SYSTEMD_UNIT="$HOME/.config/systemd/user/pre-server.service"
    if [ -f "$SYSTEMD_UNIT" ]; then
        if ! grep -q "OLLAMA_NUM_PARALLEL=2" "$SYSTEMD_UNIT" 2>/dev/null; then
            echo ""
            echo "  PRE 2.0 uses OLLAMA_NUM_PARALLEL=2 for parallel inference."
            if ask_yn "  Update systemd unit with new Ollama settings? [Y/n]" "Y"; then
                # Add or update environment lines
                if grep -q "OLLAMA_NUM_PARALLEL" "$SYSTEMD_UNIT"; then
                    sed -i 's/OLLAMA_NUM_PARALLEL=[0-9]*/OLLAMA_NUM_PARALLEL=2/' "$SYSTEMD_UNIT"
                else
                    sed -i '/\[Service\]/a Environment=OLLAMA_NUM_PARALLEL=2' "$SYSTEMD_UNIT"
                fi
                if grep -q "OLLAMA_MAX_LOADED_MODELS" "$SYSTEMD_UNIT"; then
                    sed -i 's/OLLAMA_MAX_LOADED_MODELS=[0-9]*/OLLAMA_MAX_LOADED_MODELS=2/' "$SYSTEMD_UNIT"
                else
                    sed -i '/\[Service\]/a Environment=OLLAMA_MAX_LOADED_MODELS=2' "$SYSTEMD_UNIT"
                fi
                systemctl --user daemon-reload 2>/dev/null || true
                ok "systemd unit updated (NUM_PARALLEL=2, MAX_LOADED_MODELS=2)"
            fi
        else
            ok "systemd Ollama settings are current"
        fi
    fi
fi

ok "Post-update migration complete"

# -- Restart server if it was running ----------------------------------------

if $SERVER_WAS_RUNNING; then
    step "Restarting server"
    if ask_yn "  Restart PRE server? [Y/n]" "Y"; then
        cd "$WEB_DIR"
        nohup node server.js > /tmp/pre-server.log 2>&1 &
        sleep 2
        if curl -sf "http://localhost:${PRE_PORT}/api/status" >/dev/null 2>&1; then
            ok "Server restarted on port $PRE_PORT."
        else
            warn "Server may still be starting. Check http://localhost:$PRE_PORT"
        fi
    fi
fi

# -- Done -------------------------------------------------------------------

step "Update Complete!"

NEW_VERSION="unknown"
if [ -f "$REPO_DIR/VERSION" ]; then
    NEW_VERSION="$(cat "$REPO_DIR/VERSION" | tr -d '[:space:]')"
fi

echo ""
ok "PRE updated to version $NEW_VERSION"
echo ""
echo "  Changes in this update:"
if $IS_GIT; then
    echo "    git log --oneline -10"
else
    echo "    See: ${REPO_URL}/commits/main"
fi
echo ""
echo "  If you encounter issues after updating:"
echo "    1. Restart the server: cd web && node server.js"
echo "    2. Recreate the model: ollama create pre-gemma4 -f engine/Modelfile"
echo "    3. Report issues: ${REPO_URL}/issues"
echo ""
