name: ~Build wheel template

on:
  workflow_call:
    inputs:
      runs-on:
        description: "The runner to use for the build"
        required: true
        type: string
      container-image:
        description: "Container image"
        required: true
        type: string
      upload-to-release:
        description: "Upload wheel to this release"
        required: false
        type: boolean
        default: false
      release-version:
        description: "Upload wheel to this release"
        required: false
        type: string

defaults:
  run:
    shell: bash -x -e -u -o pipefail {0}

jobs:
  build-wheel:
    runs-on: ${{ inputs.runs-on }}
    name: Build wheel (${{ inputs.container-image }})
    steps:
      - name: Move /var/lib/docker/
        run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker"

      - name: Maximize build space
        uses: easimon/maximize-build-space@master
        with:
          root-reserve-mb: 5120
          temp-reserve-mb: 32
          swap-size-mb: 10240
          remove-dotnet: "true"
          remove-android: "true"
          remove-haskell: "true"
          remove-codeql: "true"
          build-mount-path: "/var/lib/docker/"

      - name: Restore /var/lib/docker/
        run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker"

      - name: Checkout source
        uses: actions/checkout@v4
        with:
          ref: ${{ inputs.release-version }}
          submodules: recursive

      - name: Checkout build scripts
        uses: actions/checkout@v4
        with:
          path: build-scripts/

      - name: Build
        run: |
          echo "Free space:"
          df -h

      - name: Pull the container
        run: docker pull ${{ inputs.container-image }}

      - name: Set CUDA and PyTorch versions
        run: |
          cat <<'EOF' >> script.sh
          #!/bin/bash

          set -eoxu pipefail

          echo "MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
          echo "MATRIX_TORCH_VERSION=$NVIDIA_PYTORCH_VERSION" >> $GITHUB_ENV
          echo "WHEEL_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'})" >> $GITHUB_ENV
          echo "MATRIX_PYTHON_VERSION=$(python -c "import sys; print('{}.{}'.format(sys.version_info[0], sys.version_info[1]))" | awk -F \. {'print $1 $2'})" >> $GITHUB_ENV
          echo "CXX11_ABI=$(python -c 'import torch; print(str(torch._C._GLIBCXX_USE_CXX11_ABI).upper())')" >> $GITHUB_ENV

          cat $GITHUB_ENV
          EOF

          docker run \
            --rm \
            --shm-size=64g \
            --workdir /workspace \
            --volume $(pwd):/workspace \
            --volume $GITHUB_ENV:$GITHUB_ENV \
            -e GITHUB_ENV=$GITHUB_ENV \
            ${{ inputs.container-image }} bash /workspace/script.sh

      - name: Build wheel
        id: build_wheel
        env:
          CXX11_ABI: ${{ env.CXX11_ABI }}
          MATRIX_TORCH_VERSION: ${{ env.MATRIX_TORCH_VERSION}}
          WHEEL_CUDA_VERSION: ${{ env.WHEEL_CUDA_VERSION }}
          MATRIX_PYTHON_VERSION: ${{ env.MATRIX_PYTHON_VERSION }}
        run: |
          EXIT_CODE=$(docker run \
            --rm \
            --shm-size=64g \
            --workdir /workspace \
            --volume $(pwd):/workspace \
            --volume $GITHUB_ENV:$GITHUB_ENV \
            -e PIP_CONSTRAINT= \
            -e GITHUB_ENV=$GITHUB_ENV \
            -e CXX11_ABI=$CXX11_ABI \
            -e MATRIX_TORCH_VERSION=$MATRIX_TORCH_VERSION \
            -e WHEEL_CUDA_VERSION=$WHEEL_CUDA_VERSION \
            -e MATRIX_PYTHON_VERSION=$MATRIX_PYTHON_VERSION \
            ${{ inputs.container-image }} bash /workspace/build-scripts/.github/scripts/build.sh | tail -n 1)

      - name: Log Built Wheels
        run: |
          ls dist

      - name: Get Release with tag
        id: get_current_release
        uses: joutvhu/get-release@v1
        with:
          tag_name: ${{ inputs.release-version }}
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

      - name: Upload Release Asset
        id: upload_release_asset
        if: inputs.upload-to-release
        uses: actions/upload-release-asset@v1
        env:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
        with:
          upload_url: ${{ steps.get_current_release.outputs.upload_url }}
          asset_path: ./dist/${{env.wheel_name}}
          asset_name: ${{env.wheel_name}}
          asset_content_type: application/*