name: Image_oncall
on:
  workflow_call:
    inputs:
      suffix:
        description: 'The tag subfix to use'
        required: true
        type: string
      should_push:
        description: 'Whether to push the image'
        required: false
        type: boolean
        default: False
      dockerfile:
        description: 'The Dockerfile to use'
        required: false
        type: string
      quay_username:
        description: 'Quay username for pushing images'
        required: false
        type: string
      workflow_dispatch_tag:
        description: 'The tag to use for workflow dispatch'
        required: false
        type: string
    secrets:
        QUAY_PASSWORD:
            description: 'Quay password for pushing images'
            required: false

jobs:
  build-push-digest:
    name: build
    runs-on: ${{ matrix.runner }}
    strategy:
      matrix:
        include:
          - arch: linux/amd64
            runner: ubuntu-latest
            tag: amd64
          - arch: linux/arm64
            runner: ubuntu-22.04-arm
            tag: arm64
    steps:
    - uses: actions/checkout@v6
      if: ${{ github.event_name != 'workflow_dispatch' }}
      with:
        fetch-depth: 0
        persist-credentials: false

    - uses: actions/checkout@v6
      if: ${{ github.event_name == 'workflow_dispatch' }}
      with:
        fetch-depth: 0
        persist-credentials: false
        ref: ${{ inputs.workflow_dispatch_tag }}

    - name: Free up disk space
      uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be # v1.3.1
      with:
        tool-cache: true
        docker-images: false

    - name: Publish - Login to Quay Container Registry
      if: ${{ inputs.should_push }}
      uses: docker/login-action@v3
      with:
        registry: quay.io
        username: ${{ inputs.quay_username }}
        password: ${{ secrets.QUAY_PASSWORD }}

    - name: Set up Docker Buildx
      uses: docker/setup-buildx-action@v3
      with:
        install: true
        driver: docker-container
        use: true

    - name: Set cache ref
      id: cache
      run: |
        if [ "${{ github.ref_type }}" = "tag" ]; then
          # For tag events, use the images built from source branch as cache (the tag image doesn't exist yet).
          if [ -z "$branch" ]; then
            branch=$(git branch -r --contains HEAD \
              | grep -v 'HEAD' \
              | sed 's|[[:space:]]*origin/||' \
              | head -1)
          fi
          branch="${branch:-main}"
        else
          # For PR events github.ref_name is "<PR_number>/merge" which has no cached image;
          # use base_ref (target branch) instead. For push/schedule, base_ref is empty so
          # fall back to ref_name which is the actual branch name.
          branch="${{ github.base_ref || github.ref_name }}"
        fi
        # Replace / with - for use in image tags
        branch="${branch//\//-}"
        echo "ref=quay.io/ascend/vllm-ascend:${branch}-${{ inputs.suffix }}" >> $GITHUB_OUTPUT

    - name: Build and push
      uses: docker/build-push-action@v6
      id: build
      with:
        platforms: ${{ matrix.arch }}
        # use the current repo path as the build context, ensure .git is contained
        context: .
        file: ${{ inputs.dockerfile || 'Dockerfile' }}
        # only trigger when tag, branch/main push
        push: ${{ inputs.should_push }}
        outputs: type=image,name=quay.io/ascend/vllm-ascend,push-by-digest=true,name-canonical=true,push=${{ inputs.should_push }}
        build-args: |
          PIP_INDEX_URL=https://pypi.org/simple
        # use previously pushed multi-arch image as cache to speed up builds
        cache-from: type=registry,ref=${{ steps.cache.outputs.ref }}
        provenance: false

    - name: Export digest
      run: |
        mkdir -p ${{ runner.temp }}/digests
        digest="${{ steps.build.outputs.digest }}"
        touch "${{ runner.temp }}/digests/${digest#sha256:}"

    - name: Upload digest
      uses: actions/upload-artifact@v6
      with:
        name: digests-${{ inputs.suffix }}-${{ matrix.tag }}
        path: ${{ runner.temp }}/digests/*
        if-no-files-found: error
        retention-days: 1


  merge-image:
    runs-on: ubuntu-latest
    needs: build-push-digest
    if: ${{ inputs.should_push }}
    steps:
      - name: Checkout
        uses: actions/checkout@v6

      - name: Download arm64 digests
        uses: actions/download-artifact@v7
        with:
          path: ${{ runner.temp }}/digests
          pattern: digests-${{ inputs.suffix }}-arm64
          merge-multiple: true

      - name: Download amd64 digests
        uses: actions/download-artifact@v7
        with:
          path: ${{ runner.temp }}/digests
          pattern: digests-${{ inputs.suffix }}-amd64
          merge-multiple: true

      - name: Prepare suffix
        id: suffix
        run: |
          if [ -n "${{ inputs.suffix }}" ]; then
            echo "SUFFIX=-${{ inputs.suffix }}" >> $GITHUB_ENV
          else
            echo "SUFFIX=" >> $GITHUB_ENV
          fi

      - name: Docker meta
        id: meta
        uses: docker/metadata-action@v5
        with:
          # TODO(yikun): add more hub image and a note on release policy for container image
          images: |
            quay.io/ascend/vllm-ascend
          # Note for test case
          # https://github.com/marketplace/actions/docker-metadata-action#typeref
          # 1. branch job publish per main/*-dev branch commits
          # 2. main and dev pull_request is build only, so the tag pr-N-openeuler is fine
          # 3. only pep440 matched tag will be published:
          #    - v0.7.1 --> v0.7.1-openeuler
          #    - pre/post/dev: v0.7.1rc1-openeuler/v0.7.1rc1-openeuler/v0.7.1rc1.dev1-openeuler/v0.7.1.post1-openeuler, no latest
          #      which follow the rule from vLLM with prefix v
          # TODO(yikun): the post release might be considered as latest release
          tags: |
            type=ref,event=branch,suffix=${{ env.SUFFIX }}
            type=pep440,pattern={{raw}},suffix=${{ env.SUFFIX }}
            type=schedule,pattern=main,suffix=${{ env.SUFFIX }}
            type=raw,value=${{ inputs.workflow_dispatch_tag }},enable=${{ github.event_name == 'workflow_dispatch' }},suffix=${{ env.SUFFIX }}
          flavor:
            latest=false

      - name: Login to Quay
        uses: docker/login-action@v3
        with:
          registry: quay.io
          username: ${{ inputs.quay_username }}
          password: ${{ secrets.QUAY_PASSWORD }}

      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3

      - name: Merge and push multi-arch image
        env:
          IMAGE: quay.io/ascend/vllm-ascend
          TAGS: ${{ steps.meta.outputs.tags }}
        run: |
          DIGESTS=$(printf "$IMAGE@sha256:%s " $(ls ${{ runner.temp }}/digests))

          echo "Digests: $DIGESTS"
          echo "Current tags:"
          echo "$TAGS"

          for tag in $TAGS; do
            echo "Creating tag $tag"
            docker buildx imagetools create \
              -t "$tag" \
              $DIGESTS
          done