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: 'CUDA version to use'
    required: true
  image_tag:
    description: 'Custom image tag'
    required: true
  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: 'true'
  extra_build_args:
    description: 'Additional Docker build args (newline-separated list of KEY=VALUE pairs)'
    required: false
    default: ''
outputs:
  image_tag:
    description: 'Image Tag'
    value: ${{ steps.build.outputs.image_tag }}

runs:
  using: "composite"
  steps:
    - name: Build image
      id: build
      shell: bash
      env:
        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 }}"

        # Normalize platform: prepend linux/ only when not already present
        PLATFORM_FLAG="${PLATFORM}"
        if [[ "$PLATFORM_FLAG" != linux/* ]]; then
          PLATFORM_FLAG="linux/${PLATFORM_FLAG}"
        fi
        # Sanitized platform for registry tags and filenames (no slashes or commas)
        PLATFORM_SAFE=$(echo "${PLATFORM}" | sed 's|linux/||g; s/,/-/g')

        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 }}-${PLATFORM_SAFE}.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
        CACHE_ARGS=""
        if [[ "${{ inputs.target }}" != "frontend" ]]; then
          CACHE_TAG="${{ inputs.framework }}-cuda${CUDA_VERSION_MAJOR}-${PLATFORM_SAFE}-cache"
          CACHE_ARGS="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${CACHE_TAG} "
          CACHE_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:release-${CACHE_TAG} "
          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-${CACHE_TAG},mode=max "
          elif [[ "$GITHUB_REF_NAME" == "main" ]]; then
            CACHE_ARGS+="--cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${CACHE_TAG},mode=max "
          fi
        else
          CACHE_TAG="${{ inputs.target }}-cuda${CUDA_VERSION_MAJOR}-${PLATFORM_SAFE}-cache"
          CACHE_ARGS="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${CACHE_TAG} "
          if [[ "$GITHUB_REF_NAME" == "main" ]]; then
            CACHE_ARGS+="--cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${CACHE_TAG},mode=max "
          fi
        fi

        echo "$CACHE_ARGS"
        # Collect optional overrides provided by the workflow
        if [ "${{ inputs.no_cache }}" == "true" ]; then
          EXTRA_ARGS+=" --no-cache"
        fi
        if [ "${{ inputs.use_sccache }}" == "true" ]; then
          EXTRA_ARGS+=" --build-arg USE_SCCACHE=true"
          EXTRA_ARGS+=" --build-arg SCCACHE_BUCKET=${SCCACHE_S3_BUCKET}"
          EXTRA_ARGS+=" --build-arg SCCACHE_REGION=${AWS_DEFAULT_REGION}"
        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

        # Add extra build args (each as a separate --build-arg)
        EXTRA_BUILD_ARGS="${{ inputs.extra_build_args }}"
        if [ -n "$EXTRA_BUILD_ARGS" ]; then
          while IFS= read -r BUILD_ARG; do
            if [ -n "$BUILD_ARG" ]; then
              EXTRA_ARGS+=" --build-arg ${BUILD_ARG}"
            fi
          done <<< "$EXTRA_BUILD_ARGS"
        fi

        # Pass AWS credentials as build secrets for sccache S3 access.
        # Dockerfile steps reference these via --mount=type=secret,id=aws-key-id,env=...
        SECRET_ARGS=""
        if [ "${{ inputs.use_sccache }}" == "true" ] && [ -n "${AWS_ACCESS_KEY_ID:-}" ]; then
          SECRET_ARGS+=" --secret id=aws-key-id,env=AWS_ACCESS_KEY_ID"
          SECRET_ARGS+=" --secret id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY"
        fi

        docker buildx build \
          --progress=plain \
          --tag "$IMAGE_TAG" \
          --platform "${PLATFORM_FLAG}" \
          -f ./container/rendered.Dockerfile \
          $CACHE_ARGS $EXTRA_ARGS $SECRET_ARGS . 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
      env:
        AWS_DEFAULT_REGION: ${{ inputs.aws_default_region }}
      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
        IMAGE_TAG="${{ steps.build.outputs.image_tag }}"
        IMAGE_SIZE_TYPE="unknown"
        if [ -n "$IMAGE_TAG" ]; then
          if [ "${{ inputs.no_load }}" == "false" ]; then
            # Image was loaded into local Docker daemon but not pushed; docker inspect
            # reports uncompressed (virtual) size, which differs from the compressed
            # registry size reported below
            IMAGE_SIZE_BYTES=$(docker image inspect "$IMAGE_TAG" --format='{{.Size}}' 2>/dev/null || echo "0")
            IMAGE_SIZE_TYPE="uncompressed"
          elif [ "${{ inputs.push_image }}" == "true" ]; then
            # Image was pushed to ECR without loading locally; query ECR for compressed size
            REPO_NAME=$(echo "$IMAGE_TAG" | sed 's|.*amazonaws.com/||' | cut -d: -f1)
            IMAGE_TAG_NAME=$(echo "$IMAGE_TAG" | cut -d: -f2-)
            IMAGE_SIZE_BYTES=$(aws ecr describe-images \
              --repository-name "$REPO_NAME" \
              --image-ids "imageTag=${IMAGE_TAG_NAME}" \
              --query 'imageDetails[0].imageSizeInBytes' \
              --output text 2>/dev/null || echo "0")
            # Treat empty or "None" response as 0
            if [[ "$IMAGE_SIZE_BYTES" == "None" || -z "$IMAGE_SIZE_BYTES" ]]; then
              IMAGE_SIZE_BYTES=0
            fi
            IMAGE_SIZE_TYPE="compressed"
          else
            IMAGE_SIZE_BYTES=0
          fi
          echo "📦 Image size: ${IMAGE_SIZE_BYTES} bytes (${IMAGE_SIZE_TYPE})"
        else
          IMAGE_SIZE_BYTES=0
          echo "⚠️  No image tag available"
        fi

        PLATFORM_ARCH=$(echo "${{ inputs.platform }}" | sed 's|linux/||g; s/,/-/g')
        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},
          "image_size_type": "${IMAGE_SIZE_TYPE}",
          "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

