install.sh

#!/usr/bin/env bash
set -euo pipefail

# install.sh — tqstack v3 installer
# Supports pure SDPA models and hybrid GDN+SDPA models (Qwen 3.5/3.6)

REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
APP_HOME="${HOME}/ai/turboquant"
VENV_DIR="${APP_HOME}/.venv"
TQ_DIR="${APP_HOME}/turboquant-mlx"
PLIST_DIR="${HOME}/Library/LaunchAgents"
USER_ID="$(id -u)"
PYTHON_BIN="${PYTHON_BIN:-python3}"
CONFIG_FILE="${APP_HOME}/config.env"

TURBOQUANT_REPO="https://github.com/sharpner/turboquant-mlx.git"
TURBOQUANT_COMMIT="main"

mkdir -p "${APP_HOME}" "${PLIST_DIR}"

# --- Step 0: Preflight checks ---
printf '\n[0/7] Running preflight checks\n'
source "${REPO_ROOT}/scripts/preflight-check.sh"

# --- Step 1: Write config file ---
printf '\n[1/7] Writing configuration\n'

# Detect model architecture for optimal routing threshold
SIDECAR_MODEL="${TQ_SIDECAR_MODEL:-mlx-community/Qwen3.5-35B-A3B-4bit}"
IS_HYBRID="false"
ROUTING_THRESHOLD=6000

case "${SIDECAR_MODEL}" in
    *Qwen3.5*|*Qwen3.6*|*qwen3_5*|*qwen3_6*|*Qwen3-Next*|*qwen3_next*)
        IS_HYBRID="true"
        ROUTING_THRESHOLD=16000
        printf '[INFO] Hybrid GDN+SDPA model detected: %s\n' "${SIDECAR_MODEL}"
        printf '[INFO] Routing threshold set to %d tokens\n' "${ROUTING_THRESHOLD}"
        ;;
    *)
        printf '[INFO] Standard SDPA model detected: %s\n' "${SIDECAR_MODEL}"
        ;;
esac

cat > "${CONFIG_FILE}" <<EOF
# tqstack v3 configuration
# Generated by install.sh on $(date -u +"%Y-%m-%dT%H:%M:%SZ")
SIDECAR_ENABLED=${SIDECAR_ENABLED}
MLX_ENABLED=${MLX_ENABLED}
TOTAL_MEM_GB=${TOTAL_MEM_GB}

# Models
OLLAMA_MODEL=qwen3.5:35b-a3b-coding-nvfp4
SIDECAR_MODEL=${SIDECAR_MODEL}

# Ports
SIDECAR_PORT=8001
ROUTER_PORT=8000
OLLAMA_PORT=11434

# TurboQuant cache
TQ_KV_BITS=4
TQ_KV_GROUP_SIZE=64
TQ_USE_ROTATION=1
TQ_USE_NORMALIZATION=1
TQ_BOUNDARY_LAYERS=0

# Routing
PROMPT_TOKEN_THRESHOLD=${ROUTING_THRESHOLD}
TQ_IS_HYBRID_MODEL=${IS_HYBRID}
EOF
echo "[OK] Config written to ${CONFIG_FILE}"

# --- Step 2: Configure Ollama ---
printf '\n[2/7] Configuring Ollama environment\n'
export MLX_ENABLED
bash "${REPO_ROOT}/scripts/start-ollama-optimized.sh"

# --- Step 3: Create virtualenv and install Python deps ---
printf '\n[3/7] Creating virtualenv and installing dependencies\n'
"${PYTHON_BIN}" -m venv "${VENV_DIR}"
source "${VENV_DIR}/bin/activate"
python -m pip install -U pip wheel setuptools
pip install \
  "fastapi>=0.110" \
  "uvicorn[standard]>=0.29" \
  "httpx>=0.27" \
  "pydantic>=2.0,<3" \
  "mlx>=0.22" \
  "mlx-lm>=0.22" \
  "huggingface_hub>=0.23"

# --- Step 4: Install TurboQuant (only if sidecar enabled) ---
if [ "$SIDECAR_ENABLED" = "true" ]; then
  printf '\n[4/7] Installing TurboQuant MLX dependency\n'
  if [ ! -d "${TQ_DIR}" ]; then
    git clone "${TURBOQUANT_REPO}" "${TQ_DIR}"
    cd "${TQ_DIR}" && git checkout "${TURBOQUANT_COMMIT}" && cd "${REPO_ROOT}"
  else
    printf 'TurboQuant already present at %s\n' "${TQ_DIR}"
  fi
else
  printf '\n[4/7] Skipping TurboQuant (sidecar disabled for %sGB machine)\n' "${TOTAL_MEM_GB}"
fi

# --- Step 5: Write LaunchAgent plists ---
printf '\n[5/7] Writing LaunchAgent plists\n'

# Router plist
sed \
  -e "s#__REPO_ROOT__#${REPO_ROOT}#g" \
  -e "s#__HOME__#${HOME}#g" \
  "${REPO_ROOT}/launchagents/com.tqstack-router.plist.template" \
  > "${PLIST_DIR}/com.tqstack.router.plist"

# Sidecar plist (only if enabled)
if [ "$SIDECAR_ENABLED" = "true" ]; then
  sed \
    -e "s#__REPO_ROOT__#${REPO_ROOT}#g" \
    -e "s#__HOME__#${HOME}#g" \
    "${REPO_ROOT}/launchagents/com.tqstack-sidecar.plist.template" \
    > "${PLIST_DIR}/com.tqstack.sidecar.plist"
fi

chmod +x "${REPO_ROOT}/scripts/"*.sh

# --- Step 6: Pull the Ollama model ---
printf '\n[6/7] Pulling Ollama model (this may take a while for 22GB)\n'
ollama pull qwen3.5:35b-a3b-coding-nvfp4 || {
  echo '[WARN] Model pull failed. Pull it later with: ollama pull qwen3.5:35b-a3b-coding-nvfp4'
}

# --- Step 7: Load LaunchAgents ---
printf '\n[7/7] Loading LaunchAgents\n'
bash "${REPO_ROOT}/scripts/load-launchagents.sh" || true

# --- Done ---
printf '\n========================================\n'
printf '  Installation complete!\n'
printf '========================================\n\n'
printf 'Configuration:\n'
printf '  MLX backend:    %s\n' "$MLX_ENABLED"
printf '  Sidecar:        %s\n' "$SIDECAR_ENABLED"
printf '  Hybrid model:   %s\n' "$IS_HYBRID"
printf '  Memory:         %s GB\n' "$TOTAL_MEM_GB"
printf '  Ollama model:   qwen3.5:35b-a3b-coding-nvfp4\n'
printf '  Sidecar model:  %s\n' "${SIDECAR_MODEL}"
printf '  Route at:       >= %d tokens\n' "${ROUTING_THRESHOLD}"
printf '\nNext steps:\n'
printf '  1. Quit and reopen Ollama\n'
printf '  2. Test health: curl http://127.0.0.1:8000/health\n'
printf '  3. Point clients to: http://127.0.0.1:8000/v1\n'
if [ "$SIDECAR_ENABLED" = "true" ]; then
  printf '  4. Sidecar test: curl http://127.0.0.1:8001/health\n'
fi
printf '\n'