name: Test Suite

on:
  push:
    branches: [ main, master, develop ]
    paths:
      - 'src/**'
      - 'tests/**'
      - 'pyproject.toml'
      - '.github/workflows/test.yml'
  pull_request:
    branches: [ main, master, develop ]
    paths:
      - 'src/**'
      - 'tests/**'
      - 'pyproject.toml'
      - '.github/workflows/test.yml'
  workflow_dispatch:

jobs:
  test-jax:
    name: Test JAX/Flax
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.10", "3.11", "3.12"]

    steps:
    - uses: actions/checkout@v4
    
    - name: Free disk space
      run: |
        sudo rm -rf /usr/share/dotnet
        sudo rm -rf /opt/ghc
        sudo rm -rf /usr/local/share/boost
        df -h
    
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
        cache: 'pip'
    
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        pip install pytest pytest-cov
        pip install -e ".[dev]"
        pip install "jax==0.9.1" "jaxlib==0.9.1" "flax==0.12.5" optax
    
    - name: Clear pycache
      run: |
        find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
        find . -name "*.pyc" -delete 2>/dev/null || true
    
    - name: Run JAX tests
      run: |
        pytest tests/test_nnx/ tests/test_linen/ tests/integration/ -v --cov=nmn --cov-report=xml
      env:
        JAX_PLATFORM_NAME: cpu
    
    - name: Upload coverage
      uses: codecov/codecov-action@v4
      if: matrix.python-version == '3.11'
      with:
        file: ./coverage.xml
        flags: jax
        token: ${{ secrets.CODECOV_TOKEN }}
      continue-on-error: true

  test-torch:
    name: Test PyTorch
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.10", "3.11", "3.12"]

    steps:
    - uses: actions/checkout@v4
    
    - name: Free disk space
      run: |
        sudo rm -rf /usr/share/dotnet
        sudo rm -rf /opt/ghc
        sudo rm -rf /usr/local/share/boost
        df -h
    
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
        cache: 'pip'
    
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        pip install pytest pytest-cov
        pip install -e ".[dev]"
        pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
    
    - name: Clear pycache
      run: |
        find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
        find . -name "*.pyc" -delete 2>/dev/null || true
    
    - name: Run PyTorch tests
      run: |
        pytest tests/test_torch/ -v --cov=nmn --cov-report=xml
    
    - name: Upload coverage
      uses: codecov/codecov-action@v4
      if: matrix.python-version == '3.11'
      with:
        file: ./coverage.xml
        flags: torch
        token: ${{ secrets.CODECOV_TOKEN }}
      continue-on-error: true

  test-keras:
    name: Test Keras/TF
    runs-on: ubuntu-latest
    strategy:
      fail-fast: false
      matrix:
        python-version: ["3.10", "3.11", "3.12"]

    steps:
    - uses: actions/checkout@v4
    
    - name: Free disk space
      run: |
        sudo rm -rf /usr/share/dotnet
        sudo rm -rf /opt/ghc
        sudo rm -rf /usr/local/share/boost
        sudo rm -rf "$AGENT_TOOLSDIRECTORY"
        df -h
    
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
        cache: 'pip'
    
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        pip install pytest pytest-cov
        pip install -e ".[dev]"
        pip install tensorflow
    
    - name: Clear pycache
      run: |
        find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
        find . -name "*.pyc" -delete 2>/dev/null || true
    
    - name: Run Keras/TF tests
      run: |
        pytest tests/test_keras/ tests/test_tf/ -v --cov=nmn --cov-report=xml
    
    - name: Upload coverage
      uses: codecov/codecov-action@v4
      if: matrix.python-version == '3.11'
      with:
        file: ./coverage.xml
        flags: keras
        token: ${{ secrets.CODECOV_TOKEN }}
      continue-on-error: true

  lint:
    name: Lint
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v4

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

    - name: Install lint tools
      run: |
        pip install flake8 black isort

    - name: Lint with flake8 (errors only — fails build)
      run: |
        flake8 src/nmn --count --select=E9,F63,F7,F82 --show-source --statistics

    - name: Lint with flake8 (style warnings — informational)
      run: |
        flake8 src/nmn --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics

    - name: Check formatting with black
      run: black --check --diff src/nmn tests
      continue-on-error: true  # style only — warn, don't fail

    - name: Check imports with isort
      run: isort --check-only --diff src/nmn tests
      continue-on-error: true  # style only — warn, don't fail

  mypy:
    name: Type Check (MyPy)
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v4

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

    - name: Install dependencies
      run: |
        pip install mypy
        pip install -e ".[dev]"
        pip install torch --index-url https://download.pytorch.org/whl/cpu
        pip install "jax==0.9.1" "jaxlib==0.9.1" "flax==0.12.5"

    - name: Run MyPy on torch module
      run: mypy src/nmn/torch --ignore-missing-imports --no-error-summary
      continue-on-error: true  # informational for now — tighten over time

    - name: Run MyPy on nnx module
      run: mypy src/nmn/nnx --ignore-missing-imports --no-error-summary
      continue-on-error: true