name: Build Wheels

on: [workflow_call, workflow_dispatch]
jobs:
  build_sdist:
    name: Build source distribution and no binary wheel
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
        with:
          submodules: recursive
      - name: Strip unsupported tags in README
        run: |
          sed -i '/<!-- pypi-strip -->/,/<!-- \/pypi-strip -->/d' README.md
      - name: Build sdist
        run: BUILD_NO_CUDA=1 pipx run build --sdist
      - name: Build wheel with no binaries
        run: BUILD_NO_CUDA=1 python setup.py bdist_wheel --dist-dir=dist
      - uses: actions/upload-artifact@v4
        with:
          name: pypi_packages
          path: dist/*.tar.gz

  build_wheels:
    runs-on: ${{ matrix.os }}
    environment: production

    strategy:
      fail-fast: false
      matrix:
        os: [ubuntu-22.04, windows-2022]
        python-version: ['3.10']
        # Last 3 minor PyTorch versions aligned with pytorch.org
        torch-version: ['2.8.0', '2.9.1', '2.10.0']
        # CUDA versions aligned with PyTorch build matrix
        cuda-version: ['cu126', 'cu128', 'cu129', 'cu130']
        exclude:
          # PyTorch 2.8.0 does not ship cu130 wheels
          - torch-version: '2.8.0'
            cuda-version: 'cu130'
          # PyTorch 2.9.1 does not ship cu129 wheels
          - torch-version: '2.9.1'
            cuda-version: 'cu129'
          # PyTorch 2.10.0 does not ship cu129 wheels
          - torch-version: '2.10.0'
            cuda-version: 'cu129'
          # CUDA 12.6 is incompatible with MSVC >= 14.44 (VS 2022);
          # the MSVC STL requires CUDA 12.8+
          - os: windows-2022
            cuda-version: 'cu126'

    steps:
      - name: Checkout repository
        uses: actions/checkout@v4
        with:
          submodules: recursive

      - name: Set up Python ${{ matrix.python-version }}
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Free up disk space
        if: ${{ runner.os == 'Linux' }}
        run: |
          echo "Disk space before cleanup:"
          df -h
          sudo rm -rf /usr/share/dotnet
          echo "Disk space after cleanup:"
          df -h
        shell: bash

      - name: Install CUDA ${{ matrix.cuda-version }}
        if: ${{ matrix.cuda-version != 'cpu' }}
        run: |
          bash .github/workflows/cuda/${{ runner.os }}.sh ${{ matrix.cuda-version }}
        shell: bash

      - name: Install PyTorch ${{ matrix.torch-version }}+${{ matrix.cuda-version }}
        run: |
          pip install torch==${{ matrix.torch-version }} --extra-index-url https://download.pytorch.org/whl/${{ matrix.cuda-version }}
          python -c "import torch; print('PyTorch:', torch.__version__)"
          python -c "import torch; print('CUDA:', torch.version.cuda)"
          python -c "import torch; print('CUDA Available:', torch.cuda.is_available())"
        shell: bash

      - name: Set version
        if: ${{ runner.os != 'macOS' }}
        run: |
          VERSION=`sed -n 's/^__version__ = "\(.*\)"/\1/p' gsplat/version.py`
          TORCH_VERSION=`echo "pt${{ matrix.torch-version }}" | sed "s/..$//" | sed "s/\.//g"`
          CUDA_VERSION=`echo ${{ matrix.cuda-version }}`
          echo "New version name: $VERSION+$TORCH_VERSION$CUDA_VERSION"
          sed -i "s/$VERSION/$VERSION+$TORCH_VERSION$CUDA_VERSION/" gsplat/version.py
        shell: bash

      - name: Upgrade pip
        run: |
          # https://github.com/pypa/setuptools/issues/5174, pkg_resources is removed in setuptools 82 but we need it for older torch versions
          pip install --upgrade "setuptools<82"
          pip install ninja
        shell: bash

      - name: Install main package for CPU
        if: ${{ matrix.cuda-version == 'cpu' }}
        run: |
          BUILD_NO_CUDA=1 pip install .
        shell: bash

      - name: Build wheel
        env:
          # Step-level env so Windows Python sees CUDA_HOME (bash export does not propagate to python.exe)
          CUDA_HOME: ${{ runner.os == 'Windows' && (matrix.cuda-version == 'cu130' && 'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0' || matrix.cuda-version == 'cu129' && 'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.9' || matrix.cuda-version == 'cu128' && 'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8' || 'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6') || '' }}
        run: |
          pip install wheel
          if [ "$RUNNER_OS" = "Windows" ]; then
            export PATH="${CUDA_HOME}/bin:${PATH}"
            export TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0"
          else
            source .github/workflows/cuda/${{ runner.os }}-env.sh ${{ matrix.cuda-version }}
          fi
          MAX_JOBS=2 python setup.py bdist_wheel --dist-dir=dist
        shell: bash

      - name: Test wheel
        run: |
          cd dist
          ls -lah
          pip install *.whl
          python -c "import gsplat; print('gsplat:', gsplat.__version__)"
          pip install "$(ls *.whl)[lidar]"
          python -c "import gsplat; print('gsplat:', gsplat.__version__); from scipy.spatial import cKDTree; print('scipy OK')"
          cd ..
        shell: bash

      - uses: actions/upload-artifact@v4
        with:
          # Include unique matrix values to avoid name collisions.
          name: compiled_wheels_python${{ matrix.python-version }}-${{ matrix.os }}-${{ matrix.torch-version }}-${{ matrix.cuda-version }}
          path: dist/*.whl
