#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
export UV_PROJECT_ENVIRONMENT="${UV_PROJECT_ENVIRONMENT:-$ROOT_DIR/.venv}"

METAX_PYPI_INDEX="${METAX_PYPI_INDEX:-https://repos.metax-tech.com/r/maca-pypi/simple}"
METAX_PYPI_HOST="${METAX_PYPI_HOST:-repos.metax-tech.com}"
METAX_INSTALL_GPU_DEPS="${METAX_INSTALL_GPU_DEPS:-false}"

METAX_TORCH_SPEC="${METAX_TORCH_SPEC:-torch}"
METAX_TORCHAUDIO_SPEC="${METAX_TORCHAUDIO_SPEC:-}"
METAX_TORCHVISION_SPEC="${METAX_TORCHVISION_SPEC:-}"
METAX_VLLM_SPEC="${METAX_VLLM_SPEC:-vllm}"

usage() {
  cat <<EOF
Usage:
  ./scripts/sync_metax_env.sh [--list-versions PACKAGE] [uv sync args...]

Environment overrides:
  METAX_TORCH_SPEC=torch==x.y.z
  METAX_TORCHAUDIO_SPEC=torchaudio==x.y.z  # optional
  METAX_TORCHVISION_SPEC=torchvision==x.y.z # optional
  METAX_VLLM_SPEC=vllm==x.y.z
  METAX_INSTALL_GPU_DEPS=false
  METAX_PYPI_INDEX=${METAX_PYPI_INDEX}
EOF
}

cd "$ROOT_DIR"

LIST_PACKAGE=""
if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
fi
if [[ "${1:-}" == "--list-versions" ]]; then
  LIST_PACKAGE="${2:-torch}"
  shift
  [[ $# -gt 0 ]] && shift
fi

if [[ -n "$LIST_PACKAGE" ]]; then
  if command -v pip3 >/dev/null 2>&1; then
    pip3 index versions "$LIST_PACKAGE" \
      -i "$METAX_PYPI_INDEX" \
      --trusted-host "$METAX_PYPI_HOST"
  elif command -v pip >/dev/null 2>&1; then
    pip index versions "$LIST_PACKAGE" \
      -i "$METAX_PYPI_INDEX" \
      --trusted-host "$METAX_PYPI_HOST"
  else
    python3 -m pip index versions "$LIST_PACKAGE" \
      -i "$METAX_PYPI_INDEX" \
      --trusted-host "$METAX_PYPI_HOST"
  fi
  exit 0
fi

uv venv "$UV_PROJECT_ENVIRONMENT"
uv sync --project "$ROOT_DIR/environments/metax" --no-dev --no-install-project "$@"

gpu_install_args=(pip install --reinstall)
if [[ "$METAX_INSTALL_GPU_DEPS" != "true" ]]; then
  gpu_install_args+=(--no-deps)
fi
gpu_install_args+=(
  --index-url "$METAX_PYPI_INDEX" \
  --trusted-host "$METAX_PYPI_HOST" \
)
for package_spec in "$METAX_TORCH_SPEC" "$METAX_TORCHAUDIO_SPEC" "$METAX_TORCHVISION_SPEC" "$METAX_VLLM_SPEC"; do
  [[ -n "$package_spec" ]] && gpu_install_args+=("$package_spec")
done

uv "${gpu_install_args[@]}"

if [[ "$METAX_INSTALL_GPU_DEPS" != "true" ]]; then
  echo "[sync_metax_env] Installed MetaX GPU stack with --no-deps to avoid PyPI/NVIDIA fallback."
  echo "[sync_metax_env] Set METAX_INSTALL_GPU_DEPS=true only if the MetaX index hosts all required GPU-stack dependencies."
fi

"$UV_PROJECT_ENVIRONMENT/bin/python" - <<'PY'
import importlib
import torch

print("torch", torch.__version__)
print("torch.cuda.is_available", torch.cuda.is_available())
for name in ("vllm",):
    module = importlib.import_module(name)
    print(name, getattr(module, "__version__", "unknown"))
PY