name: "Track Size"

on:
  workflow_call:
    inputs:
      CONCURRENCY:
        required: true
        type: string
      PR_NUMBER:
        required: false
        type: number
      WITH_EXAMPLES:
        required: true
        type: boolean

defaults:
  run:
    shell: bash --noprofile --norc -euo pipefail {0}

permissions:
  contents: write
  id-token: write
  deployments: write
  pull-requests: write

jobs:
  track-sizes:
    name: "Track Sizes"
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
        with:
          fetch-depth: 0 # we need full history
          ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.ref || '' }}

      - name: Get context
        id: context
        run: |
          echo "short_sha=$(echo ${{ github.sha }} | cut -c1-7)" >> "$GITHUB_OUTPUT"

      - id: "auth"
        uses: google-github-actions/auth@v2
        with:
          workload_identity_provider: ${{ secrets.GOOGLE_WORKLOAD_IDENTITY_PROVIDER }}
          service_account: ${{ secrets.GOOGLE_SERVICE_ACCOUNT }}

      - name: "Set up Cloud SDK"
        uses: "google-github-actions/setup-gcloud@v2"
        with:
          version: ">= 363.0.0"

      - name: Download web_viewer
        uses: actions/download-artifact@v4
        with:
          name: web_viewer
          path: "crates/viewer/re_web_viewer_server/web_viewer"

      - name: Download examples
        if: ${{ inputs.WITH_EXAMPLES }}
        uses: actions/download-artifact@v4
        with:
          name: example_data
          path: example_data

      - name: Download base results
        run: |
          # Get base commit:
          # 1. From the index file
          # 2. From the latest commit in the base branch of the PR
          # 3. From the latest commit in the current branch
          index_path="gs://rerun-builds/sizes/index"
          if [ "$(gsutil -q stat $index_path ; echo $?)" = 0 ]; then
            gsutil cp $index_path "/tmp/base_index"
            base_commit=$(cat /tmp/base_index)
          else
            if [ -n ${{ inputs.PR_NUMBER }} ]; then
              base_commit=$(echo ${{ github.event.pull_request.base.sha }} | cut -c1-7)
            else
              base_commit=${{ steps.context.outputs.short_sha }}
            fi
          fi
          echo "base commit: $base_commit"

          # Download data for base commit, or default to empty file
          data_path="gs://rerun-builds/sizes/commit/$base_commit/data.json"
          if [ "$(gsutil -q stat $data_path ; echo $?)" = 0 ]; then
            gsutil cp $data_path "/tmp/prev.json"
          else
            echo "[]" > "/tmp/prev.json"
          fi

      - name: Measure sizes
        id: measure
        run: |
          entries=()

          entries+=("Wasm:crates/viewer/re_web_viewer_server/web_viewer/re_viewer_bg.wasm:MiB")
          entries+=("JS:crates/viewer/re_web_viewer_server/web_viewer/re_viewer.js:kiB")

          if [ ${{ inputs.WITH_EXAMPLES }} = "true" ]; then
            for file in example_data/*.rrd; do
              name=$(basename "$file")
              entries+=("$name:$file:MiB")
            done
          fi

          python3 scripts/ci/count_bytes.py "${entries[@]}"  > /tmp/sizes.json

          python3 scripts/ci/count_dependencies.py -p re_sdk    --no-default-features > /tmp/deps1.json
          python3 scripts/ci/count_dependencies.py -p re_viewer --all-features        > /tmp/deps2.json
          python3 scripts/ci/count_dependencies.py -p rerun     --all-features        > /tmp/deps3.json

          # Merge the results, putting dependencies first (on top):
          jq -s '.[0] + .[1] + .[2] + .[3]' /tmp/deps1.json /tmp/deps2.json /tmp/deps3.json /tmp/sizes.json > /tmp/data.json

          comparison=$(
            python3 scripts/ci/compare.py \
              --threshold=2% \
              --before-header=${{ (inputs.PR_NUMBER && github.event.pull_request.base.ref) || 'Before' }} \
              --after-header=${{ github.ref_name }} \
              "/tmp/prev.json" "/tmp/data.json"
          )
          {
            echo 'comparison<<EOF'
            echo "$comparison"
            echo EOF
          } >> "$GITHUB_OUTPUT"

          if [ -n "$comparison" ]; then
            echo "is_comparison_set=true" >> "$GITHUB_OUTPUT"
            echo "is_comparison_set=true"
          else
            echo "is_comparison_set=false" >> "$GITHUB_OUTPUT"
            echo "is_comparison_set=false"
          fi

          echo "$entries"
          echo "previous: $(cat /tmp/prev.json)"
          echo "current: $(cat /tmp/data.json)"
          echo "$comparison"

      - name: Upload data to GCS (commit)
        uses: google-github-actions/upload-cloud-storage@v3
        with:
          path: /tmp/data.json
          destination: "rerun-builds/sizes/commit/${{ steps.context.outputs.short_sha }}"
          process_gcloudignore: false

      - name: Create index file
        if: github.ref == 'refs/heads/main'
        run: |
          echo "${{ steps.context.outputs.short_sha }}" > "/tmp/index"

      - name: Upload index
        if: github.ref == 'refs/heads/main'
        uses: google-github-actions/upload-cloud-storage@v3
        with:
          path: /tmp/index
          destination: "rerun-builds/sizes"
          process_gcloudignore: false

      - name: Create PR comment
        if: inputs.PR_NUMBER != '' && steps.measure.outputs.is_comparison_set == 'true'
        # https://github.com/mshick/add-pr-comment
        uses: mshick/add-pr-comment@v3.9.1
        # GitHub API is very unreliable. We don't want to fail the entire workflow just because we couldn't post a comment.
        continue-on-error: true
        with:
          repo-token: ${{ secrets.GITHUB_TOKEN }}
          message: |
            # Size changes

            ${{ steps.measure.outputs.comparison }}

      - uses: prefix-dev/setup-pixi@v0.10.0
        with:
          pixi-version: v0.71.3

      - name: Render benchmark result
        if: github.ref == 'refs/heads/main'
        run: |
          pixi run python scripts/ci/render_bench.py sizes \
            --after $(date -d"180 days ago" +%Y-%m-%d) \
            --output "gs://rerun-builds/graphs"