name: PyPI Release

on:
  push:
    tags:
      - 'v*'
  workflow_dispatch:
    inputs:
      dev_release:
        description: "Do a dev release or regular release"
        required: true
        default: "false"

permissions:
  contents: read

jobs:
  setup:
    runs-on: ubuntu-latest
    steps:
      - name: Set publishing variables
        run: echo "Publishing setup complete"

  build_documentation:
    if: github.repository == 'ml-explore/mlx'
    runs-on: [self-hosted, macos]
    steps:
      - uses: actions/checkout@v5
      - uses: ./.github/actions/build-docs
    
  deploy_documentation:
    needs: build_documentation
    permissions:
      pages: write
      id-token: write
    runs-on: ubuntu-latest
    environment:
      name: github-pages
      url: ${{ steps.deployment.outputs.page_url }}
    steps:
      - name: Deploy to GitHub Pages
        id: deployment
        uses: actions/deploy-pages@v4

  build_linux_release:
    if: github.repository == 'ml-explore/mlx'
    strategy:
      matrix:
        python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
        arch: ['x86_64', 'aarch64']
    runs-on: ${{ matrix.arch == 'x86_64' && 'ubuntu-22.04' || 'ubuntu-22.04-arm' }}
    env:
      PYPI_RELEASE: 1
      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
    steps:
      - uses: actions/checkout@v5
      - uses: ./.github/actions/setup-linux
        with:
          python-version: ${{ matrix.python_version }}
      - uses: ./.github/actions/build-linux-release
        with:
          build-backend: ${{ matrix.python-version == '3.10' }}
          arch: ${{ matrix.arch }}
      - name: Upload MLX artifacts
        uses: actions/upload-artifact@v5
        with:
          overwrite: true
          name: linux-wheels-${{ matrix.python_version }}
          path: wheelhouse/mlx-*.whl
      - name: Upload CPU artifacts
        if: matrix.python_version == '3.10'
        uses: actions/upload-artifact@v5
        with:
          overwrite: true
          name: mlx-cpu
          path: wheelhouse/mlx_cpu-*.whl
  
  build_mac_release:
    if: github.repository == 'ml-explore/mlx'
    strategy:
      matrix:
        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
    runs-on: [self-hosted, macos]
    env:
      PYPI_RELEASE: 1
      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}

    steps:
      - uses: actions/checkout@v5
      - uses: ./.github/actions/setup-macos
        with:
          python-version: ${{ matrix.python-version }}

      - name: Install dependencies
        shell: bash -l {0}
        run: |
          pip install --upgrade pip
          pip install cmake setuptools nanobind==2.4.0
          pip install -e . -v
      - name: Generate package stubs
        shell: bash -l {0}
        run: |
          pip install typing_extensions
          python setup.py generate_stubs
      - name: Build macOS 14 package
        uses: ./.github/actions/build-macos-release
        with:
          macos-target: 14.0
          build-backend: ${{ matrix.python-version == '3.10' }}
      - name: Build macOS 15 package
        uses: ./.github/actions/build-macos-release
        with:
          macos-target: 15.0
          build-backend: ${{ matrix.python-version == '3.10' }}
      - name: Upload MLX artifacts
        uses: actions/upload-artifact@v5
        with:
          overwrite: true
          name: mac-wheels-${{ matrix.python-version }}
          path: dist/mlx-*.whl
      - name: Upload Metal artifacts
        if: matrix.python-version == '3.10'
        uses: actions/upload-artifact@v5
        with:
          overwrite: true
          name: mlx-metal
          path: dist/mlx_metal-*.whl

  build_cuda_release:
    if: github.repository == 'ml-explore/mlx'
    runs-on: ubuntu-22-large
    env:
      PYPI_RELEASE: 1
      DEV_RELEASE: ${{ github.event.inputs.dev_release == 'true' && 1 || 0 }}
    steps:
      - uses: actions/checkout@v5
      - uses: ./.github/actions/setup-linux
        with:
          toolkit: 'cuda-12.9'
      - name: Build Python package
        uses: ./.github/actions/build-cuda-release
        with:
          toolkit: 'cuda-12.9'
      - name: Upload artifacts
        uses: actions/upload-artifact@v5
        with:
          overwrite: true
          name: mlx-cuda
          path: wheelhouse/mlx_cuda-*.whl

  pypi-publish:
    name: Upload release to PyPI
    runs-on: ubuntu-latest
    needs: [setup, build_linux_release, build_mac_release]
    permissions:
      id-token: write
    environment:
      name: pypi
      url: https://pypi.org/p/mlx
    steps:
      - uses: actions/download-artifact@v6
        with:
          pattern: linux-wheels-*
          merge-multiple: true
          path: dist
      - uses: actions/download-artifact@v6
        with:
          pattern: mac-wheels-*
          merge-multiple: true
          path: dist
      - name: Display structure of downloaded files
        run: ls -R dist
      - name: Publish package distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/
  
  pypi-publish-cuda:
    name: Upload CUDA release to PyPI
    runs-on: ubuntu-latest
    needs: [setup, build_cuda_release]
    permissions:
      id-token: write
    environment:
      name: pypi
      url: https://pypi.org/p/mlx-cuda
    steps:
      - uses: actions/download-artifact@v6
        with:
          name: mlx-cuda
          path: dist
      - name: Display structure of downloaded files
        run: ls -R dist
      - name: Publish package distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/

  pypi-publish-cpu:
    name: Upload CPU release to PyPI
    runs-on: ubuntu-latest
    needs: [setup, build_linux_release]
    permissions:
      id-token: write
    environment:
      name: pypi
      url: https://pypi.org/p/mlx-cpu
    steps:
      - uses: actions/download-artifact@v6
        with:
          name: mlx-cpu
          path: dist
      - name: Display structure of downloaded files
        run: ls -R dist
      - name: Publish package distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/

  pypi-publish-metal:
    name: Upload Metal release to PyPI
    runs-on: ubuntu-latest
    needs: [setup, build_mac_release]
    permissions:
      id-token: write
    environment:
      name: pypi
      url: https://pypi.org/p/mlx-metal
    steps:
      - uses: actions/download-artifact@v6
        with:
          name: mlx-metal
          path: dist
      - name: Display structure of downloaded files
        run: ls -R dist
      - name: Publish package distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          repository-url: https://upload.pypi.org/legacy/