#!/bin/bash
set -e

# 🧾 Config
# APEX_COMMIT=f3a960f80244cf9e80558ab30f7f7e8cbf03c0a0  # Known stable commit

echo "🧹 Cleaning any previous apex build..."
rm -rf apex
rm -rf *.egg-info build dist

echo "📥 Cloning NVIDIA/apex..."
git clone https://github.com/NVIDIA/apex.git
cd apex
# git checkout $APEX_COMMIT

echo "🛠 Installing build dependencies..."
sudo apt-get update
sudo apt-get install -y \
  build-essential \
  ninja-build \
  python3-dev \
  libffi-dev \
  libncurses5-dev \
  libncursesw5-dev \
  libreadline-dev \
  libssl-dev \
  libsqlite3-dev \
  zlib1g-dev \
  libbz2-dev \
  liblzma-dev \
  git

echo "🐍 Upgrading pip and wheel..."
pip install -U pip setuptools wheel

echo "🧪 Verifying PyTorch + CUDA availability..."
python -c "import torch; print('PyTorch:', torch.__version__, '| CUDA:', torch.version.cuda)"

echo "🔧 Building Apex with CUDA and C++ extensions..."
python setup.py bdist_wheel --cuda_ext --cpp_ext

echo "✅ Done! Built wheel:"
ls -lh dist/*.whl

cd ..