parameters:
- name: gpuTarget
  type: string
  default: ''

steps:
- task: Bash@3
  name: downloadCKBuild
  displayName: Download specific CK build
  retryCountOnTaskFailure: 3
  env:
    CXX: $(Agent.BuildDirectory)/rocm/llvm/bin/amdclang++
    CC: $(Agent.BuildDirectory)/rocm/llvm/bin/amdclang
  inputs:
    targetType: inline
    workingDirectory: $(Agent.BuildDirectory)/s
    script: |
      AZ_API="https://dev.azure.com/ROCm-CI/ROCm-CI/_apis"
      GH_API="https://api.github.com/repos/ROCm"
      EXIT_CODE=0

      # Try to find an Azure build for the specific CK commit called out in MIOpen's requirements.txt
      CK_COMMIT=$(grep 'ROCm/composable_kernel' requirements.txt | sed -E 's/.*@([a-f0-9]{40}).*/\1/')
      echo "Fetching CK build ID for commit $CK_COMMIT"
      CK_CHECKS_URL="$GH_API/composable_kernel/commits/${CK_COMMIT}/check-runs"
      CK_BUILD_ID=$(curl -s $CK_CHECKS_URL | \
        jq '.check_runs[] | select(.name == "composable_kernel" and .app.slug == "azure-pipelines" and .conclusion == "success") | .details_url' | \
        tr -d '"' | grep -oP 'buildId=\K\d+')

      # If none found, use latest successful CK build instead
      if [[ -z "$CK_BUILD_ID" ]]; then
        echo "Did not find specific CK build ID"
        LATEST_BUILD_URL="$AZ_API/build/builds?definitions=$(COMPOSABLE_KERNEL_PIPELINE_ID)&statusFilter=completed&resultFilter=succeeded&\$top=1&api-version=7.1"
        CK_BUILD_ID=$(curl -s $LATEST_BUILD_URL | jq '.value[0].id')
        echo "Found latest CK build ID: $CK_BUILD_ID"
        EXIT_CODE=1
      else
        echo "Found specific CK build ID: $CK_BUILD_ID"
      fi

      AZURE_URL="$AZ_API/build/builds/$CK_BUILD_ID/artifacts?api-version=7.1"
      ARTIFACT_URL=$(curl -s $AZURE_URL | \
        jq --arg gfx "${{ parameters.gpuTarget }}" '
          .value
          | map(select(.name | test($gfx)))
          | max_by(.name | capture("_(?<dropNumber>\\d+)").dropNumber | tonumber)
          | .resource.downloadUrl
        ' | \
        tr -d '"')

      # If using the specific CK commit and it doesn't have any valid artifacts, use latest successful CK build instead
      if { [[ -z "$ARTIFACT_URL" ]] || [[ "$ARTIFACT_URL" == "null" ]]; } && [[ $EXIT_CODE -eq 0 ]]; then
        echo "Did not find valid specific CK build artifact"
        LATEST_BUILD_URL="$AZ_API/build/builds?definitions=$(COMPOSABLE_KERNEL_PIPELINE_ID)&statusFilter=completed&resultFilter=succeeded&\$top=1&api-version=7.1"
        CK_BUILD_ID=$(curl -s $LATEST_BUILD_URL | jq '.value[0].id')
        echo "Found latest CK build ID: $CK_BUILD_ID"
        AZURE_URL="$AZ_API/build/builds/$CK_BUILD_ID/artifacts?api-version=7.1"
        ARTIFACT_URL=$(curl -s $AZURE_URL | \
          jq --arg os "ubuntu2204" --arg gfx "${{ parameters.gpuTarget }}" '
            .value
            | map(select(.name | test($os) and test($gfx)))
            | max_by(.name | capture("_(?<dropNumber>\\d+)").dropNumber | tonumber)
            | .resource.downloadUrl
          ' | \
          tr -d '"')
        EXIT_CODE=2
      fi

      echo "Downloading CK artifact from $ARTIFACT_URL"

      RETRIES=0
      MAX_RETRIES=5
      SUCCESS=false
      while [ $RETRIES -lt $MAX_RETRIES ]; do
        wget -nv $ARTIFACT_URL -O $(System.ArtifactsDirectory)/ck.zip && \
          unzip $(System.ArtifactsDirectory)/ck.zip -d $(System.ArtifactsDirectory) && \
          mkdir -p $(Agent.BuildDirectory)/rocm && \
          tar -zxvf $(System.ArtifactsDirectory)/composable_kernel*/*.tar.gz -C $(Agent.BuildDirectory)/rocm && \
          rm -r $(System.ArtifactsDirectory)/ck.zip $(System.ArtifactsDirectory)/composable_kernel*

        if [ $? -eq 0 ]; then
          SUCCESS=true
          echo "Successfully downloaded CK."
          break
        else
          RETRIES=$((RETRIES + 1))
          echo "Failed to download CK on attempt $RETRIES/$MAX_RETRIES, retrying..."
          sleep 1
        fi
      done

      if [ "$SUCCESS" = false ]; then
        echo "ERROR: failed to download CK after $MAX_RETRIES attempts."
        exit 1
      fi

      if [[ $EXIT_CODE -ne 0 ]]; then
        BUILD_COMMIT=$(curl -s $AZ_API/build/builds/$CK_BUILD_ID | jq '.sourceVersion' | tr -d '"')
        if [[ $EXIT_CODE -eq 1 ]]; then
          echo "WARNING: couldn't find a CK build for commit $CK_COMMIT"
        elif [[ $EXIT_CODE -eq 2 ]]; then
          echo "WARNING: couldn't find a valid CK artifact for commit $CK_COMMIT"
        fi
        echo "Instead used latest CK build $CK_BUILD_ID for commit $BUILD_COMMIT"
      fi
