Android studio rename app bundle and apk file based on flavours

Question

I am working on an app that has flavours like SIT, UAT and Prod. I want to rename the aab and apk files in the following format:

project-name_versionCode_versionName-flavour.aab or apk

I added a method in the app/build.gradle to get the versionCodes depending on flavour

def static getBuildNo(String variant) {
    def SIT = 111
    def UAT = 2
    def PROD = 3

    if(variant.equalsIgnoreCase("SIT")) { return SIT }
    if(variant.equalsIgnoreCase("UAT")) { return UAT }
    if(variant.equalsIgnoreCase("PROD")) { return PROD }
}

def static fetchVersion() {
    return "1.1.0"
}

And to rename the files, I added this method for all the flavours. For eg

productFlavors {

  sit {
       versionCode getBuildNo("SIT")
       setProperty("archivesBaseName", "MyApp_${getBuildNo("sit")}_${fetchVersion()}")
      }
}

But instead of printing the file name as

MyApp_111_1.1.0-sit-release.aab

it keeps printing

MyApp_3_1.1.0-sit-release.aab

Somehow it keeps taking up the PROD build no instead of SIT.

IS there a better way to rename aab and apk files with flavours?

I am using Android Studio Electric Eel and Gradle 7.5

I tried to write gradle methods to rename. But the renaming is not fully correct as it keeps taking the wrong build number