# Copyright (c) 2020-2021, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
name: Update build cache
on:
  schedule:
    - cron: 0 0 * * *
  push:
    branches:
      - main
  workflow_dispatch:
    inputs:
      runner:
        required: false
        default: self-hosted-azure-builder
        type: string
        description: VM to use for build

jobs:
  pre-flight:
    runs-on: ubuntu-latest
    outputs:
      build_args: ${{ steps.manifest.outputs.BUILD_ARGS }}
      cache-from: ${{ steps.cache_from.outputs.LAST_PRS }}
    steps:
      - name: Checkout branch
        uses: actions/checkout@v4

      - name: Parse manifest.json
        id: manifest
        run: |
          BUILD_ARGS=$(cat << EOF
          BASE_IMAGE=$(cat requirements/manifest.json | jq -r '."ngc-pytorch"')
          TRTLLM_REPO=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."trt-llm".repo')
          TRTLLM_TAG=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."trt-llm".ref')
          MLM_REPO=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."megatron-lm".repo')
          MLM_TAG=$(cat requirements/manifest.json | jq -r '."vcs-dependencies"."megatron-lm".ref')
          TE_REPO=$(cat requirements/manifest.json | jq -r '."vcs-dependencies".transformer_engine.repo')
          TE_TAG=$(cat requirements/manifest.json | jq -r '."vcs-dependencies".transformer_engine.ref')
          APEX_REPO=$(cat requirements/manifest.json | jq -r '."vcs-dependencies".apex.repo')
          APEX_TAG=$(cat requirements/manifest.json | jq -r '."vcs-dependencies".apex.ref')
          EOF
          )

          echo "BUILD_ARGS<<EOF" >> $GITHUB_OUTPUT
          echo "$BUILD_ARGS" >> $GITHUB_OUTPUT
          echo "EOF" >> $GITHUB_OUTPUT

      - name: Get last merged PR
        id: cache_from
        env:
          GH_TOKEN: ${{ github.token }}
        run: |
          LAST_PRS=$(gh api graphql -f query='
            query {
              repository(owner: "NVIDIA", name: "NeMo") {
                pullRequests(states: MERGED, first: 100, orderBy: {field: UPDATED_AT, direction: DESC}) {
                  nodes {
                    number
                  }
                }
              }
            }' | jq -r '.data.repository.pullRequests.nodes[].number' | while read -r number; do
              echo "nemoci.azurecr.io/nemo_container-buildcache:$number"
            done)

          echo "LAST_PRS<<EOF" >> $GITHUB_OUTPUT
          echo "$LAST_PRS" >> $GITHUB_OUTPUT
          echo "EOF" >> $GITHUB_OUTPUT

  cicd-test-container-build:
    needs: [pre-flight]
    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_build_container.yml@v0.27.0
    strategy:
      fail-fast: false
      matrix:
        include:
          - dockerfile: docker/Dockerfile.ci
            image-name: nemo_container_automodel
          - dockerfile: docker/Dockerfile.ci
            image-name: nemo_container_nemo2
          - dockerfile: docker/Dockerfile.ci
            image-name: nemo_container_speech
          - dockerfile: docker/Dockerfile.ci
            image-name: nemo_container
          - dockerfile: docker/Dockerfile.ci.export_deploy
            image-name: nemo_container_export_deploy
    with:
      image-name: ${{ matrix.image-name }}
      dockerfile: ${{ matrix.dockerfile }}
      image-label: nemo-core
      build-args: |
        IMAGE_LABEL=nemo-core
        NEMO_TAG=${{ github.sha }}
        NEMO_REPO=https://github.com/NVIDIA/NeMo
        ${{ needs.pre-flight.outputs.BUILD_ARGS }}
      runner: ${{ inputs.runner || 'self-hosted-azure-builder' }}
      use-inline-cache: false
      prune-filter-timerange: 24h
      cache-from: |
        nemoci.azurecr.io/${{ matrix.image-name }}-buildcache:main
        ${{ needs.pre-flight.outputs.cache-from }}