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'
  image_tag:
    description: 'Custom image tag (optional, defaults to framework:latest)'
    required: false
  ngc_ci_access_token:
    description: 'NGC CI Access Token'
    required: false
  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
  base_image_tag:
    description: 'Optional override for base image tag passed to build.sh'
    required: false
  runtime_image_tag:
    description: 'Optional override for RUNTIME_IMAGE_TAG build-arg'
    required: false
  cuda_version:
    description: 'Optional override for CUDA_VERSION build-arg'
    required: false
  torch_backend:
    description: 'Optional override for TORCH_BACKEND build-arg (e.g., cu129)'
    required: false

outputs:
  image_tag:
    description: 'Image Tag'
    value: ${{ steps.build.outputs.image_tag }}

runs:
  using: "composite"
  steps:
    - name: Set up Docker Buildx
      uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 #v3.11.1
      with:
        # This is not # of CPUs, this is CPU priority, 1024 is the default
        driver-opts: 'cpu-shares=100'
    - name: Install awscli
      shell: bash
      run: |
        curl "https://awscli.amazonaws.com/awscli-exe-linux-$(uname -m).zip" -o "awscliv2.zip"
        unzip awscliv2.zip
        sudo ./aws/install
    - name: Login to ECR
      shell: bash
      env:
        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
      run: |
        aws ecr get-login-password --region ${{ inputs.aws_default_region }} | docker login --username AWS --password-stdin ${ECR_HOSTNAME}
    - name: Login to NGC
      if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name == 'push'
      shell: bash
      run: |
        echo "${{ inputs.ngc_ci_access_token }}" | docker login nvcr.io -u '$oauthtoken' --password-stdin
    - name: Cleanup
      if: always()
      shell: bash
      run: |
        docker system prune -af
    - 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 }}
      run: |
        # Determine image tag
        if [ -n "${{ inputs.image_tag }}" ]; then
          IMAGE_TAG="${{ inputs.image_tag }}"
        else
          IMAGE_TAG="${{ inputs.framework }}:latest"
        fi

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

        echo "image_tag=$IMAGE_TAG" >> $GITHUB_OUTPUT
        # Collect optional overrides provided by the workflow
        EXTRA_ARGS=""
        if [ -n "${{ inputs.base_image_tag }}" ]; then
          EXTRA_ARGS+=" --base-image-tag ${{ inputs.base_image_tag }}"
        fi
        if [ -n "${{ inputs.runtime_image_tag }}" ]; then
          EXTRA_ARGS+=" --build-arg RUNTIME_IMAGE_TAG=${{ inputs.runtime_image_tag }}"
        fi
        if [ -n "${{ inputs.cuda_version }}" ]; then
          EXTRA_ARGS+=" --build-arg CUDA_VERSION=${{ inputs.cuda_version }}"
        fi
        if [ -n "${{ inputs.torch_backend }}" ]; then
          EXTRA_ARGS+=" --build-arg TORCH_BACKEND=${{ inputs.torch_backend }}"
        fi

        ./container/build.sh --tag "$IMAGE_TAG" \
          --target ${{ inputs.target }} \
          --vllm-max-jobs 10 \
          --framework ${{ inputs.framework }} \
          --platform ${{ inputs.platform }} \
          --use-sccache \
          --sccache-bucket "$SCCACHE_S3_BUCKET" \
          --sccache-region "$AWS_DEFAULT_REGION" $EXTRA_ARGS

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

    - name: Capture Build Metrics
      id: metrics
      shell: bash
      run: |
        echo "📊 Capturing build metrics for ${{ inputs.framework }}..."

        # 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

        echo "📊 Final metrics captured"

        # Create consolidated metrics JSON file
        echo "🔍 Debug: inputs.platform = '${{ inputs.platform }}'"
        PLATFORM_ARCH=$(echo "${{ inputs.platform }}" | sed 's/linux\///')
        echo "🔍 Debug: PLATFORM_ARCH = '${PLATFORM_ARCH}'"
        echo "PLATFORM_ARCH=${PLATFORM_ARCH}" >> $GITHUB_ENV
        JOB_KEY="${{ inputs.framework }}-${PLATFORM_ARCH}"
        echo "🔍 Debug: JOB_KEY = '${JOB_KEY}'"

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

        # Create the job metrics file directly
        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

        echo "📁 Created build metrics file for ${JOB_KEY}:"
        cat "$METRICS_FILE"

        # Metrics captured and saved to JSON file

    # Upload job-specific build metrics as artifact
    - name: Upload Build Metrics
      uses: actions/upload-artifact@v4
      with:
        name: build-metrics-${{ inputs.framework }}-${{ env.PLATFORM_ARCH }}
        path: build-metrics/metrics-${{ inputs.framework }}-${{ env.PLATFORM_ARCH }}.json
        retention-days: 7