name: Build and Publish Prebuilts Artifacts

on:
  workflow_dispatch:
    inputs:
      rocm_version:
        description: 'ROCm version'
        required: true
        default: '6.2'
      torch_version:
        description: 'Torch version (e.g., rocm6.2)'
        required: true
        default: 'rocm6.2'

jobs:
  build:
    runs-on: ubuntu-latest

    steps:
      - name: Free Disk Space (Ubuntu)
        uses: jlumbroso/free-disk-space@main
        with:
          tool-cache: false

      # Step 1: Checkout the repository
      - name: Checkout repository
        uses: actions/checkout@v4

      # Step 2: Set up Docker Buildx
      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3

      # Step 3: Build the Docker image
      - name: Build Docker image
        run: |
          echo ${{ inputs.rocm_version }}
          echo ${{ inputs.torch_version }}
          
          docker build \
            --build-arg ROCM_VERSION=${{ inputs.rocm_version }} \
            --build-arg TORCH_VERSION=${{ inputs.torch_version }} \
            -t prebuilt-wheels ./prebuilts

      # Step 4: Create a container and run the script
      - name: Run Docker container and generate wheels
        run: |
          docker create --name prebuilt-container prebuilt-wheels
          docker start -a prebuilt-container

      # Step 5: Copy bitsandbytes wheel artifact to host
      - name: Copy bitsandbytes wheel to host
        run: |
          mkdir -p /tmp/bitsandbytes/
          docker cp prebuilt-container:/tmp/bitsandbytes/dist/ /tmp/bitsandbytes/

      # Step 6: Copy llama-cpp-python wheel artifact to host
      - name: Copy llama-cpp-python wheel to host
        run: |
          mkdir -p /tmp/llama-cpp-python/
          docker cp prebuilt-container:/tmp/llama-cpp-python/dist/ /tmp/llama-cpp-python/

      # Step 6: Copy xformers wheel artifact to host
      - name: Copy xformers wheel to host
        run: |
          mkdir -p /tmp/xformers/
          docker cp prebuilt-container:/tmp/xformers/dist/ /tmp/xformers/

      # Step 7: Upload bitsandbytes wheel artifact
      - name: Upload bitsandbytes wheel
        uses: actions/upload-artifact@v4
        with:
          name: bitsandbytes-wheels
          path: /tmp/bitsandbytes/dist/*.whl

      # Step 8: Upload llama-cpp-python wheel artifact
      - name: Upload llama-cpp-python wheel
        uses: actions/upload-artifact@v4
        with:
          name: llama-cpp-python-wheels
          path: /tmp/llama-cpp-python/dist/*.whl

      # Step 9: Upload xformers wheel artifact
      - name: Upload xformers wheel
        uses: actions/upload-artifact@v4
        with:
          name: xformers-wheels
          path: /tmp/xformers/dist/*.whl

      # Step 10: Cleanup Docker container
      - name: Cleanup
        run: |
          docker rm prebuilt-container