#!/bin/bash
# Recreate the inference environment for the Moshi RAG student.
#
#   bash setup_env.sh            # -> /workspace/menv (or ./menv if no /workspace)
#   bash setup_env.sh /opt/menv  # explicit location
#
# Needs python3.11 + CUDA 12.1-capable driver.  A100 / 3090 / 4090 / A10 all
# work with these wheels; on newer drivers they still run (no JIT kernels).
set -euo pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"
if [ -n "${1:-}" ]; then
  ENV_DIR="$1"
elif [ -n "${ENV_DIR:-}" ]; then
  ENV_DIR="$ENV_DIR"
elif [ -d /workspace ]; then
  ENV_DIR=/workspace/menv
else
  ENV_DIR="$HERE/menv"
fi
PYTHON="${PYTHON:-python3.11}"

command -v "$PYTHON" >/dev/null 2>&1 || {
  echo "ERROR: $PYTHON not found. Install Python 3.11 first." >&2
  exit 1
}

echo "== venv: $ENV_DIR"
"$PYTHON" -m venv "$ENV_DIR"
"$ENV_DIR/bin/pip" install --upgrade pip wheel

echo "== torch 2.5.1 / cu121"
"$ENV_DIR/bin/pip" install torch==2.5.1 torchaudio==2.5.1 \
  --index-url https://download.pytorch.org/whl/cu121

echo "== server requirements"
"$ENV_DIR/bin/pip" install -r "$HERE/requirements-server.txt"

echo "== moshi (no-deps: it would pull sphn 0.2 whose opus API the server does not use)"
"$ENV_DIR/bin/pip" install --no-deps moshi==0.2.13

echo "== check"
"$ENV_DIR/bin/python" - <<'PY'
import aiohttp
import numpy as np
import sentence_transformers
import sentencepiece
import sphn
import torch
print("torch", torch.__version__, "| cuda", torch.cuda.is_available(),
      "| devices", torch.cuda.device_count())
print("sentencepiece", sentencepiece.__version__)

# audio path smoke test (the server streams opus both ways)
w = sphn.OpusStreamWriter(24000)
w.append_pcm(np.zeros(1920, dtype=np.float32))
blob = w.read_bytes()
r = sphn.OpusStreamReader(24000)
r.append_bytes(blob)
print("opus ok:", len(blob), "bytes ->", r.read_pcm().shape)
PY
echo
echo "OK. Environment ready: $ENV_DIR"
echo "Next: HF_TOKEN=... bash download_assets.sh   (model + index)"
echo "Then: bash serve_s28_web.sh"