#!/bin/bash
# Build all CUDA extension packages from source for sm_89 (RTX 4090)
# Must run inside the activated venv with CUDA toolkit installed
#
# Usage: source venv/bin/activate && bash build_cuda_extensions.sh

set -e

export PATH="/usr/local/cuda/bin:$PATH"
export TORCH_CUDA_ARCH_LIST="8.9"

echo "=== Building CUDA extensions for sm_89 (RTX 4090) ==="
echo ""

# 1. utils3d (MoGe dependency — GitHub fork, NOT pip package)
echo "[1/5] utils3d (EasternJournalist fork)..."
pip install --quiet "git+https://github.com/EasternJournalist/utils3d.git@3fab839f0be9931dac7c8488eb0e1600c236e183"
echo "  ✓ done"

# 2. flex_gemm
echo "[2/5] flex_gemm..."
pip uninstall -y flex_gemm 2>/dev/null || true
pip install --quiet "git+https://github.com/JeffreyXiang/FlexGEMM.git" --no-build-isolation
echo "  ✓ done"

# 3. o_voxel (visualbruno fork — has tiled_flexible_dual_grid_to_mesh)
echo "[3/5] o_voxel..."
pip uninstall -y o_voxel 2>/dev/null || true
pip install --quiet "git+https://github.com/visualbruno/TRELLIS.2.git#subdirectory=o-voxel" --no-build-isolation
echo "  ✓ done"

# 4. cumesh (visualbruno fork — has remesh_narrow_band_dc_quad, needs Eigen)
echo "[4/5] cumesh..."
pip uninstall -y cumesh 2>/dev/null || true
CPLUS_INCLUDE_PATH="/usr/include/eigen3:$CPLUS_INCLUDE_PATH" \
  pip install --quiet "git+https://github.com/visualbruno/CuMesh.git" --no-build-isolation
echo "  ✓ done"

# 5. nvdiffrast
echo "[5/5] nvdiffrast..."
pip uninstall -y nvdiffrast 2>/dev/null || true
pip install --quiet "git+https://github.com/NVlabs/nvdiffrast.git" --no-build-isolation
echo "  ✓ done"

echo ""
echo "=== Verification ==="
python -c "
import utils3d; print('utils3d:', utils3d.__version__)
import flex_gemm; print('flex_gemm: OK')
import o_voxel.convert; assert hasattr(o_voxel.convert, 'tiled_flexible_dual_grid_to_mesh'); print('o_voxel: OK (tiled_flexible_dual_grid_to_mesh found)')
import cumesh.remeshing; assert hasattr(cumesh.remeshing, 'remesh_narrow_band_dc_quad'); print('cumesh: OK (remesh_narrow_band_dc_quad found)')
import nvdiffrast.torch; print('nvdiffrast: OK (RasterizeCudaContext available)')
"
echo "=== All CUDA extensions built successfully ==="