#!/usr/bin/env bash

set -euo pipefail

OUTPUT_ROOT="${OUTPUT_ROOT:-/home/admin/cpfs/wjh/frontier-community-qwen235-smoke-20260715}"
FRONTIER_ROOT="${FRONTIER_ROOT:-${OUTPUT_ROOT}/Frontier-d9cfeb6}"
VENV_ROOT="${VENV_ROOT:-/tmp/wjh-frontier-vllm0102-smoke/.venv}"
LOG_DIR="${OUTPUT_ROOT}/logs"

export UV_HTTP_TIMEOUT="${UV_HTTP_TIMEOUT:-300}"

mkdir -p "${LOG_DIR}" "${OUTPUT_ROOT}/provenance"
exec > >(tee -a "${LOG_DIR}/setup_env.log") 2>&1

echo "SETUP_ENV output=${OUTPUT_ROOT} frontier=${FRONTIER_ROOT} venv=${VENV_ROOT} python=/usr/local/bin/python3.12 vllm=0.10.2 transformers=4.55.2"

if [[ ! -f "${FRONTIER_ROOT}/pyproject.toml" ]]; then
  echo "ERROR: Frontier source is missing at ${FRONTIER_ROOT}" >&2
  exit 1
fi

uv venv --clear --python /usr/local/bin/python3.12 "${VENV_ROOT}"
uv pip install \
  --python "${VENV_ROOT}/bin/python" \
  "vllm==0.10.2" \
  "transformers==4.55.2" \
  "flashinfer-python>=0.3,<0.4" \
  -e "${FRONTIER_ROOT}[test]"

"${VENV_ROOT}/bin/python" - <<'PY' | tee "${OUTPUT_ROOT}/provenance/environment.json"
import importlib.metadata as metadata
import json
import platform

import flashinfer
import frontier
import torch
import vllm

record = {
    "python": platform.python_version(),
    "torch": torch.__version__,
    "torch_cuda": torch.version.cuda,
    "vllm_import_version": vllm.__version__,
    "vllm_metadata_version": metadata.version("vllm"),
    "vllm_path": vllm.__file__,
    "transformers_metadata_version": metadata.version("transformers"),
    "tokenizers_metadata_version": metadata.version("tokenizers"),
    "flashinfer_metadata_version": metadata.version("flashinfer-python"),
    "flashinfer_path": flashinfer.__file__,
    "frontier_metadata_version": metadata.version("frontier-simulator"),
    "frontier_path": frontier.__path__[0],
}
print(json.dumps(record, indent=2, sort_keys=True))
assert record["vllm_import_version"] == "0.10.2", record
assert record["vllm_metadata_version"] == "0.10.2", record
assert record["transformers_metadata_version"] == "4.55.2", record
PY

uv pip freeze --python "${VENV_ROOT}/bin/python" > "${OUTPUT_ROOT}/provenance/requirements.freeze.txt"
sha256sum \
  "${OUTPUT_ROOT}/provenance/environment.json" \
  "${OUTPUT_ROOT}/provenance/requirements.freeze.txt" \
  > "${OUTPUT_ROOT}/provenance/environment.sha256"

echo "SETUP_ENV_COMPLETE"