name: 'Docker Build'
description: 'Build Dynamo container images'
inputs:
  framework:
    description: 'Framework to build'
    required: true
    default: 'vllm'
  target:
    description: 'Target to build'
    required: false
    default: 'runtime'
  platform:
    description: 'Docker platform to build on, ie. linux/amd64'
    required: false
    default: 'linux/amd64'
  cuda_version:
    description: 'Optional override for CUDA_VERSION build-arg'
    required: true
  image_tag:
    description: 'Custom image tag'
    required: true
  ci_token:
    description: 'CI Token'
    required: false
  aws_default_region:
    description: 'AWS Default Region'
    required: false
  sccache_s3_bucket:
    description: 'SCCache S3 Bucket'
    required: false
  aws_account_id:
    description: 'AWS Account ID'
    required: false
  aws_access_key_id:
    description: 'AWS Access Key ID'
    required: false
  aws_secret_access_key:
    description: 'AWS Secret Access Key'
    required: false
  no_cache:
    description: 'Disable Docker build cache'
    required: false
  extra_tags:
    description: 'Additional image tags (newline-separated list of full image:tag references)'
    required: false
    default: ''
  push_image:
    description: 'Push the image to the registry'
    required: false
    default: 'false'
  no_load:
    description: 'Do not load the image into docker (useful for validation-only builds)'
    required: false
    default: 'true'
  use_sccache:
    description: 'Use SCCache for caching'
    required: false
    default: 'false'
  ci:
    description: 'CI mode: for frontend target, uses existing buildx builder and pushes EPP image to ECR'
    required: false
    default: 'false'
outputs:
  image_tag:
    description: 'Image Tag'
    value: ${{ steps.build.outputs.image_tag }}

runs:
  using: "composite"
  steps:
    - name: Build image
      id: build
      shell: bash
      env:
        GITHUB_TOKEN: ${{ inputs.ci_token }}
        AWS_DEFAULT_REGION: ${{ inputs.aws_default_region }}
        SCCACHE_S3_BUCKET:  ${{ inputs.sccache_s3_bucket }}
        AWS_ACCESS_KEY_ID: ${{ inputs.aws_access_key_id }}
        AWS_SECRET_ACCESS_KEY: ${{ inputs.aws_secret_access_key }}
        PLATFORM: ${{ inputs.platform }}
        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
        GITHUB_RUN_ID: ${{ github.run_id }}
        GITHUB_JOB: ${{ github.job }}
        GITHUB_REF_NAME: ${{ github.ref_name }}
        CUDA_VERSION: ${{ inputs.cuda_version }}
      run: |
        set -x

        IMAGE_TAG="${{ inputs.image_tag }}"


        CUDA_VERSION_MAJOR=${CUDA_VERSION%%.*}

        BUILD_START_TIME=$(date -u +%Y-%m-%dT%H:%M:%SZ)
        echo "BUILD_START_TIME=${BUILD_START_TIME}" >> $GITHUB_ENV

        echo "image_tag=$IMAGE_TAG" >> $GITHUB_OUTPUT

        # Create build logs directory
        mkdir -p build-logs
        BUILD_LOG_FILE="build-logs/build-${{ inputs.framework }}-$(echo '${{ inputs.platform }}' | sed 's/linux\///').log"
        echo "BUILD_LOG_FILE=${BUILD_LOG_FILE}" >> $GITHUB_ENV
        echo "📝 Build log will be saved to: ${BUILD_LOG_FILE}"

        # Collect optional overrides provided by the workflow
        # Set base cache args and set --cache-to if this is a main commit
        # TODO: Fix this - Skip cache for frontend target - a different docker driver is used for the EPP build, which causes issues with cache export
        CACHE_ARGS=""
        if [[ "${{ inputs.target }}" != "frontend" ]]; then
          CACHE_ARGS="--cache-to type=inline "
          CACHE_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:${{ inputs.framework }}-cuda${CUDA_VERSION_MAJOR}-${PLATFORM##*/}-cache "
          CACHE_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${{ inputs.framework }}-cuda${CUDA_VERSION_MAJOR}-${PLATFORM##*/} "
          CACHE_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:release-${{ inputs.framework }}-cuda${CUDA_VERSION_MAJOR}-${PLATFORM##*/}-cache "
          if [[ "$GITHUB_REF_NAME" =~ ^release ]]; then
            # Release branches also use release cache
            CACHE_ARGS+="--cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:release-${{ inputs.framework }}-cuda${CUDA_VERSION_MAJOR}-${PLATFORM##*/}-cache,mode=max "
          elif [[ "$GITHUB_REF_NAME" == "main" ]]; then
            CACHE_ARGS+="--cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:${{ inputs.framework }}-cuda${CUDA_VERSION_MAJOR}-${PLATFORM##*/}-cache,mode=max "
          fi
        fi
        echo "$CACHE_ARGS"
        # Collect optional overrides provided by the workflow

        if [[ "${{ inputs.ci }}" == "true" ]]; then
          # CI mode for frontend: use existing buildx builder, push EPP to registry
          EXTRA_ARGS+=" --ci"
        fi

        if [ "${{ inputs.no_cache }}" == "true" ]; then
          EXTRA_ARGS+=" --no-cache"
        fi
        if [ "${{ inputs.use_sccache }}" == "true" ]; then
          EXTRA_ARGS+=" --build-arg CARGO_BUILD_JOBS=4 --use-sccache"
        fi
        if [ "${{ inputs.push_image }}" == "true" ]; then
          EXTRA_ARGS+=" --push"
        elif [ "${{ inputs.no_load }}" == "false" ]; then
          EXTRA_ARGS+=" --load"
        fi

        # Add extra tags (each as a separate --tag argument)
        EXTRA_TAGS="${{ inputs.extra_tags }}"
        if [ -n "$EXTRA_TAGS" ]; then
          while IFS= read -r EXTRA_TAG; do
            if [ -n "$EXTRA_TAG" ]; then
              EXTRA_ARGS+=" --tag ${EXTRA_TAG}"
            fi
          done <<< "$EXTRA_TAGS"
        fi

        docker buildx build \
          --progress=plain \
          --tag "$IMAGE_TAG" \
          --platform linux/${{ inputs.platform }} \
          -f ./container/rendered.Dockerfile \
          $CACHE_ARGS \
          $EXTRA_ARGS \
          $EPP_IMAGE_ARG . 2>&1 | tee "${BUILD_LOG_FILE}"

        BUILD_EXIT_CODE=${PIPESTATUS[0]}

        BUILD_END_TIME=$(date -u +%Y-%m-%dT%H:%M:%SZ)
        echo "BUILD_END_TIME=${BUILD_END_TIME}" >> $GITHUB_ENV

        # Exit with the build's exit code
        exit ${BUILD_EXIT_CODE}


    - name: Capture Build Metrics
      id: metrics
      shell: bash
      run: |

        # Create metrics directory
        mkdir -p build-metrics

        # Get accurate build timing
        BUILD_START_TIME="${{ env.BUILD_START_TIME }}"
        BUILD_END_TIME="${{ env.BUILD_END_TIME }}"

        # Calculate duration
        START_EPOCH=$(date -d "$BUILD_START_TIME" +%s)
        END_EPOCH=$(date -d "$BUILD_END_TIME" +%s)
        BUILD_DURATION_SEC=$((END_EPOCH - START_EPOCH))

        echo "🕐 Build timing:"
        echo "  Start: ${BUILD_START_TIME}"
        echo "  End: ${BUILD_END_TIME}"
        echo "  Duration: ${BUILD_DURATION_SEC} seconds"

        # Get image size using docker inspect
        IMAGE_TAG="${{ steps.build.outputs.image_tag }}"
        if [ -n "$IMAGE_TAG" ]; then
          IMAGE_SIZE_BYTES=$(docker image inspect "$IMAGE_TAG" --format='{{.Size}}' 2>/dev/null || echo "0")
          echo "📦 Image size: ${IMAGE_SIZE_BYTES} bytes"
        else
          IMAGE_SIZE_BYTES=0
          echo "⚠️  No image tag available"
        fi

        PLATFORM_ARCH=$(echo "${{ inputs.platform }}" | sed 's/linux\///')
        echo "  Architecture: ${PLATFORM_ARCH}"
        echo "PLATFORM_ARCH=${PLATFORM_ARCH}" >> $GITHUB_ENV
        JOB_KEY="${{ inputs.framework }}-${PLATFORM_ARCH}"
        echo "  Job Key: ${JOB_KEY}"

        # Create job-specific metrics file
        mkdir -p build-metrics
        METRICS_FILE="build-metrics/metrics-${{ inputs.framework }}-${PLATFORM_ARCH}-${{ github.run_id }}-${{ job.check_run_id }}.json"

        # Create the job metrics file
        cat > "$METRICS_FILE" << EOF
        {
          "framework": "${{ inputs.framework }}",
          "target": "${{ inputs.target }}",
          "platform": "${{ inputs.platform }}",
          "platform_arch": "${PLATFORM_ARCH}",
          "image_size_bytes": ${IMAGE_SIZE_BYTES},
          "build_start_time": "${BUILD_START_TIME}",
          "build_end_time": "${BUILD_END_TIME}",
          "build_duration_sec": ${BUILD_DURATION_SEC}
        }
        EOF

        cat "$METRICS_FILE"

    - name: Generate Comprehensive Build Metrics
      id: comprehensive-metrics
      if: always()
      shell: bash
      run: |
        echo "=========================================="
        echo "📊 GENERATING COMPREHENSIVE BUILD METRICS"
        echo "=========================================="

        # Create metrics directory
        mkdir -p build-metrics

        PLATFORM_ARCH="${{ env.PLATFORM_ARCH }}"
        WORKFLOW_ID="${{ github.run_id }}"
        JOB_ID="${{ job.check_run_id }}"
        FRAMEWORK_LOWER=$(echo "${{ inputs.framework }}" | tr '[:upper:]' '[:lower:]')

        # Make parser executable
        chmod +x .github/scripts/parse_buildkit_output.py

        # Check for build logs and build stage arguments dynamically
        # Use the BUILD_LOG_FILE set during the build step
        BUILD_LOG="${{ env.BUILD_LOG_FILE }}"

        # Path to container metadata created in previous step
        CONTAINER_METADATA="build-metrics/metrics-${{ inputs.framework }}-${PLATFORM_ARCH}-${WORKFLOW_ID}-${JOB_ID}.json"

        # Output single comprehensive JSON with all build stages
        COMPREHENSIVE_JSON="build-metrics/build-${{ inputs.framework }}-${PLATFORM_ARCH}-${WORKFLOW_ID}-${JOB_ID}.json"

        echo "🚀 Parsing BuildKit outputs and merging with container metrics..."

        # Build stage arguments dynamically based on which logs exist
        STAGE_ARGS=()

        if [ -f "$BUILD_LOG" ]; then
          echo "  ✓ Found base image log: ${BUILD_LOG}"
          STAGE_ARGS+=("runtime:${BUILD_LOG}")
        else
          echo "  ℹ️  No image log found"
        fi

        # Check for any additional stage logs (e.g., build-logs/stage3-*.log)
        for extra_log in build-logs/stage*.log; do
          if [ -f "$extra_log" ]; then
            stage_name=$(basename "$extra_log" .log)
            echo "  ✓ Found additional stage log: ${extra_log} (${stage_name})"
            STAGE_ARGS+=("${stage_name}:${extra_log}")
          fi
        done

        echo "Container Metadata: ${CONTAINER_METADATA}"
        echo "Output: ${COMPREHENSIVE_JSON}"
        echo ""

        # Run parser with all discovered stages
        # Usage: parse_buildkit_output.py <output_json> <stage1_name:log_file> [stage2_name:log_file] ... [--metadata=<file>]
        set +e
        python3 .github/scripts/parse_buildkit_output.py \
          "$COMPREHENSIVE_JSON" \
          "${STAGE_ARGS[@]}" \
          "--metadata=${CONTAINER_METADATA}"
        PARSER_EXIT_CODE=$?
        set -e

        echo ""
        echo "📊 Parser exit code: ${PARSER_EXIT_CODE}"

        if [ ${PARSER_EXIT_CODE} -eq 0 ] && [ -f "$COMPREHENSIVE_JSON" ]; then
          echo "✅ Comprehensive build metrics generated successfully"
          echo "📄 Output file: ${COMPREHENSIVE_JSON}"
        else
          echo "⚠️  Metrics generation had issues but continuing..."
        fi

    # Upload comprehensive build metrics as artifact
    - name: Upload Comprehensive Build Metrics
      uses: actions/upload-artifact@b7c566a772e6b6bfb58ed0dc250532a479d7789f #v6
      if: always()
      with:
        name: build-metrics-${{ inputs.framework }}-${{ inputs.target }}-${{ env.PLATFORM_ARCH }}-${{ github.run_id }}-${{ job.check_run_id }}
        path: build-metrics/build-${{ inputs.framework }}-${{ env.PLATFORM_ARCH }}-${{ github.run_id }}-${{ job.check_run_id }}.json
        retention-days: 7