#!/usr/bin/env bash
# 在「已激活的 fishserver」中，将 torch/torchvision 换成指定 CUDA 索引的构建。
# 若从未单独装过 fishserver，请直接跑 bootstrap_fishserver.sh（已内含 GPU 轮子）。
set -euo pipefail

TORCH_INDEX="${TORCH_INDEX:-https://download.pytorch.org/whl/cu124}"
export PIP_DEFAULT_TIMEOUT="${PIP_DEFAULT_TIMEOUT:-120}"

# 必须用「目标环境」里的 python：仅靠 `command -v python` 在激活失败时仍会指向 base。
if [[ -n "${FISHSERVER_PY:-}" ]]; then
  PY="$FISHSERVER_PY"
elif [[ -n "${CONDA_PREFIX:-}" && -x "${CONDA_PREFIX}/bin/python" ]]; then
  PY="${CONDA_PREFIX}/bin/python"
else
  PY="$(command -v python)"
fi

echo "Using PyTorch wheel index: $TORCH_INDEX"
echo "conda env: ${CONDA_DEFAULT_ENV:-"(unset)"}  CONDA_PREFIX: ${CONDA_PREFIX:-"(unset)"}"
echo "Using python: $PY ($("$PY" --version 2>&1))"

if [[ -z "${FISHSERVER_PY:-}" && "${CONDA_DEFAULT_ENV:-}" == "base" ]]; then
  echo "" >&2
  echo "WARNING: 当前 conda 仍是 base，刚才那类输出会把 torch 装进 base（常见为 py3.12），不是 fishserver（py3.11）。" >&2
  echo "  请先: conda activate fishserver   再跑本脚本，或:" >&2
  echo "  FISHSERVER_PY=\$HOME/miniforge3/envs/fishserver/bin/python bash packaging/patch_cuda_torch.sh" >&2
  echo "" >&2
fi

# 已装过 PyPI/cu130 的 torch 时，不加 uninstall / force 会一律 "already satisfied"，版本不会变。
# 全局 pip.conf 里的 extra-index-url（如 pypi.ngc.nvidia.com）也会参与解析，用空配置避开。
echo "Removing any existing torch / torchvision …"
"$PY" -m pip uninstall -y torch torchvision 2>/dev/null || true

echo "Installing torch / torchvision from $TORCH_INDEX only …"
PIP_CONFIG_FILE=/dev/null PIP_EXTRA_INDEX_URL="" \
  "$PY" -m pip install --no-cache-dir torch torchvision --index-url "$TORCH_INDEX"

"$PY" -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())"