name: Build AutoAWQ Wheels with CUDA

on:
  push:
    tags:
      - "v*"

jobs:
  release:
    # Retrieve tag and create release
    name: Create Release
    runs-on: ubuntu-latest
    outputs:
      upload_url: ${{ steps.create_release.outputs.upload_url }}
    steps:
      - name: Checkout
        uses: actions/checkout@v3

      - name: Extract branch info
        shell: bash
        run: |
          echo "release_tag=${GITHUB_REF#refs/*/}" >> $GITHUB_ENV

      - name: Create Release
        id: create_release
        uses: "actions/github-script@v6"
        env:
          RELEASE_TAG: ${{ env.release_tag }}
        with:
          github-token: "${{ secrets.GITHUB_TOKEN }}"
          script: |
            const script = require('.github/workflows/scripts/github_create_release.js')
            await script(github, context, core)

  build_wheels:
    name: Build AWQ
    runs-on: ${{ matrix.os }}
    needs: release
    
    strategy:
      matrix:
        os: [ubuntu-20.04, windows-latest]
        pyver: ["3.8", "3.9", "3.10", "3.11"]
        cuda: ["11.8.0", "12.1.1"]
    defaults:
      run:
        shell: pwsh
    env:
      CUDA_VERSION: ${{ matrix.cuda }}
    
    steps:
      - name: Free Disk Space
        uses: jlumbroso/free-disk-space@v1.3.0
        if: runner.os == 'Linux'
        with:
          tool-cache: false
          android: true
          dotnet: true
          haskell: true
          large-packages: false
          docker-images: true
          swap-storage: false
        
      - uses: actions/checkout@v3

      - uses: actions/setup-python@v3
        with:
          python-version: ${{ matrix.pyver }}

      - name: Setup Mamba
        uses: conda-incubator/setup-miniconda@v2.2.0
        with:
          activate-environment: "build"
          python-version: ${{ matrix.pyver }}
          miniforge-variant: Mambaforge
          miniforge-version: latest
          use-mamba: true
          add-pip-as-python-dependency: true
          auto-activate-base: false
          
      - name: Install Dependencies
        run: |
          # Install CUDA toolkit
          mamba install -y 'cuda' -c "nvidia/label/cuda-${env:CUDA_VERSION}"

          # Env variables
          $env:CUDA_PATH = $env:CONDA_PREFIX
          $env:CUDA_HOME = $env:CONDA_PREFIX
          
          # Install torch
          $cudaVersion = $env:CUDA_VERSION.Replace('.', '')
          $cudaVersionPytorch = $cudaVersion.Substring(0, $cudaVersion.Length - 1)
          if ([int]$cudaVersionPytorch -gt 118) { $pytorchVersion = "torch==2.1.0" } else {$pytorchVersion = "torch==2.0.1"}
          python -m pip install --upgrade --no-cache-dir $pytorchVersion+cu$cudaVersionPytorch --index-url https://download.pytorch.org/whl/cu$cudaVersionPytorch
          python -m pip install build setuptools wheel ninja

          # Print version information
          python --version
          python -c "import torch; print('PyTorch:', torch.__version__)"
          python -c "import torch; print('CUDA:', torch.version.cuda)"
          python -c "import os; print('CUDA_HOME:', os.getenv('CUDA_HOME', None))"
          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"

      - name: Build Wheel
        run: |
          $env:CUDA_PATH = $env:CONDA_PREFIX
          $env:CUDA_HOME = $env:CONDA_PREFIX

          python setup.py sdist bdist_wheel
      
      - name: Upload Assets
        uses: shogo82148/actions-upload-release-asset@v1
        with:
          upload_url: ${{ needs.release.outputs.upload_url }}
          asset_path: ./dist/*.whl