name: PR Test

on:
  push:
    branches: [main]
    paths:
      - "fastvideo/**/*.py"
      - ".github/workflows/pr-test.yml"
  pull_request:
    branches: [main]
    types: [opened, ready_for_review, synchronize, reopened]
    paths:
      - "fastvideo/**/*.py"
      - ".github/workflows/pr-test.yml"
  workflow_dispatch:
    inputs:
      custom_image:
        description: "Custom image from this repository (default: fastvideo-dev:latest)"
        required: false
        default: "fastvideo-dev:latest"
        type: string
      run_encoder_test:
        description: "Run encoder-test"
        required: false
        default: false
        type: boolean
      run_vae_test:
        description: "Run vae-test"
        required: false
        default: false
        type: boolean
      run_transformer_test:
        description: "Run transformer-test"
        required: false
        default: false
        type: boolean
      run_ssim_test:
        description: "Run ssim-test"
        required: false
        default: false
        type: boolean

env:
  PYTHONUNBUFFERED: "1"

concurrency:
  group: pr-test-${{ github.ref }}
  cancel-in-progress: true

jobs:
  pre-commit:
    uses: ./.github/workflows/pre-commit.yml

  change-filter:
    runs-on: ubuntu-latest
    needs: pre-commit
    if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }}
    outputs:
      encoder-test: ${{ steps.filter.outputs.encoder-test }}
      vae-test: ${{ steps.filter.outputs.vae-test }}
      transformer-test: ${{ steps.filter.outputs.transformer-test }}
    steps:
      - uses: actions/checkout@v4
      - uses: dorny/paths-filter@v3
        id: filter
        with:
          filters: |
            encoder-test:
              - 'fastvideo/v1/models/encoders/**'
              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/encoders/**'
            vae-test:
              - 'fastvideo/v1/models/vaes/**'
              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/vaes/**'
            transformer-test:
              - 'fastvideo/v1/models/dits/**'
              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/transformers/**'

  encoder-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.encoder-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_encoder_test == 'true')
    runs-on: ubuntu-latest
    environment: runpod-runners
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.10"

      - name: Set up SSH key
        run: |
          mkdir -p ~/.ssh
          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
          chmod 600 ~/.ssh/id_rsa
          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub

      - name: Install dependencies
        run: pip install requests

      - name: Run tests on RunPod
        env:
          JOB_ID: "encoder-test"
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        timeout-minutes: 30
        run: >-
          python .github/scripts/runpod_api.py 
          --gpu-type "NVIDIA A40"
          --gpu-count 1
          --volume-size 100
          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"

      - name: Terminate RunPod Instances
        if: ${{ always() }}
        env:
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
          JOB_ID: "encoder-test"
        run: python .github/scripts/runpod_cleanup.py

  vae-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true')
    runs-on: ubuntu-latest
    environment: runpod-runners
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.10"

      - name: Set up SSH key
        run: |
          mkdir -p ~/.ssh
          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
          chmod 600 ~/.ssh/id_rsa
          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub

      - name: Install dependencies
        run: pip install requests

      - name: Run tests on RunPod
        env:
          JOB_ID: "vae-test"
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        timeout-minutes: 30
        run: >-
          python .github/scripts/runpod_api.py 
          --gpu-type "NVIDIA A40"
          --gpu-count 1
          --volume-size 100
          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"

      - name: Terminate RunPod Instances
        if: ${{ always() }}
        env:
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
          JOB_ID: "vae-test"
        run: python .github/scripts/runpod_cleanup.py

  transformer-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
    runs-on: ubuntu-latest
    environment: runpod-runners
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.10"

      - name: Set up SSH key
        run: |
          mkdir -p ~/.ssh
          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
          chmod 600 ~/.ssh/id_rsa
          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub

      - name: Install dependencies
        run: pip install requests

      - name: Run tests on RunPod
        env:
          JOB_ID: "transformer-test"
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        timeout-minutes: 30
        run: >-
          python .github/scripts/runpod_api.py 
          --gpu-type "NVIDIA L40S"
          --gpu-count 1
          --volume-size 100
          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"

      - name: Terminate RunPod Instances
        if: ${{ always() }}
        env:
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
          JOB_ID: "transformer-test"
        run: python .github/scripts/runpod_cleanup.py
      
  ssim-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
    runs-on: ubuntu-latest
    environment: runpod-runners
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.10"

      - name: Set up SSH key
        run: |
          mkdir -p ~/.ssh
          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
          chmod 600 ~/.ssh/id_rsa
          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub

      - name: Install dependencies
        run: pip install requests

      - name: Run tests on RunPod
        env:
          JOB_ID: "ssim-test"
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        timeout-minutes: 45
        run: >-
          python .github/scripts/runpod_api.py 
          --gpu-type "NVIDIA A40"
          --gpu-count 2
          --disk-size 200
          --volume-size 200
          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"

      - name: Terminate RunPod Instances
        if: ${{ always() }}
        env:
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
          JOB_ID: "ssim-test"
        run: python .github/scripts/runpod_cleanup.py

  runpod-cleanup:
    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
    if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
    runs-on: ubuntu-latest
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.10"

      - name: Install dependencies
        run: pip install requests

      - name: Cleanup all RunPod instances
        env:
          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test"]' # JSON array of job IDs
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        run: python .github/scripts/runpod_cleanup.py