name: CI
on: [pull_request]

concurrency:
  group: "${{ github.workflow }}-${{ github.ref }}"
  cancel-in-progress: ${{ github.event_name == 'pull_request' }}

env:
  PYTHON_VERSION: '3.12'
  VENV_DIR: tilelang_ci

jobs:
  format-check:
    runs-on: [self-hosted, nvidia]

    permissions:
      contents: write

    steps:
    - name: Checkout repository
      uses: actions/checkout@v5
      with:
        fetch-depth: 0

    - name: Set up Python
      uses: actions/setup-python@v6
      with:
        python-version: ${{ env.PYTHON_VERSION }}

    - name: Ensure venv (local & persistent)
      run: |
        set -e
        REQS_HASH=$(sha256sum requirements-test.txt 2>/dev/null | awk '{print $1}' || echo "no_requirements")
        MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}"

        if [[ -f "$MARKER" ]] && [[ -f "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" ]]; then
          echo "venv exists and hash matches – reuse it"
        else
          echo "venv stale or missing – recreating"
          rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" "$MARKER"
          python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}"
          # shellcheck source=/dev/null
          source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
          python -m pip install --upgrade pip --no-user
          [[ -f requirements-test.txt ]] && \
            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-test.txt --no-user
          pip install flash_attn==2.5.8 --no-user --no-build-isolation
          touch "$MARKER"
        fi

    - name: Run format check
      run: |
        source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
        git submodule update --init --recursive
        mkdir -p build
        # run cmake to create the build directory with compile_commands.json
        cd build; cmake .. -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DUSE_CUDA=ON; cd ..
        if ! output=$(./format.sh 2>&1); then
          echo "------------------------------------"
          echo "message:"
          echo "$output"
          printf '%s\n' "$output" | grep "Please review and stage the changes."
          echo "------------------------------------"
          exit 1
        fi
        rm -rf build

  build-test-nvidia:
    runs-on: [self-hosted, nvidia]
    needs: format-check
    permissions:
      contents: read
    steps:
    - name: Checkout repository
      uses: actions/checkout@v5
      with:
        fetch-depth: 0
        repository: ${{ github.event.pull_request.head.repo.full_name }}
        ref: ${{ github.event.pull_request.head.ref }}

    - name: Set up Python
      uses: actions/setup-python@v6
      with:
        python-version: ${{ env.PYTHON_VERSION }}

    - name: Ensure venv (local & persistent)
      run: |
        set -e
        REQS_HASH=$(cat requirements-test.txt 2>/dev/null || true)
        MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}"

        if [[ -f "$MARKER" ]] && [[ -f "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" ]]; then
          echo "venv exists and hash matches – reuse it"
        else
          echo "venv stale or missing – recreating"
          rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" "$MARKER"
          python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}"
          source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
          python -m pip install --upgrade pip --no-user
          [[ -f requirements-test.txt ]] && \
            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-test.txt --no-user
          # flash attention usually requires no isolation build
          pip install flash_attn==2.5.8 --no-user --no-build-isolation
          pip install . --no-user
          touch "$MARKER"
        fi

    - name: Install project (wheel form)
      run: |
        source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
        pip install . --no-user -v

    - name: Run examples
      run: |
        source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
        cd examples
        unset PYTHONPATH
        python -m pytest -n 4 **/test*.py -v -r fE --durations=0

    - name: Run tests
      run: |
        source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
        cd testing/python
        unset PYTHONPATH
        python -m pytest -n 4 -v -r fE --durations=0 --timeout=3600