name: Concrete Python Release
on:
  workflow_dispatch:
    inputs:
      instance_id:
        description: 'Instance ID'
        type: string
      instance_image_id:
        description: 'Instance AMI ID'
        type: string
      instance_type:
        description: 'Instance product type'
        type: string
      runner_name:
        description: 'Action runner name'
        type: string
      request_id:
        description: 'Slab request ID'
        type: string
      user_inputs:
        description: 'either "nightly" or "public" or "private" to specify the release type'
        required: true
        default: 'nightly'
        type: string


env:
  DOCKER_IMAGE_TEST: ghcr.io/zama-ai/concrete-compiler
  GLIB_VER: 2_28
  RELEASE_TYPE: ${{ inputs.user_inputs }}

jobs:
  build-linux-x86:
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11"]

    runs-on: ${{ github.event.inputs.runner_name }}
    steps:
      - name: Log instance configuration
        run: |
          echo "IDs: ${{ inputs.instance_id }}"
          echo "AMI: ${{ inputs.instance_image_id }}"
          echo "Type: ${{ inputs.instance_type }}"
          echo "Request ID: ${{ inputs.request_id }}"
          echo "User Inputs: ${{ inputs.user_inputs }}"

      - name: Set up GitHub environment
        run: |
          echo "HOME=/home/ubuntu" >> "${GITHUB_ENV}"

      - name: Checkout
        uses: actions/checkout@v3
        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: cp frontends/concrete-python/version.txt frontends/concrete-python/concrete/fhe/version.py

      - name: Build wheel
        uses: addnab/docker-run-action@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
            -v ${{ env.SSH_AUTH_SOCK }}:/ssh.socket
            -e SSH_AUTH_SOCK=/ssh.socket
            ${{ env.DOCKER_GPU_OPTION }}
          shell: bash
          run: |
            set -e

            rustup toolchain install nightly-2024-01-31
            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
            
            cd /concrete/compilers/concrete-compiler/compiler
            make BUILD_DIR=/build CCACHE=ON DATAFLOW_EXECUTION_ENABLED=ON Python3_EXECUTABLE=$(which python) 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@v4
        with:
          name: ${{ format('wheel-{0}-linux-x86', matrix.python-version) }}
          path: frontends/concrete-python/dist/*manylinux*.whl
          retention-days: 3

  build-macos:
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11"]
        runs-on: ["aws-mac1-metal", "aws-mac2-metal"]

    runs-on: ${{ matrix.runs-on }}
    steps:
      - name: Checkout
        uses: actions/checkout@v3
        with:
          submodules: recursive
          fetch-depth: 0

      - name: Install OS Dependencies
        run: |
          brew install ninja ccache

      - name: Setup rust toolchain for concrete-cpu
        uses: ./.github/workflows/setup_rust_toolchain_for_concrete_cpu

      - 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

      - 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: cp frontends/concrete-python/version.txt frontends/concrete-python/concrete/fhe/version.py

      - name: Build wheel
        run: |
          export CONCRETE_PYTHON=$(pwd)/frontends/concrete-python
          export 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
          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
          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@v4
        with:
          name: ${{ format('wheel-{0}-{1}', matrix.python-version, matrix.runs-on) }}
          path: frontends/concrete-python/dist/*macos*.whl
          retention-days: 3

  push:
    needs: [build-linux-x86, build-macos]
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3

      - uses: actions/download-artifact@v4
        with:
          path: wheels
          merge-multiple: true

      - 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}/concrete-python
          # update indexes and invalidate cloudfront cache
          python .github/workflows/scripts/s3_update_html_indexes.py

      - name: Start pushing Docker images
        if: ${{ env.RELEASE_TYPE == 'public' }}
        run: |
          export TAG=$(git describe --tags --abbrev=0)
          curl -L \
            -X POST \
            -H "Accept: application/vnd.github+json" \
            -H "Authorization: Bearer ${{ secrets.GITHUB_TOKEN }}" \
            -H "X-GitHub-Api-Version: 2022-11-28" \
            https://api.github.com/repos/zama-ai/concrete/actions/workflows/concrete_python_push_docker_image.yml/dispatches \
            -d "{\"ref\": \"$TAG\", \"inputs\": {\"tag\":\"$TAG\"}}"

  test-linux-x86:
    needs: [build-linux-x86]
    continue-on-error: true
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11"]
    runs-on: ${{ github.event.inputs.runner_name }}
    steps:
      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
      - name: Download wheels
        uses: actions/download-artifact@v4
        with:
          name: ${{ format('wheel-{0}-linux-x86', matrix.python-version) }}
          path: ${{ format('wheel-{0}-linux-x86', matrix.python-version) }}
      - name: Checkout the repository
        uses: actions/checkout@v3
        with:
          submodules: recursive
          path: repo
      - name: Test wheel
        run: |
          WHEEL_DIR=$(pwd)/${{ format('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/pytest.ini .

          # Running tests
          pytest tests -svv -n auto

  test-macos:
    needs: [build-macos]
    continue-on-error: true
    strategy:
      matrix:
        python-version: ["3.8", "3.9", "3.10", "3.11"]
        runs-on: ["aws-mac1-metal", "aws-mac2-metal"]
    runs-on: ${{ matrix.runs-on }}
    steps:
      - name: Download wheels
        uses: actions/download-artifact@v4
        with:
          name: ${{ format('wheel-{0}-{1}', matrix.python-version, matrix.runs-on) }}
          path: ${{ format('wheel-{0}-{1}', matrix.python-version, matrix.runs-on) }}
      - name: Checkout the repository
        uses: actions/checkout@v3
        with:
          submodules: recursive
          path: repo
      - name: Test wheel
        run: |
          WHEEL_DIR=$(pwd)/${{ format('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
          export TEST_TMP_DIR=$(mktemp -d)
          echo "TEST_TMP_DIR=$TEST_TMP_DIR" >> "${GITHUB_ENV}"
          cd $TEST_TMP_DIR

          # 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

          # MacOS x86 have conflict between our OpenMP library, and one from torch
          # we fix it by using a single one (from torch)
          # see discussion: https://discuss.python.org/t/conflicting-binary-extensions-in-different-packages/25332/8
          export SITE_PACKAGES=$(pwd)/.testenv/lib/$PYTHON_VERSION/site-packages/
          find $SITE_PACKAGES -not \( -path $SITE_PACKAGES/concrete -prune \) -name 'lib*omp5.dylib' -or -name 'lib*omp.dylib' | xargs -n 1  ln -f -s $SITE_PACKAGES/concrete/.dylibs/libomp.dylib

          # Copy test files
          cp -R $CONCRETE_PYTHON/tests .
          cp -R $CONCRETE_PYTHON/pytest.ini .
          cp $CONCRETE_PYTHON/Makefile .

          # Running tests
          mkdir ./KeySetCache
          pytest tests -svv -n auto --key-cache "./KeySetCache" -m "not dataflow and not graphviz"

      - name: Cleanup host
        if: success() || failure()
        run: |
          rm -rf $TEST_TMP_DIR