name: Publish FastVideo Kernel to PyPI on Version Change

on:
  push:
    branches:
      - main
    paths:
      - "fastvideo-kernel/pyproject.toml"
  workflow_dispatch:

jobs:
  check-version-change:
    runs-on: ubuntu-latest
    outputs:
      version-changed: ${{ steps.check-version.outputs.changed }}
      new-version: ${{ steps.check-version.outputs.new-version }}
    steps:
      - name: Checkout code
        uses: actions/checkout@v4
        with:
          fetch-depth: 2

      - name: Check if version changed
        id: check-version
        run: |
          cd fastvideo-kernel
          # Get current commit's version from pyproject.toml
          # Use ^ to match start of line to avoid matching minimum-version
          NEW_VERSION=$(grep -oP '^version\s*=\s*"\K[^"]+' pyproject.toml)
          echo "New version: $NEW_VERSION"
          
          # Get previous version from git history
          # Note: git show expects path relative to repo root
          OLD_VERSION=$(git show HEAD~1:fastvideo-kernel/pyproject.toml | grep -oP '^version\s*=\s*"\K[^"]+' || echo "0.0.0")
          echo "Old version: $OLD_VERSION"
          
          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
            echo "changed=true" >> "$GITHUB_OUTPUT"
            echo "new-version=$NEW_VERSION" >> "$GITHUB_OUTPUT"
          else
            echo "Version did not change"
            echo "changed=false" >> "$GITHUB_OUTPUT"
          fi

  build_wheels:
    name: Build Wheel
    needs: check-version-change
    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
    runs-on: ${{ matrix.platform.os }}

    strategy:
      fail-fast: false
      matrix:
          python-version: ['3.12']
          torch-cuda:
            # torch 2.12 dropped cu128; cu130 (CUDA 13) is the default wheel, cu126 covers older drivers.
            - torch-version: '2.12.0'
              cuda-version: '12.6.3'
              torch-cuda-short: 'cu126'
            - torch-version: '2.12.0'
              cuda-version: '13.0.0'
              torch-cuda-short: 'cu130'
          platform:
            # x86_64 builds the full cu126 + cu130 set (cu130 ships the consumer
            # Blackwell sm_120a FP4 kernels).
            - os: ubuntu-22.04
              arch: x86_64
              wheel-plat: manylinux_2_35_x86_64
          # aarch64 is Blackwell (GB200 sm_100a + DGX Spark / consumer sm_120a), not
          # Hopper, and Blackwell needs CUDA >= 12.8 — so only the cu130 leg applies.
          # Added via include so x86 keeps cu126 + cu130 while aarch64 stays cu130-only.
          include:
            - python-version: '3.12'
              torch-cuda:
                torch-version: '2.12.0'
                cuda-version: '13.0.0'
                torch-cuda-short: 'cu130'
              platform:
                os: ubuntu-22.04-arm
                arch: aarch64
                wheel-plat: manylinux_2_35_aarch64

    steps:
      - name: Free up disk space
        run: |
          echo "Initial disk space:"
          df -h
          
          # Remove large directories
          sudo rm -rf /usr/share/dotnet
          sudo rm -rf /usr/local/lib/android
          sudo rm -rf /opt/ghc
          sudo rm -rf /usr/local/share/boost
          sudo rm -rf /usr/share/swift
          sudo rm -rf /usr/local/lib/node_modules
          sudo rm -rf /usr/local/share/powershell
          sudo rm -rf /usr/share/rust
          sudo rm -rf /usr/local/.ghcup
          
          # Remove cached files
          sudo rm -rf /var/lib/apt/lists/*
          sudo rm -rf /var/cache/apt/archives/*
          
          echo "Disk space after cleanup:"
          df -h

      - name: Checkout
        uses: actions/checkout@v4

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

      - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
        uses: Jimver/cuda-toolkit@v0.2.35
        id: cuda-toolkit
        with:
          cuda: ${{ matrix.torch-cuda.cuda-version }}
          linux-local-args: '["--toolkit"]'
          method: 'network'

      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
        run: |
          sudo apt update
          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11

          # Allow Git to Access Safe Directory
          git config --global --add safe.directory /__w/FastVideo/FastVideo

          # Set CUDA environment variables
          export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }}
          export PATH=${CUDA_HOME}/bin:${PATH}
          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
          
          # Verify installation
          gcc --version
          g++ --version
          clang-11 --version
          nvcc --version

      - name: Install uv
        uses: astral-sh/setup-uv@v3

      - name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }}
        run: |
          uv pip install --system typing-extensions==4.12.2
          uv pip install --system --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}}
          nvcc --version
          python --version
          python -c "import torch; print('PyTorch:', torch.__version__)"
          python -c "import torch; print('CUDA:', torch.version.cuda)"
          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"

      - name: Build wheel
        run: |
          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH

          uv pip install --system setuptools ninja packaging wheel triton scikit-build-core cmake build

          cd fastvideo-kernel
          git submodule update --init --recursive  # Ensure ThunderKittens submodule is initialized
          # Release builds run on GPU-less runners, so set kernels + arch explicitly:
          #   * aarch64 = Blackwell (GB200 sm_100a + DGX Spark/consumer sm_120a), NOT
          #     Hopper, so TK (sm_90a wgmma) is OFF. The C++ FP4 (attn_qat_infer, SM120)
          #     covers sm_120a; turbodiffusion covers sm_100a+sm_120a. The sm_100 FP4
          #     forward is the FA4 CuTe DSL path in the fastvideo package (PR #1221),
          #     JIT-compiled at runtime — not built into this wheel.
          #   * x86_64 cu130 = Hopper TK + consumer Blackwell sm_120a FP4.
          #   * x86_64 cu126 = Hopper TK only (older drivers; CUDA < 12.8 has no FP4).
          # The per-arch split in CMakeLists pins the FP4 targets to sm_120a and builds
          # the main extension for the full arch list. CMAKE_BUILD_PARALLEL_LEVEL caps
          # Ninja so heavy CUTLASS/TK template TUs don't OOM the 16 GB runner (exit 143).
          if [ "${{ matrix.platform.arch }}" = "aarch64" ]; then
            export TORCH_CUDA_ARCH_LIST="10.0a;12.0a"
            export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON"
            export CMAKE_BUILD_PARALLEL_LEVEL=1
          elif [ "${{ matrix.torch-cuda.torch-cuda-short }}" = "cu130" ]; then
            export TORCH_CUDA_ARCH_LIST="9.0a;12.0a"
            export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DFASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
            # A single FP4 TU (attn_qat_infer) can use ~8-12 GB on its own, so serialize.
            export CMAKE_BUILD_PARALLEL_LEVEL=1
          else
            export TORCH_CUDA_ARCH_LIST="9.0a"
            export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a"
            # Hopper-only: the two TK TUs are the heavy ones (single-arch) and fit side
            # by side; -j4 overlaps the light TUs without OOMing (~12-14 GB peak).
            export CMAKE_BUILD_PARALLEL_LEVEL=4
          fi

          # Build standard wheel (no local version suffix) for PyPI
          python -m build --wheel --outdir dist

          # Fix the wheel to be manylinux compliant
          uv pip install --system auditwheel
          # Point auditwheel at torch libs, but do not vendor them into the wheel.
          TORCH_LIB_DIR=$(python - <<'PY'
          import os
          import torch

          print(os.path.join(os.path.dirname(torch.__file__), "lib"))
          PY
          )
          export LD_LIBRARY_PATH="${TORCH_LIB_DIR}:${LD_LIBRARY_PATH}"
          # Target manylinux_2_35 (Ubuntu 22.04 native), per-arch plat tag.
          auditwheel repair dist/*.whl --plat ${{ matrix.platform.wheel-plat }} -w fixed_dist \
            --exclude libtorch_cuda.so \
            --exclude libtorch_cpu.so \
            --exclude libtorch.so \
            --exclude libc10.so \
            --exclude libc10_cuda.so \
            --exclude libtorch_python.so
          # Move fixed wheels back to dist for upload consistency
          rm dist/*.whl
          mv fixed_dist/*.whl dist/

      - name: Upload wheel artifact
        # Upload every matrix leg as a distinct artifact (for inspection / GitHub releases).
        # The publish job below selects which CUDA build is pushed to PyPI.
        uses: actions/upload-artifact@v4
        with:
          name: fastvideo_kernel-py${{ matrix.python-version }}-${{ matrix.torch-cuda.torch-cuda-short }}-${{ matrix.platform.arch }}-torch${{ matrix.torch-cuda.torch-version }}
          path: fastvideo-kernel/dist/*.whl
          retention-days: 90

  publish_package:
    name: Publish package
    needs: [build_wheels, check-version-change]
    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
    runs-on: ubuntu-22.04
    permissions:
      id-token: write  # Needed for OIDC Trusted Publishing

    steps:
      - uses: actions/checkout@v4

      - uses: actions/setup-python@v5
        with:
          python-version: '3.12'

      - name: Download PyPI wheels
        # Publish the cu130 (CUDA 13) wheels to PyPI for both architectures:
        #   x86_64  — Hopper sm_90a TK + consumer Blackwell sm_120a FP4
        #   aarch64 — Blackwell: turbodiffusion (sm_100a/sm_120a) + C++ FP4 (sm_120a);
        #             no TK (Hopper). sm_100 FP4 forward is the FA4 CuTe DSL path in the
        #             fastvideo package (#1221), shipped/JIT separately.
        # The x86_64 cu126 wheel stays available as a build artifact / GitHub-release asset.
        uses: actions/download-artifact@v4
        with:
          path: fastvideo-kernel/dist/
          pattern: 'fastvideo_kernel-py*-cu130-*'
          merge-multiple: true

      - name: Install uv
        uses: astral-sh/setup-uv@v3

      - name: Build source distribution
        run: |
          uv pip install --system build scikit-build-core cmake ninja

          cd fastvideo-kernel
          # We don't need full CUDA/Torch to just package the source (sdist)
          python -m build --sdist --outdir dist

      - name: Publish release distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
          packages-dir: fastvideo-kernel/dist/