#!/usr/bin/env bash
# Rebuild an isolated VBench scoring environment without modifying sglang/deploy/vbench.
set -Eeuo pipefail

CONDA=${CONDA:-/root/.miniconda3/bin/conda}
SOURCE_ENV=${SOURCE_ENV:-vbench}
TARGET_ENV=${TARGET_ENV:-vbench-score}
TARGET_PREFIX=${TARGET_PREFIX:-/root/.miniconda3/envs/$TARGET_ENV}
VBENCH_REPO=${VBENCH_REPO:-/data/wxy/VBench}
CHECKER=${CHECKER:-/data/wxy/check_vbench_score_env.py}
LOG_ROOT=${LOG_ROOT:-/data/wxy/vbench_score_setup_logs}
PIP_INDEX_URL=${PIP_INDEX_URL:-https://pypi.org/simple}
HF_ENDPOINT=${HF_ENDPOINT:-https://hf-mirror.com}
export PIP_INDEX_URL HF_ENDPOINT

mkdir -p "$LOG_ROOT"
log() { printf '[%s] %s\n' "$(date '+%F %T')" "$*"; }
die() { log "ERROR: $*" >&2; exit 1; }

[[ -x "$CONDA" ]] || die "missing conda: $CONDA"
[[ -d "$VBENCH_REPO/vbench" ]] || die "missing VBench repo: $VBENCH_REPO"
[[ -f "$CHECKER" ]] || die "missing acceptance checker: $CHECKER"

if [[ ! -x "$TARGET_PREFIX/bin/python" ]]; then
  log "cloning $SOURCE_ENV into isolated environment $TARGET_ENV"
  "$CONDA" create --name "$TARGET_ENV" --clone "$SOURCE_ENV" --yes
else
  log "reusing existing isolated prefix $TARGET_PREFIX"
fi

PYTHON="$TARGET_PREFIX/bin/python"
CONSTRAINTS="$LOG_ROOT/constraints.txt"
cat > "$CONSTRAINTS" <<'EOF'
torch==2.11.0
torchvision==0.26.0
torchaudio==2.11.0
numpy==1.26.4
transformers==4.36.2
tokenizers==0.15.2
timm==0.9.12
omegaconf==2.3.0
iopath==0.1.9
setuptools==80.9.0
scipy==1.14.1
tifffile==2024.9.20
EOF

log "recording inherited CUDA stack"
"$PYTHON" - <<'PY' > "$LOG_ROOT/inherited_versions.txt"
import cv2, numpy, sys, torch, torchvision
print(sys.version)
print("torch", torch.__version__, "cuda", torch.version.cuda)
print("torchvision", torchvision.__version__)
print("numpy", numpy.__version__)
print("cv2", cv2.__version__)
print("capability", torch.cuda.get_device_capability())
PY

log "installing compatibility pins required by the official VBench repository"
"$PYTHON" -m pip install --upgrade --upgrade-strategy only-if-needed -c "$CONSTRAINTS" \
  'setuptools==80.9.0' \
  'numpy==1.26.4' \
  'scipy==1.14.1' \
  'tifffile==2024.9.20' \
  '/data/wxy/vbench_wheels/tokenizers-0.15.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl' \
  '/data/wxy/vbench_wheels/transformers-4.36.2-py3-none-any.whl' \
  'timm==0.9.12' \
  'omegaconf==2.3.0' \
  'iopath==0.1.9'

log "installing the missing runtime dependency closure"
"$PYTHON" -m pip install --upgrade --upgrade-strategy only-if-needed -c "$CONSTRAINTS" \
  'pyparsing>=3.1' 'contourpy>=1.2' 'cycler>=0.12' 'fonttools>=4.50' 'kiwisolver>=1.4' \
  'cython>=3.0' 'black>=24.0' 'hydra-core>=1.1,<1.4' 'tensorboard>=2.16' \
  'botocore>=1.43,<1.44' 'jmespath>=1.0' 's3transfer>=0.15,<0.16' \
  'future>=1.0'

log "running the full scoring-environment acceptance test"
"$PYTHON" "$CHECKER" --vbench-repo "$VBENCH_REPO" \
  > "$LOG_ROOT/acceptance.json" 2> "$LOG_ROOT/acceptance.stderr.log"

"$PYTHON" -m pip freeze > "$LOG_ROOT/pip-freeze.txt"
"$CONDA" list --name "$TARGET_ENV" > "$LOG_ROOT/conda-list.txt"
"$PYTHON" -m pip check > "$LOG_ROOT/pip-check.txt" 2>&1 || true
touch "$LOG_ROOT/DONE"
log "isolated VBench scoring environment ready: $TARGET_PREFIX"