#!/usr/bin/env bash

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SHARED_ROOT="${SHARED_ROOT:-/shared_work/xuan}"
BOOTSTRAP_CONDA_EXE="${BOOTSTRAP_CONDA_EXE:-/work/xuan/miniconda3/bin/conda}"
CONDA_BASE_PREFIX="${CONDA_BASE_PREFIX:-$SHARED_ROOT/miniforge3}"
CONDA_EXE="${CONDA_EXE:-$CONDA_BASE_PREFIX/bin/conda}"
ENV_NAME="${ENV_NAME:-xvla-stable}"
ENV_PREFIX="${ENV_PREFIX:-$SHARED_ROOT/conda-envs/$ENV_NAME}"
CONDARC_PATH="${CONDARC_PATH:-$REPO_ROOT/.condarc.shared}"

mkdir -p \
  "$SHARED_ROOT/conda-envs" \
  "$SHARED_ROOT/conda-pkgs" \
  "$SHARED_ROOT/.cache/pip" \
  "$SHARED_ROOT/.cache/huggingface" \
  "$SHARED_ROOT/.cache/torch" \
  "$SHARED_ROOT/.cache/matplotlib" \
  "$SHARED_ROOT/.cache/xdg" \
  "$SHARED_ROOT/.config"

export CONDARC="$CONDARC_PATH"
export CONDA_PKGS_DIRS="$SHARED_ROOT/conda-pkgs"
export CONDA_ENVS_PATH="$SHARED_ROOT/conda-envs"
export PIP_CACHE_DIR="$SHARED_ROOT/.cache/pip"
export HF_HOME="$SHARED_ROOT/.cache/huggingface"
export TRANSFORMERS_CACHE="$HF_HOME/transformers"
export HUGGINGFACE_HUB_CACHE="$HF_HOME/hub"
export TORCH_HOME="$SHARED_ROOT/.cache/torch"
export MPLCONFIGDIR="$SHARED_ROOT/.cache/matplotlib"
export XDG_CACHE_HOME="$SHARED_ROOT/.cache/xdg"
export PYTHONNOUSERSITE=1

if [[ ! -x "$CONDA_EXE" ]]; then
  echo "Shared conda bootstrap not found at $CONDA_EXE" >&2
  echo "Install Miniforge under $CONDA_BASE_PREFIX first, then rerun this script." >&2
  exit 1
fi

if [[ -d "$ENV_PREFIX" ]]; then
  echo "Conda environment already exists at $ENV_PREFIX"
else
  "$CONDA_EXE" env create --prefix "$ENV_PREFIX" -f "$REPO_ROOT/environment.yml"
fi

echo
echo "Shared X-VLA environment is ready at $ENV_PREFIX"
echo "Activate it with:"
echo "  source \"$REPO_ROOT/scripts/activate_shared_env.sh\""