name: concrete-python release

on:
  workflow_dispatch:
  push:
    tags:
      - 'v[0-9]+.[0-9]+.[0-9]+*'
  schedule:
    # Nightly Release @ 3AM after each work day
    - cron: "0 3 * * 2-6"

env:
  DOCKER_IMAGE_TEST: ghcr.io/zama-ai/concrete-compiler
  CUDA_PATH: /usr/local/cuda-11.8
  ACTION_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
  SLACK_CHANNEL: ${{ secrets.SLACK_CHANNEL }}
  SLACK_USERNAME: ${{ secrets.BOT_USERNAME }}
  SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK }}
  RELEASE_TYPE: ${{ (github.event_name == 'push' && contains(github.ref, 'refs/tags/')) && 'public' || 'nightly' }}
  PYTEST_MAXPROCESSES: "100"

concurrency:
  group: concrete_python_release_${{ github.ref }}
  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}

jobs:
  check-duplicate-nightly:
    runs-on: ubuntu-latest
    steps:
      - name: Checkout
        if: ${{ env.RELEASE_TYPE == 'nightly' }}
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          submodules: recursive
          fetch-depth: 0
      - name: Avoid releasing nightly twice
        if: ${{ env.RELEASE_TYPE == 'nightly' }}
        run: |
          # will fail if last commit has a nightly tag
          git log -1 --format=%D | grep "tag: nightly" && exit 1
          # will succeed otherwise
          exit 0

  setup-instance:
    needs: check-duplicate-nightly
    runs-on: ubuntu-latest
    outputs:
      runner-name: ${{ steps.start-instance.outputs.label }}
    steps:
      - name: Start instance
        id: start-instance
        uses: zama-ai/slab-github-runner@f26b8d611b2e695158fb0a6980834f0612f65ef8 # v1.4.0
        with:
          mode: start
          github-token: ${{ secrets.SLAB_ACTION_TOKEN }}
          slab-url: ${{ secrets.SLAB_BASE_URL }}
          job-secret: ${{ secrets.JOB_SECRET }}
          backend: aws
          profile: release

  build-linux-x86:
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
        hw: ["cpu", "gpu"]
    needs: setup-instance
    runs-on: ${{ needs.setup-instance.outputs.runner-name }}
    steps:
      - name: Set up GitHub environment
        run: |
          echo "HOME=/home/ubuntu" >> "${GITHUB_ENV}"
      - name: Checkout
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          submodules: recursive
          fetch-depth: 0
      - name: Set release version (nightly)
        if: ${{ env.RELEASE_TYPE == 'nightly' }}
        run: |
          NIGHTLY_VERSION=$(date +"%Y.%m.%d")
          NIGHTLY_VERSION_ONE_NUMBER=$(date +"%Y%m%d")
          LATEST_RELEASE_VERSION=$(git tag -l |grep "v.*" |sort |tail -n 1 | grep -e '[0-9].*' -o)
          echo "__version__ = \"${LATEST_RELEASE_VERSION}-dev${NIGHTLY_VERSION_ONE_NUMBER}\"" >| frontends/concrete-python/version.txt
          git tag "nightly-${NIGHTLY_VERSION}" || true
          git push origin "nightly-${NIGHTLY_VERSION}" || true
      - name: Set release version (public)
        if: ${{ env.RELEASE_TYPE == 'public' }}
        run: echo "__version__ = \"$(git describe --tags --abbrev=0 | grep -e '[0-9].*' -o)\"" >| frontends/concrete-python/version.txt
      - name: Expose release version from Python
        run: |
          # remove old version
          sed '/^__version__/d' -i frontends/concrete-python/concrete/fhe/version.py
          # add new version
          cat frontends/concrete-python/version.txt >> frontends/concrete-python/concrete/fhe/version.py
      - name: Set GPU Options
        if: ${{ matrix.hw == 'gpu' }}
        run: |
          {
            echo "CUDA_SUPPORT=ON"
            echo "TIMING_ENABLED=ON"
            echo "CUDA_PATH=${{ env.CUDA_PATH }}"
            echo "DATAFLOW_EXECUTION_ENABLED=OFF"
          } >> "${GITHUB_ENV}"
      - name: Set CPU Options
        if: ${{ matrix.hw == 'cpu' }}
        run: |
          {
            echo "CUDA_SUPPORT=OFF"
            echo "TIMING_ENABLED=OFF"
            echo "CUDA_PATH="
            echo "DATAFLOW_EXECUTION_ENABLED=ON"
          } >> "${GITHUB_ENV}"
      - name: Build wheel
        uses: addnab/docker-run-action@4f65fabd2431ebc8d299f8e5a018d79a769ae185 # v3
        id: build-compiler-bindings
        with:
          registry: ghcr.io
          image: ${{ env.DOCKER_IMAGE_TEST }}
          username: ${{ secrets.GHCR_LOGIN }}
          password: ${{ secrets.GHCR_PASSWORD }}
          options: >-
            -v ${{ github.workspace }}:/concrete
            -v ${{ github.workspace }}/build:/build
          shell: bash
          run: |
            set -e

            rm -rf /build/*

            export PYTHON=${{ format('python{0}', matrix.python-version) }}
            echo "Using $PYTHON"

            dnf -y install graphviz graphviz-devel

            cd /concrete/frontends/concrete-python
            make PYTHON=$PYTHON venv
            source .venv/bin/activate

            ccache -z

            cd /concrete/compilers/concrete-compiler/compiler
            make BUILD_DIR=/build \
              CCACHE=ON \
              DATAFLOW_EXECUTION_ENABLED=${{ env.DATAFLOW_EXECUTION_ENABLED }} \
              Python3_EXECUTABLE=$(which python) \
              CUDA_SUPPORT=${{ env.CUDA_SUPPORT }} \
              TIMING_ENABLED=${{ env.TIMING_ENABLED }} \
              CUDA_PATH=${{ env.CUDA_PATH }} \
              python-bindings

            echo "Debug: ccache statistics (after the build):"
            ccache -s



            cd /concrete/frontends/concrete-python

            export COMPILER_BUILD_DIRECTORY="/build"
            make whl

            deactivate
      - name: Upload wheel
        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
        with:
          name: ${{ format('{0}-wheel-{1}-linux-x86', matrix.hw, matrix.python-version) }}
          path: frontends/concrete-python/dist/*manylinux*.whl
          retention-days: 3
      - name: Build concrete rust
        # Note: That arbitrary if to not build in every python version
        if: ${{ matrix.python-version == '3.10' }}
        uses: addnab/docker-run-action@4f65fabd2431ebc8d299f8e5a018d79a769ae185 # v3
        with:
          registry: ghcr.io
          image: ${{ env.DOCKER_IMAGE_TEST }}
          username: ${{ secrets.GHCR_LOGIN }}
          password: ${{ secrets.GHCR_PASSWORD }}
          options: >-
            -v ${{ github.workspace }}:/concrete
            -v ${{ github.workspace }}/build:/build
          shell: bash
          run: |
            cargo install cxxbridge-cmd
            cd /concrete/compilers/concrete-compiler/compiler
            make BUILD_DIR=/build \
              CCACHE=ON \
              DATAFLOW_EXECUTION_ENABLED=${{ env.DATAFLOW_EXECUTION_ENABLED }} \
              Python3_EXECUTABLE=$(which python) \
              CUDA_SUPPORT=${{ env.CUDA_SUPPORT }} \
              TIMING_ENABLED=${{ env.TIMING_ENABLED }} \
              CUDA_PATH=${{ env.CUDA_PATH }} \
              concrete-rust
      - name: Upload binary libraries
        # Note: That arbitrary if to not build in every python version
        if: ${{ matrix.python-version == '3.10' }}
        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
        with:
          name: ${{ format('{0}-binaries-linux-x86_64', matrix.hw) }}
          path: |
            ${{ github.workspace }}/build/lib/libConcrete*.so
            ${{ github.workspace }}/build/lib/libhpx*.so
          retention-days: 3
      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "build-linux-x86(${{ matrix.hw }}/py${{matrix.python-version}}) finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

  release-checks:
    strategy:
      matrix:
        python-version: ["3.10"]
    needs: [ build-linux-x86 ]
    runs-on: ubuntu-latest
    steps:
      - name: Checkout
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          submodules: recursive
          fetch-depth: 0
      - name: Setup Python
        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
        with:
          python-version: ${{ matrix.python-version }}
      - name: Download wheels
        uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          name: ${{ format('cpu-wheel-{0}-linux-x86', matrix.python-version) }}
          path: ${{ format('cpu-wheel-{0}-linux-x86', matrix.python-version) }}
      - name: Check documentation
        run: |
          WHEEL_DIR=$(pwd)/${{ format('cpu-wheel-{0}-linux-x86', matrix.python-version) }}/
          CONCRETE_WHEEL="${WHEEL_DIR}/*.whl" .github/workflows/scripts/make_apidocs.sh
      - name: Upload docs.patch if failure
        if: ${{ failure() }}
        uses: actions/upload-artifact@50769540e7f4bd5e21e526ee35c689e35e0d6874 # v4.4.0
        with:
          name: docs-artifacts
          path: |
            docs.patch
            docs/dev/api
          retention-days: 3
      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "release-checks finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

  build-macos:
    needs: check-duplicate-nightly
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
        runs-on: ["aws-mac1-metal", "aws-mac2-metal"]
    runs-on: ${{ matrix.runs-on }}
    steps:
      - name: Checkout
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          submodules: recursive
          fetch-depth: 0
      - name: Install OS Dependencies
        run: |
          brew install ninja ccache
      - name: Setup rust toolchain for concrete-cpu
        uses: ./.github/actions/setup_rust_toolchain_for_concrete_cpu
      - name: Set release version (nightly)
        if: ${{ env.RELEASE_TYPE == 'nightly' }}
        run: |
          NIGHTLY_VERSION=$(date +"%Y%m%d")
          LATEST_RELEASE_VERSION=$(git tag -l |grep "v.*" |sort |tail -n 1 | grep -e '[0-9].*' -o)
          echo "__version__ = \"${LATEST_RELEASE_VERSION}-dev${NIGHTLY_VERSION}\"" >| frontends/concrete-python/version.txt
      - name: Set release version (public)
        if: ${{ env.RELEASE_TYPE == 'public' }}
        run: echo "__version__ = \"$(git describe --tags --abbrev=0 | grep -e '[0-9].*' -o)\"" >| frontends/concrete-python/version.txt
      - name: Expose release version from Python
        run: |
          # remove old version
          sed -e '/^__version__/d' -i '' frontends/concrete-python/concrete/fhe/version.py
          # add new version
          cat frontends/concrete-python/version.txt >> frontends/concrete-python/concrete/fhe/version.py
      - name: Build wheel
        run: |
          # Hack for mac2 host...
          export PATH=$PATH:/Users/ec2-user/.cargo/bin:/Users/ec2-user/.pyenv/versions/3.8.20/bin

          CONCRETE_PYTHON=$(pwd)/frontends/concrete-python
          CONCRETE_COMPILER=$(pwd)/compilers/concrete-compiler/compiler
          export COMPILER_BUILD_DIRECTORY=$CONCRETE_COMPILER/build
          export PYTHON=${{ format('python{0}', matrix.python-version) }}
          echo "Using $PYTHON"

          # Setup pkg-config to find OpenBLAS (scipy need it)
          export PKG_CONFIG_PATH="/opt/homebrew/opt/openblas/lib/pkgconfig"

          # Setup vitual environment
          rm -rf .venv
          $PYTHON -m venv .venv && . .venv/bin/activate

          # Install requirements
          pip install -r "${CONCRETE_PYTHON}"/requirements.txt
          pip install -r "${CONCRETE_PYTHON}"/requirements.dev.txt

          # Build python bindings of concrete compiler
          cd "${CONCRETE_COMPILER}" || exit
          echo "Debug: ccache statistics (prior to the build):" && ccache -s
          make Python3_EXECUTABLE="$(which python)" python-bindings
          echo "Debug: ccache statistics (after the build):" && ccache -s

          # Build wheel
          cd "${CONCRETE_PYTHON}" || exit
          rm -rf dist
          mkdir -p dist
          pip wheel -v --no-deps -w dist .
          delocate-wheel -v dist/*macos*.whl

          deactivate
      - name: Upload wheel
        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
        with:
          name: ${{ format('cpu-wheel-{0}-{1}', matrix.python-version, matrix.runs-on) }}
          path: frontends/concrete-python/dist/*macos*.whl
          retention-days: 3
      - name: Build concrete rust
        # Note: That arbitrary if to not build in every python version
        if: ${{ matrix.python-version == '3.10' }}
        run: |
          cargo install cxxbridge-cmd
          cd compilers/concrete-compiler/compiler
          make concrete-rust
      - name: Upload binary libraries mac2
        # Note: That arbitrary if to not build in every python version
        if: ${{ matrix.python-version == '3.10' && matrix.runs-on == 'aws-mac2-metal' }}
        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
        with:
          name: cpu-binaries-macosx_arm64
          path: compilers/concrete-compiler/compiler/build/lib/libConcrete*.dylib
          retention-days: 3
      - name: Upload binary libraries mac1
        if: ${{ matrix.python-version == '3.10' && matrix.runs-on == 'aws-mac1-metal' }}
        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
        with:
          name: cpu-binaries-macosx_x86_64
          path: compilers/concrete-compiler/compiler/build/lib/libConcrete*.dylib
          retention-days: 3
      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "build-macos finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

  hash-cpu:
    # Generate hashes for the wheels, used later for provenance.
    needs: [build-linux-x86, build-macos]
    runs-on: ubuntu-latest
    outputs:
      hash: ${{ steps.hash.outputs.hash }}
    steps:
      - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          path: frontends/concrete-python/dist
          pattern: 'cpu-wheel-*'
          merge-multiple: true
      - name: generate hash
        id: hash
        run: cd frontends/concrete-python/dist && echo "hash=$(sha256sum ./*.whl | base64 -w0)" >> "${GITHUB_OUTPUT}"

  provenance-cpu:
    needs: [hash-cpu]
    permissions:
      actions: read
      id-token: write
      contents: write
    # Can't pin with hash due to how this workflow works.
    uses: slsa-framework/slsa-github-generator/.github/workflows/generator_generic_slsa3.yml@v2.0.0
    with:
      base64-subjects: ${{ needs.hash-cpu.outputs.hash }}
      provenance-name: cpu-wheels.intoto.jsonl

  hash-gpu:
    # Generate hashes for the wheels, used later for provenance.
    needs: [build-linux-x86, build-macos]
    runs-on: ubuntu-latest
    outputs:
      hash: ${{ steps.hash.outputs.hash }}
    steps:
      - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          path: frontends/concrete-python/dist
          pattern: 'gpu-wheel-*'
          merge-multiple: true
      - name: generate hash
        id: hash
        run: cd frontends/concrete-python/dist && echo "hash=$(sha256sum ./*.whl | base64 -w0)" >> "${GITHUB_OUTPUT}"

  provenance-gpu:
    needs: [hash-gpu]
    permissions:
      actions: read
      id-token: write
      contents: write
    # Can't pin with hash due to how this workflow works.
    uses: slsa-framework/slsa-github-generator/.github/workflows/generator_generic_slsa3.yml@v2.0.0
    with:
      base64-subjects: ${{ needs.hash-gpu.outputs.hash }}
      provenance-name: gpu-wheels.intoto.jsonl

  push:
    needs: [build-linux-x86, build-macos, provenance-cpu, provenance-gpu]
    runs-on: ubuntu-latest
    strategy:
      matrix:
        hw: ["cpu", "gpu"]
    steps:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          fetch-depth: 0
      - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          path: wheels
          pattern: '${{ matrix.hw }}-wheel-*'
          merge-multiple: true
      - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          pattern: '${{ matrix.hw }}-binaries-*'
      - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          pattern: '${{ matrix.hw }}*.intoto.jsonl'
      # When building a new public tag, create a new draft release.
      - name: create draft release
        if: ${{ env.RELEASE_TYPE == 'public'}}
        run: |
          TAG=$(git describe --tags --abbrev=0)
          HW=${{ matrix.hw }}
          export TAG
          echo "${TAG}"
          # Create zip of binaries folders
          for i in ./*binaries*; do
            zip "${i}.zip" "${i}"/*
          done
          gh release create --draft --repo ${{ github.repository }} \
          --verify-tag "${TAG}" \
          --title "${TAG} - ${HW^^}" \
          wheels/* ./*.intoto.jsonl/* *binaries*.zip
        env:
          GH_TOKEN: ${{ github.token }}
      - name: Upload wheels to S3
        if: ${{ env.RELEASE_TYPE == 'public' || env.RELEASE_TYPE == 'nightly' }}
        env:
          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_IAM_ID }}
          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_IAM_KEY }}
          AWS_DEFAULT_REGION: ${{ secrets.AWS_REGION }}
          S3_BUCKET_NAME: ${{ secrets.AWS_S3_PYPI_BUCKET_NAME }}
          CLOUDFRONT_DISTRIBUTION_ID: ${{ secrets.AWS_CLOUDFRONT_PYPI_DISTRIBUTION_ID }}
        run: |
          pip install boto3 bigtree
          # upload wheels
          aws s3 sync ./wheels/ "s3://${S3_BUCKET_NAME}/${{ matrix.hw }}/concrete-python"
          # update indexes and invalidate cloudfront cache
          python .github/workflows/scripts/s3_update_html_indexes.py
      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "push finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

  test-linux-x86-cpu:
    needs: [setup-instance, build-linux-x86]
    continue-on-error: true
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
        # FIXME: we need to investigate why this tests are taking too long
        # in the meantime, we only run them for one version when in nightly
        exclude:
          - python-version: ${{ (github.event_name == 'push' && contains(github.ref, 'refs/tags/')) && '' || '3.9' }}
          - python-version: ${{ (github.event_name == 'push' && contains(github.ref, 'refs/tags/')) && '' || '3.10'}}
          - python-version: ${{ (github.event_name == 'push' && contains(github.ref, 'refs/tags/')) && '' || '3.11'}}
          - python-version: ${{ (github.event_name == 'push' && contains(github.ref, 'refs/tags/')) && '' || '3.12'}}
    runs-on: ${{ needs.setup-instance.outputs.runner-name }}
    steps:
      # HOME is needed by actions-rs/toolchain
      - run: |
          echo "HOME=/home/ubuntu" >> "${GITHUB_ENV}"
      - name: Install rust
        uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # v1.0.7
        with:
          toolchain: nightly
          default: true
      - name: Setup Python
        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
        with:
          python-version: ${{ matrix.python-version }}
      - name: Download wheels
        uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          name: ${{ format('cpu-wheel-{0}-linux-x86', matrix.python-version) }}
          path: ${{ format('cpu-wheel-{0}-linux-x86', matrix.python-version) }}
      - name: Checkout the repository
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          submodules: recursive
          path: repo
      - name: Test wheel
        run: |
          WHEEL_DIR=$(pwd)/${{ format('cpu-wheel-{0}-linux-x86', matrix.python-version) }}/
          CONCRETE_PYTHON=$(pwd)/repo/frontends/concrete-python

          # Initialize an empty test environment
          cd "$(mktemp -d)"
          python -m venv .testenv && source .testenv/bin/activate

          # Install the concrete-python wheel
          pip install "${WHEEL_DIR}"/*.whl

          # Install extra requirements for tests
          sudo apt update -y
          sudo apt install -y graphviz libgraphviz-dev
          pip install -r "${CONCRETE_PYTHON}"/requirements.extra-full.txt
          pip install -r "${CONCRETE_PYTHON}"/requirements.dev.txt

          # TODO - check for version

          # Copy test files
          cp -R "${CONCRETE_PYTHON}"/tests .
          cp -R "${CONCRETE_PYTHON}"/examples .
          cp -R "${CONCRETE_PYTHON}"/pytest.ini .
          cp "${CONCRETE_PYTHON}"/Makefile .

          # Running tests
          make tfhers-utils
          pytest tests -svv -n auto --maxprocesses="${PYTEST_MAXPROCESSES}"
      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "test-linux-x86-cpu (${{ matrix.python-version }}) finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

  teardown-instance:
    needs: [ setup-instance, test-linux-x86-cpu ]
    if: ${{ always() && needs.setup-instance.result != 'skipped' }}
    runs-on: ubuntu-latest
    steps:
      - name: Stop instance
        id: stop-instance
        uses: zama-ai/slab-github-runner@f26b8d611b2e695158fb0a6980834f0612f65ef8 # v1.4.0
        with:
          mode: stop
          github-token: ${{ secrets.SLAB_ACTION_TOKEN }}
          slab-url: ${{ secrets.SLAB_BASE_URL }}
          job-secret: ${{ secrets.JOB_SECRET }}
          label: ${{ needs.setup-instance.outputs.runner-name }}

      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "Instance teardown finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

  setup-gpu-test-instance:
    runs-on: ubuntu-latest
    needs: [build-linux-x86]
    outputs:
      runner-name: ${{ steps.start-instance.outputs.label }}
    steps:
      - name: Start instance
        id: start-instance
        uses: zama-ai/slab-github-runner@f26b8d611b2e695158fb0a6980834f0612f65ef8 # v1.4.0
        with:
          mode: start
          github-token: ${{ secrets.SLAB_ACTION_TOKEN }}
          slab-url: ${{ secrets.SLAB_BASE_URL }}
          job-secret: ${{ secrets.JOB_SECRET }}
          backend: aws
          profile: gpu-test

  test-linux-x86-gpu:
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
      fail-fast: false
    needs: [setup-gpu-test-instance, build-linux-x86]
    runs-on: ${{ needs.setup-gpu-test-instance.outputs.runner-name }}
    steps:
      # HOME is needed by actions-rs/toolchain
      - run: |
          echo "HOME=/home/ubuntu" >> "${GITHUB_ENV}"
      - name: Install rust
        uses: actions-rs/toolchain@16499b5e05bf2e26879000db0c1d13f7e13fa3af # v1.0.7
        with:
          toolchain: nightly
          default: true
      - name: Setup Python
        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
        with:
          python-version: ${{ matrix.python-version }}

      - name: Download wheels
        uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          name: ${{ format('gpu-wheel-{0}-linux-x86', matrix.python-version) }}
          path: ${{ format('gpu-wheel-{0}-linux-x86', matrix.python-version) }}

      - name: Install concrete-python
        run: |
          WHEEL_DIR=$(pwd)/${{ format('gpu-wheel-{0}-linux-x86', matrix.python-version) }}/
          pip install "${WHEEL_DIR}"/*.whl

      - name: Checkout the repository
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          path: repo

      - name: Test wheel
        run: |
          CONCRETE_PYTHON=$(pwd)/repo/frontends/concrete-python

          # Install extra requirements for tests
          sudo apt update -y
          sudo apt install -y graphviz libgraphviz-dev
          pip install -r "${CONCRETE_PYTHON}"/requirements.extra-full.txt
          pip install -r "${CONCRETE_PYTHON}"/requirements.dev.txt

          # Running tests
          cd "${CONCRETE_PYTHON}"
          make pytest-gpu

      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "test-linux-x86-gpu (${{ matrix.python-version }}) finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"


  teardown-gpu-test-instance:
    needs: [ setup-gpu-test-instance, test-linux-x86-gpu ]
    if: ${{ always() && needs.setup-gpu-test-instance.result != 'skipped' }}
    runs-on: ubuntu-latest
    steps:
      - name: Stop instance
        id: stop-instance
        uses: zama-ai/slab-github-runner@f26b8d611b2e695158fb0a6980834f0612f65ef8 # v1.4.0
        with:
          mode: stop
          github-token: ${{ secrets.SLAB_ACTION_TOKEN }}
          slab-url: ${{ secrets.SLAB_BASE_URL }}
          job-secret: ${{ secrets.JOB_SECRET }}
          label: ${{ needs.setup-gpu-test-instance.outputs.runner-name }}

      - name: Slack Notification
        if: ${{ failure() }}
        continue-on-error: true
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ job.status }}
          SLACK_MESSAGE: "Instance teardown finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})"

  test-macos:
    needs: [build-macos]
    continue-on-error: true
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
        runs-on: ["aws-mac1-metal", "aws-mac2-metal"]
    runs-on: ${{ matrix.runs-on }}
    outputs:
      slack_message: ${{ steps.prepare_slack_notif.outputs.slack_message }}
      slack_color: ${{ steps.prepare_slack_notif.outputs.slack_color }}
    steps:
      - name: Download wheels
        uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8
        with:
          name: ${{ format('cpu-wheel-{0}-{1}', matrix.python-version, matrix.runs-on) }}
          path: ${{ format('cpu-wheel-{0}-{1}', matrix.python-version, matrix.runs-on) }}
      - name: Checkout the repository
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          submodules: recursive
          path: repo
      - name: Test wheel
        run: |
          WHEEL_DIR=$(pwd)/${{ format('cpu-wheel-{0}-{1}', matrix.python-version, matrix.runs-on) }}/
          CONCRETE_PYTHON=$(pwd)/repo/frontends/concrete-python
          PYTHON=${{ format('python{0}', matrix.python-version) }}

          # Initialize an empty test environment
          TEST_TMP_DIR=$(mktemp -d)
          echo "TEST_TMP_DIR=${TEST_TMP_DIR}" >> "${GITHUB_ENV}"
          cd "${TEST_TMP_DIR}" || exit

          # Activate virtual environment
          $PYTHON -m venv .testenv && source .testenv/bin/activate

          # Install extra requirements for tests
          pip install "${WHEEL_DIR}"/*macos*.whl
          pip install -r "${CONCRETE_PYTHON}"/requirements.dev.txt

          "${CONCRETE_PYTHON}"/../../.github/workflows/scripts/fix_multi_omp_bug_macos.sh

          # Copy test files
          cp -R "${CONCRETE_PYTHON}"/tests .
          cp -R "${CONCRETE_PYTHON}"/examples .
          cp -R "${CONCRETE_PYTHON}"/pytest.ini .
          cp "${CONCRETE_PYTHON}"/Makefile .

          # Fix: this is because cargo cannot be found
          . "$HOME/.cargo/env"

          # Running tests
          make tfhers-utils
          mkdir ./KeySetCache
          pytest tests -svv -n auto --maxprocesses="${PYTEST_MAXPROCESSES}" --key-cache "./KeySetCache" -m "not dataflow and not graphviz"
      - name: Cleanup host
        if: success() || failure()
        run: |
          rm -rf "${TEST_TMP_DIR}"
      - name: Prepare Slack Notification
        id: prepare_slack_notif
        if: ${{ failure() }}
        continue-on-error: true
        run: |
          echo "slack_message=test-macos (${{matrix.runs-on}}/${{ matrix.python-version }}) finished with status: ${{ job.status }}. (${{ env.ACTION_RUN_URL }})" >> "$GITHUB_OUTPUT"
          echo "slack_color=${{ job.status }}" >> "$GITHUB_OUTPUT"

  slack-notif-macos:
    needs: ["test-macos"]
    runs-on: "ubuntu-latest"
    if: always()
    steps:
      - name: Slack Notification
        # we want to check that prepare_slack_notif was run
        if: ${{ needs.test-macos.outputs.slack_color != '' }}
        uses: rtCamp/action-slack-notify@c33737706dea87cd7784c687dadc9adf1be59990
        env:
          SLACK_COLOR: ${{ needs.test-macos.outputs.slack_color }}
          SLACK_MESSAGE: ${{ needs.test-macos.outputs.slack_message }}
