#!/bin/bash
# switch-model.sh — Switch PRE between q8_0 and q4_K_M quantization
#
# For existing installs that want to change model size without
# re-running the full installer.
#
# Usage:
#   ./switch-model.sh q4     Switch to q4_K_M (~15GB, lighter on RAM)
#   ./switch-model.sh q8     Switch to q8_0 (~28GB, near-lossless)
#   ./switch-model.sh        Show current model and available options

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"
MODELFILE="$ENGINE_DIR/Modelfile"
PORT="${PRE_PORT:-11434}"

ok()   { echo -e "${GREEN}$1${RESET}"; }
warn() { echo -e "${YELLOW}$1${RESET}"; }
fail() { echo -e "${RED}$1${RESET}"; exit 1; }

# Detect current model by size (reliable — blob paths don't contain quant name)
CURRENT=""
if ollama list 2>/dev/null | grep -q "pre-gemma4"; then
    MODEL_GB=$(ollama list 2>/dev/null | grep "pre-gemma4" | awk '{print $3}')
    if [ "$MODEL_GB" -ge 25 ] 2>/dev/null; then
        CURRENT="q8_0"
    elif [ "$MODEL_GB" -ge 12 ] 2>/dev/null; then
        CURRENT="q4_K_M"
    else
        CURRENT="unknown (${MODEL_GB}GB)"
    fi
fi

if [ -z "$1" ]; then
    echo ""
    echo -e "${BOLD}PRE Model Switcher${RESET}"
    echo ""
    if [ -n "$CURRENT" ]; then
        echo -e "  Current model: ${BOLD}${CURRENT}${RESET}"
    else
        echo -e "  Current model: ${YELLOW}pre-gemma4 not found${RESET}"
    fi
    echo ""
    echo -e "  ${BOLD}Usage:${RESET}"
    echo -e "    ${CYAN}./switch-model.sh q4${RESET}   Switch to q4_K_M (~15GB, lighter on RAM)"
    echo -e "    ${CYAN}./switch-model.sh q8${RESET}   Switch to q8_0  (~28GB, near-lossless)"
    echo ""
    exit 0
fi

TARGET="$1"
case "$TARGET" in
    q4|q4_K_M)
        TARGET_BASE="gemma4:26b-a4b-it-q4_K_M"
        TARGET_QUANT="q4_K_M"
        TARGET_SIZE="15GB"
        ;;
    q8|q8_0)
        TARGET_BASE="gemma4:26b-a4b-it-q8_0"
        TARGET_QUANT="q8_0"
        TARGET_SIZE="28GB"
        ;;
    *)
        fail "Unknown quantization: $TARGET (use q4 or q8)"
        ;;
esac

if [ "$CURRENT" = "$TARGET_QUANT" ]; then
    ok "Already running $TARGET_QUANT — nothing to do."
    exit 0
fi

echo ""
echo -e "${BOLD}Switching PRE from ${CURRENT:-unknown} → ${TARGET_QUANT}${RESET}"
echo ""

# Step 1: Pull the base model if needed
MODEL_TAG=$(echo "$TARGET_BASE" | sed 's/.*://')
if ollama list 2>/dev/null | grep -q "gemma4.*${MODEL_TAG}"; then
    ok "  Base model $TARGET_BASE already available."
else
    echo "  Pulling $TARGET_BASE (~${TARGET_SIZE})..."
    ollama pull "$TARGET_BASE" || fail "Failed to pull $TARGET_BASE"
    ok "  Base model downloaded."
fi

# Step 2: Create the custom model from Modelfile
if [ ! -f "$MODELFILE" ]; then
    fail "Modelfile not found at $MODELFILE"
fi

EFFECTIVE_MODELFILE="$MODELFILE"
if [ "$TARGET_QUANT" != "q8_0" ]; then
    EFFECTIVE_MODELFILE="/tmp/pre-Modelfile-${TARGET_QUANT}"
    sed "s|FROM gemma4:26b-a4b-it-q8_0|FROM $TARGET_BASE|" "$MODELFILE" > "$EFFECTIVE_MODELFILE"
fi

echo "  Recreating pre-gemma4 from ${TARGET_QUANT} base..."
ollama create pre-gemma4 -f "$EFFECTIVE_MODELFILE" || fail "Failed to create pre-gemma4"
ok "  pre-gemma4 rebuilt with $TARGET_QUANT."

# Step 3: Update context window based on new model size
RAM_BYTES=$(sysctl -n hw.memsize 2>/dev/null || echo 0)
RAM_GB=$((RAM_BYTES / 1073741824))
MODEL_SIZE_GB=$(echo "$TARGET_SIZE" | sed 's/GB//')
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"

echo "$OPTIMAL_CTX" > "$HOME/.pre/context"
ok "  Context window: ${CTX_HUMAN} (${HEADROOM}GB headroom after ${TARGET_QUANT})"

# Step 4: Reload the model
echo "  Loading new model into GPU..."
curl -sf --max-time 300 "http://localhost:${PORT}/api/chat" \
    -d "{\"model\":\"pre-gemma4\",\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"keep_alive\":\"24h\",\"options\":{\"num_predict\":1,\"num_ctx\":$OPTIMAL_CTX}}" \
    >/dev/null 2>&1 || true

sleep 3
if ollama ps 2>/dev/null | grep -q "pre-gemma4"; then
    MEM=$(ollama ps 2>/dev/null | grep "pre-gemma4" | awk '{print $4}')
    ok "  Model loaded (${MEM})."
else
    warn "  Model will load on next request."
fi

echo ""
ok "Done! PRE is now running ${TARGET_QUANT} (~${TARGET_SIZE})."
echo -e "${DIM}Context window recalculated: ${CTX_HUMAN} (${OPTIMAL_CTX} tokens)${RESET}"
