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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"

export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
export PIP_INDEX_URL="${PIP_INDEX_URL:-https://mirrors.aliyun.com/pypi/simple/}"
export UV_DEFAULT_INDEX="${UV_DEFAULT_INDEX:-${PIP_INDEX_URL}}"
export UV_INDEX_URL="${UV_INDEX_URL:-${PIP_INDEX_URL}}"

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip setuptools wheel uv

UV=(python -m uv pip)
"${UV[@]}" install -U pip setuptools wheel

if [[ "${INSTALL_SWIFT_ALL:-0}" == "1" ]]; then
  "${UV[@]}" install -e 'third_party/modelscope-swift[all]' --torch-backend=auto
else
  "${UV[@]}" install -e third_party/modelscope-swift --torch-backend=auto
fi

"${UV[@]}" install -e .

if [[ "${INSTALL_TRAINING_EXTRAS:-1}" == "1" ]]; then
  "${UV[@]}" install -U 'deepspeed<0.19' liger-kernel tensorboard nvitop
fi

if [[ -n "${EXTRA_UV_PACKAGES:-}" ]]; then
  # shellcheck disable=SC2206
  EXTRA_PACKAGES=(${EXTRA_UV_PACKAGES})
  "${UV[@]}" install -U "${EXTRA_PACKAGES[@]}"
fi

mkdir -p data/raw data/processed models outputs runs logs
python - <<'PY'
import swift, sys
print("python", sys.version)
print("swift", getattr(swift, "__version__", "unknown"))
PY