#!/usr/bin/env bash
# Recreate the verified CPython 3.12 / CUDA 13.0 / vLLM 0.30.0 environment.
set -Eeuo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PY="${PY:-python3.12}"
VENV="${VENV:-$ROOT/.venv}"
CUDA_INDEX="${CUDA_INDEX:-https://download.pytorch.org/whl/cu130}"
USE_BUNDLED_WHEELS="${USE_BUNDLED_WHEELS:-1}"
RUN_TESTS="${RUN_TESTS:-1}"

if [[ "$(uname -s)" != "Linux" || "$(uname -m)" != "x86_64" ]]; then
  echo "[error] The bundled wheels require Linux x86_64; use a compatible machine or rebuild them." >&2
  exit 2
fi
"$PY" - <<'PY'
import sys
if sys.version_info[:2] != (3, 12):
    raise SystemExit(f"Python 3.12 is required, got {sys.version}")
PY

if [[ ! -x "$VENV/bin/python" ]]; then
  "$PY" -m venv "$VENV"
fi
PYTHON_BIN="$VENV/bin/python"
"$PYTHON_BIN" -m pip install --upgrade pip
"$PYTHON_BIN" -m pip install --index-url "$CUDA_INDEX" \
  torch==2.13.0 torchaudio==2.11.0 torchvision==0.28.0
"$PYTHON_BIN" -m pip install -r "$ROOT/requirements.txt"

if [[ "$USE_BUNDLED_WHEELS" == "1" ]]; then
  shopt -s nullglob
  wheels=("$ROOT"/wheels/*.whl)
  if (( ${#wheels[@]} != 3 )); then
    echo "[error] Expected three bundled kernel wheels, found ${#wheels[@]}" >&2
    exit 2
  fi
  "$PYTHON_BIN" -m pip install --no-deps --force-reinstall "${wheels[@]}"
else
  # Source builds require nvcc, ninja and a CUDA 13.0-compatible compiler toolchain.
  "$PYTHON_BIN" -m pip install packaging ninja
  MAX_JOBS="${MAX_JOBS:-8}" "$PYTHON_BIN" -m pip install --no-build-isolation \
    causal-conv1d==1.7.0 mamba-ssm==2.3.2.post1 flash-attn==2.8.3
fi

"$PYTHON_BIN" "$ROOT/service/verify_install.py" --require-gpu
"$PYTHON_BIN" "$ROOT/service/verify_bundle.py"
if [[ "$RUN_TESTS" == "1" ]]; then
  "$PYTHON_BIN" "$ROOT/service/selftest.py"
fi

echo "[done] Activate with: source \"$VENV/bin/activate\""
echo "[done] Start with:    GPU_ID=0 bash \"$ROOT/service/start.sh\""