name: 'Docker Build'
description: 'Build Dynamo container images'
inputs:
  # --- Common Docker 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 (optional, defaults to framework:latest)'
    required: false

  # --- Secret Inputs
  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

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:
        driver: docker-container
        # Enable BuildKit for enhanced metadata
        buildkitd-flags: --debug
        version: v0.14.1
    - name: Cleanup
      if: always()
      shell: bash
      run: |
        docker system prune -af
    - name: Set up Python
      uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 #v6.2.0
      with:
        python-version: '3.12'
        pip-install: jinja2 pyyaml
    - name: Generate Dockerfile
      shell: bash
      run: |
        echo "::group::Generating Dockerfile"
        echo "Generating Dockerfile for target: ${{ inputs.target }} and framework: ${{ inputs.framework }}"
        python ./container/render.py \
            --target=${{ inputs.target }} \
            --framework=${{ inputs.framework }} \
            --platform=${{ inputs.platform }} \
            --cuda-version=${{ inputs.cuda_version }} \
            --show-result \
            --output-short-filename
        echo "::endgroup::"
    - name: Build EPP image
      if: ${{ inputs.target == 'frontend' }}
      shell: bash
      env:
        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
      run: |
        sudo apt-get update && sudo apt-get install -y git build-essential protobuf-compiler libclang-dev
        curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain stable
        . "$HOME/.cargo/env"
        echo "$HOME/.cargo/bin" >> "$GITHUB_PATH"
        cargo install cbindgen

        DOCKER_PROXY="${ECR_HOSTNAME}/dockerhub/"
        pushd deploy/inference-gateway/epp
        make all DOCKER_PROXY=${DOCKER_PROXY}
        popd

        EPP_GIT_TAG=$(git describe --tags --dirty --always 2>/dev/null || echo "dev")
        EPP_IMAGE="dynamo/dynamo-epp:${EPP_GIT_TAG}"
        echo "EPP_IMAGE=${EPP_IMAGE}" >> $GITHUB_ENV
    - 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
        # Determine image tag
        if [ -n "${{ inputs.image_tag }}" ]; then
          IMAGE_TAG="${{ inputs.image_tag }}"
        else
          IMAGE_TAG="${{ inputs.framework }}:latest"
        fi

        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}"

        # 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

        EPP_IMAGE_ARG=""
        if [[ ${{ inputs.target }} == "frontend" ]]; then
          EPP_IMAGE_ARG="--build-arg EPP_IMAGE=${EPP_IMAGE}"
        fi

        docker buildx build \
          --progress=plain \
          --tag "$IMAGE_TAG" \
          --load \
          -f ./container/rendered.Dockerfile \
          $CACHE_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: Run Sanity Check on Runtime Image
      if: inputs.target == 'runtime'
      shell: bash
      run: |
        IMAGE_TAG="${{ steps.build.outputs.image_tag }}"
        echo "Running sanity check on image: $IMAGE_TAG"

        # Run the sanity check script inside the container
        # The script is located in /workspace/deploy/sanity_check.py in runtime containers
        export WORKSPACE=/workspace

        set +e
        docker run --rm "$IMAGE_TAG" python ${WORKSPACE}/deploy/sanity_check.py --runtime-check --no-gpu-check
        SANITY_CHECK_EXIT_CODE=$?
        set -e
        if [ ${SANITY_CHECK_EXIT_CODE} -ne 0 ]; then
          echo "ERROR: Sanity check failed - ai-dynamo packages not properly installed"
          exit ${SANITY_CHECK_EXIT_CODE}
        else
          echo "✅ Sanity check passed"
        fi

    - 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
        BUILD_LOG="build-logs/single-stage-build.log"

        # 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@v4
      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
