#!/usr/bin/env bash
# Fresh-vllm + Mooncake install for microbench (Phase 0).
#
# Idempotent. Runs from any dash machine that mounts /home/admin/cpfs/wjh/.
# Single venv at FRESH_ROOT/.venv is shared by all dash hosts (cpfs is mounted
# at the same path on dash0/1/2).
#
# Usage on each host:
#   bash microbench/fresh_setup/install.sh

set -eo pipefail
# Note: NOT -u; bashrc / venv activation reference unset vars (PS1, etc.).

FRESH_ROOT="/home/admin/cpfs/wjh/agentic-kv-fresh"
VENV_DIR="${FRESH_ROOT}/.venv"
VLLM_VERSION="0.18.1"

mkdir -p "${FRESH_ROOT}"
cd "${FRESH_ROOT}"

# 1. uv venv (Python 3.12). Skip if already exists and is valid.
if [[ ! -f "${VENV_DIR}/bin/python" ]]; then
    echo "[install] creating venv at ${VENV_DIR}"
    uv venv --python 3.12 "${VENV_DIR}"
else
    echo "[install] venv exists at ${VENV_DIR}, reusing"
fi

# 2. Set proxy + activate venv. Avoid sourcing bashrc directly (it has
#    interactive-only stuff); define proxyOn locally instead.
proxyOn() {
    local p="http://ipads:ipads123@127.0.0.1:11235"
    export http_proxy="$p" https_proxy="$p" all_proxy="$p"
    echo "[install] proxy set: $p"
}
proxyOn

# shellcheck disable=SC1090
source "${VENV_DIR}/bin/activate"

echo "[install] python: $(which python)  ($(python --version))"

# 3. Install vllm wheel.
if ! python -c "import vllm" 2>/dev/null; then
    echo "[install] installing vllm==${VLLM_VERSION}"
    uv pip install "vllm==${VLLM_VERSION}"
else
    INSTALLED=$(python -c "import vllm; print(vllm.__version__)")
    echo "[install] vllm already installed: ${INSTALLED}"
fi

# 4. Install Mooncake transfer engine (vllm's standard kv_connector backend).
# vllm 0.18.x mooncake connector imports mooncake.engine; the upstream package
# is `mooncake-transfer-engine`.
if ! python -c "import mooncake" 2>/dev/null; then
    echo "[install] installing mooncake-transfer-engine"
    uv pip install mooncake-transfer-engine
else
    echo "[install] mooncake already installed"
fi

# 5. Common test deps.
uv pip install httpx pytest numpy

# 6. Smoke-test imports.
python - <<'PY'
import vllm
import mooncake
print(f"[smoke] vllm        = {vllm.__version__}")
print(f"[smoke] mooncake    = {getattr(mooncake, '__version__', 'unknown')}")
import vllm.distributed.kv_transfer as kvt
print(f"[smoke] kv_transfer = {kvt.__file__}")
PY

echo "[install] DONE"