build.gradle

/*
 * This file was generated by the Gradle 'init' task.
 *
 * This is a general purpose Gradle build.
 * Learn more about Gradle by exploring our samples at https://docs.gradle.org/7.5.1/samples
 * This project uses @Incubating APIs which are subject to change.
 */

import org.opensearch.gradle.test.RestIntegTestTask
import org.opensearch.gradle.testclusters.OpenSearchCluster

import java.nio.file.Paths
import java.util.concurrent.Callable

buildscript {
    ext {
        opensearch_version = System.getProperty("opensearch.version", "3.7.1-SNAPSHOT")
        buildVersionQualifier = System.getProperty("build.version_qualifier", "")
        isSnapshot = "true" == System.getProperty("build.snapshot", "true")
        version_tokens = opensearch_version.tokenize('-')
        opensearch_build = version_tokens[0] + '.0'
        plugin_no_snapshot = opensearch_build
        if (buildVersionQualifier) {
            opensearch_build += "-${buildVersionQualifier}"
            plugin_no_snapshot += "-${buildVersionQualifier}"
        }
        if (isSnapshot) {
            opensearch_build += "-SNAPSHOT"
        }
        opensearch_group = "org.opensearch"
        opensearch_no_snapshot = opensearch_build.replace("-SNAPSHOT","")
    }

    repositories {
        mavenLocal()
        maven { url "https://ci.opensearch.org/ci/dbc/snapshots/maven/" }
        maven { url "https://ci.opensearch.org/maven2/" }
        mavenCentral()
        maven { url "https://plugins.gradle.org/m2/" }
    }

    dependencies {
        classpath "${opensearch_group}.gradle:build-tools:${opensearch_version}"
        classpath "com.diffplug.spotless:spotless-plugin-gradle:6.25.0"
        classpath "io.freefair.gradle:lombok-plugin:8.14"

        configurations.all {
            resolutionStrategy {
                force("org.eclipse.platform:org.eclipse.core.runtime:4.29.0") // CVE for < 4.29
                force("org.eclipse.platform:org.eclipse.core.resources:4.20.0") // CVE for < 4.20
            }
        }
    }
}

plugins{
    id 'eclipse'
    id "de.undercouch.download" version "5.3.0"
}
apply plugin: 'java'
apply plugin: 'java-library'
apply plugin: 'java-test-fixtures'
apply plugin: 'idea'
apply plugin: 'opensearch.opensearchplugin'
apply plugin: 'opensearch.pluginzip'
apply plugin: 'jacoco'
apply plugin: "com.diffplug.spotless"
apply plugin: 'io.freefair.lombok'
apply from: 'formatter/formatting.gradle'
apply plugin: 'opensearch.java-agent'

ext.opensearch_tmp_dir = rootProject.file('build/private/opensearch_tmp').absoluteFile
opensearch_tmp_dir.mkdirs()

ext {
    default_bwc_version = System.getProperty("bwc.version")
    default_bwc_bundle_version= System.getProperty("bwc.bundle.version")
    bwcBundleTest = (project.findProperty('customDistributionDownloadType') != null && project.properties['customDistributionDownloadType'] == "bundle")
    neural_search_bwc_version = bwcBundleTest ? System.getProperty("tests.bwc.bundle.version",rootProject.ext.default_bwc_bundle_version): System.getProperty("tests.bwc.version", rootProject.ext.default_bwc_version)
    currentBundleVersion = opensearch_version.replace("-SNAPSHOT","")

    // Config below including files are copied from security demo configuration
    runIntegTestWithSecurityPlugin= System.getProperty("security.enabled")
    if (runIntegTestWithSecurityPlugin == "true"){
        ['esnode.pem', 'esnode-key.pem', 'root-ca.pem','kirk.pem','kirk-key.pem'].forEach { file ->
            File local = Paths.get(opensearch_tmp_dir.absolutePath, file).toFile()
            download.run {
                src "https://raw.githubusercontent.com/opensearch-project/security/main/bwc-test/src/test/resources/security/" + file
                dest local
                overwrite false
            }
        }
    }

    isSnapshot = "true" == System.getProperty("build.snapshot", "true")
    projectSubstitutions = [:]

    configureSecurityPlugin = { OpenSearchCluster cluster ->
        configurations.secureIntegTestPluginArchive.asFileTree.each {
            if(it.name.contains("opensearch-security")){
                cluster.plugin(provider(new Callable<RegularFile>() {
                    @Override
                    RegularFile call() throws Exception {
                        return new RegularFile() {
                            @Override
                            File getAsFile() {
                                return it
                            }
                        }
                    }
                }))
            }
        }

        cluster.getNodes().forEach { node ->
            var creds = node.getCredentials()
            if (creds.isEmpty()) {
                creds.add(Map.of('username', 'admin', 'password', 'admin'))
            } else {
                creds.get(0).putAll(Map.of('username', 'admin', 'password', 'admin'))
            }
        }

        // // Config below including files are copied from security demo configuration
        cluster.extraConfigFile("esnode.pem", file("$opensearch_tmp_dir/esnode.pem"))
        cluster.extraConfigFile("esnode-key.pem", file("$opensearch_tmp_dir/esnode-key.pem"))
        cluster.extraConfigFile("root-ca.pem", file("$opensearch_tmp_dir/root-ca.pem"))

        // This configuration is copied from the security plugins demo install:
        // https://github.com/opensearch-project/security/blob/2.11.1.0/tools/install_demo_configuration.sh#L365-L388
        cluster.setting("plugins.security.ssl.transport.pemcert_filepath", "esnode.pem")
        cluster.setting("plugins.security.ssl.transport.pemkey_filepath", "esnode-key.pem")
        cluster.setting("plugins.security.ssl.transport.pemtrustedcas_filepath", "root-ca.pem")
        cluster.setting("plugins.security.ssl.transport.enforce_hostname_verification", "false")
        cluster.setting("plugins.security.ssl.http.enabled", "true")
        cluster.setting("plugins.security.ssl.http.pemcert_filepath", "esnode.pem")
        cluster.setting("plugins.security.ssl.http.pemkey_filepath", "esnode-key.pem")
        cluster.setting("plugins.security.ssl.http.pemtrustedcas_filepath", "root-ca.pem")
        cluster.setting("plugins.security.allow_unsafe_democertificates", "true")
        cluster.setting("plugins.security.allow_default_init_securityindex", "true")
        cluster.setting("plugins.security.unsupported.inject_user.enabled", "true")

        cluster.setting("plugins.security.authcz.admin_dn", "\n- CN=kirk,OU=client,O=client,L=test, C=de")
        cluster.setting('plugins.security.restapi.roles_enabled', '["all_access", "security_rest_api_access"]')
        cluster.setting('plugins.security.system_indices.enabled', "true")
        cluster.setting('plugins.security.system_indices.indices', '[' +
                '".plugins-ml-config", ' +
                '".plugins-ml-connector", ' +
                '".plugins-ml-model-group", ' +
                '".plugins-ml-model", ".plugins-ml-task", ' +
                '".plugins-ml-conversation-meta", ' +
                '".plugins-ml-conversation-interactions", ' +
                ']'
        )
        cluster.setSecure(true)
    }
}

allprojects {
    group = opensearch_group
    version = "${opensearch_build}"

    apply from: rootProject.file('repositories.gradle').absoluteFile
    plugins.withId('java') {
        java {
            sourceCompatibility = JavaVersion.VERSION_21
            targetCompatibility = JavaVersion.VERSION_21
        }
    }

    configurations.all {
        resolutionStrategy {
            force("com.google.errorprone:error_prone_annotations:2.21.1")
        }
    }

    afterEvaluate {
        project.dependencyLicenses.enabled = false
        project.thirdPartyAudit.enabled = false
        project.loggerUsageCheck.enabled = false
        project.forbiddenApis.ignoreFailures = false
        project.forbiddenPatterns {
            setEnabled(false)
        }
        project.testingConventions.enabled = false
        project.validateNebulaPom.enabled = false
        project.licenseFile = rootProject.file('LICENSE.txt')
        project.noticeFile = rootProject.file('NOTICE.txt')
    }
}

configurations {
    zipArchive
    secureIntegTestPluginArchive
}

tasks.register("preparePluginPathDirs") {
    mustRunAfter clean
    doLast {
        def newPath = pathToPlugin.replace(".", "/")
        mkdir "src/main/java/org/opensearch/$newPath"
        mkdir "src/test/java/org/opensearch/$newPath"
    }
}

def pluginName = 'opensearch-neural-search'
def pluginDescription = 'A plugin that adds dense neural retrieval into the OpenSearch ecosytem'
def projectPath = 'org.opensearch'
def pathToPlugin = 'neuralsearch.plugin'
def pluginClassName = 'NeuralSearch'

publishing {
    repositories {
        maven {
            name = "Snapshots"
            url = System.getenv("MAVEN_SNAPSHOTS_S3_REPO")
            credentials(AwsCredentials) {
                accessKey = System.getenv("AWS_ACCESS_KEY_ID")
                secretKey = System.getenv("AWS_SECRET_ACCESS_KEY")
                sessionToken = System.getenv("AWS_SESSION_TOKEN")
            }
        }
    }
    publications {
        pluginZip(MavenPublication) { publication ->
            pom {
                name = pluginName
                description = pluginDescription
                groupId = "org.opensearch.plugin"
                licenses {
                    license {
                        name = "The Apache License, Version 2.0"
                        url = "http://www.apache.org/licenses/LICENSE-2.0.txt"
                    }
                }
                developers {
                    developer {
                        name = "OpenSearch"
                        url = "https://github.com/opensearch-project/neural-search"
                    }
                }
            }
        }
    }
}
opensearchplugin {
    name pluginName
    description pluginDescription
    classname "${projectPath}.${pathToPlugin}.${pluginClassName}"
    licenseFile rootProject.file('LICENSE')
    noticeFile rootProject.file('NOTICE')
    extendedPlugins = ['opensearch-knn', 'transport-grpc']
}

dependencyLicenses.enabled = false
thirdPartyAudit.enabled = false
loggerUsageCheck.enabled = false
// No need to validate pom, as we do not upload to maven/sonatype
validateNebulaPom.enabled = false

def knnJarDirectory = "$buildDir/dependencies/opensearch-knn"

dependencies {
    api "org.opensearch:opensearch:${opensearch_version}"
    implementation group: 'org.opensearch.plugin', name:'mapper-extras-client', version: "${opensearch_version}"
    zipArchive group: 'org.opensearch.plugin', name:'opensearch-job-scheduler', version: "${opensearch_build}"
    zipArchive group: 'org.opensearch.plugin', name:'opensearch-knn', version: "${opensearch_build}"
    zipArchive group: 'org.opensearch.plugin', name:'opensearch-ml-plugin', version: "${opensearch_build}"
    secureIntegTestPluginArchive group: 'org.opensearch.plugin', name:'opensearch-security', version: "${opensearch_build}"
    compileOnly fileTree(dir: knnJarDirectory, include: ["opensearch-knn-${opensearch_build}.jar", "remote-index-build-client-${opensearch_build}.jar"])
    compileOnly "org.opensearch.plugin:transport-grpc-spi:${opensearch_version}"
    compileOnly "org.opensearch:protobufs:${versions.opensearchprotobufs}"
    compileOnly (group: 'com.google.guava', name: 'guava', version:"${versions.guava}") {
        exclude group: 'com.google.errorprone', module: 'error_prone_annotations'
    }
    api group: 'org.opensearch', name:'opensearch-ml-client', version: "${opensearch_build}"
    testFixturesImplementation "org.opensearch.test:framework:${opensearch_version}"
    compileOnly group: 'org.apache.commons', name: 'commons-lang3', version: "${versions.commonslang}"
    implementation group: 'org.apache.commons', name: 'commons-math3', version: '3.6.1'
    implementation group: 'com.google.code.gson', name: 'gson', version: '2.10.1'
    implementation 'org.apache.commons:commons-collections4:4.5.0'
    implementation 'com.googlecode.concurrentlinkedhashmap:concurrentlinkedhashmap-lru:1.4.2'
    // ml-common excluded reflection for runtime so we need to add it by ourselves.
    // https://github.com/opensearch-project/ml-commons/commit/464bfe34c66d7a729a00dd457f03587ea4e504d9
    // TODO: Remove following three lines of dependencies if ml-common include them in their jar
    runtimeOnly group: 'org.reflections', name: 'reflections', version: '0.9.12'
    runtimeOnly group: 'org.javassist', name: 'javassist', version: '3.29.2-GA'
    runtimeOnly group: 'org.opensearch', name: 'common-utils', version: "${opensearch_build}"
    runtimeOnly group: 'org.apache.commons', name: 'commons-text', version: '1.10.0'
    runtimeOnly group: 'org.json', name: 'json', version: '20231013'
    // json-path 2.9.0 depends on slf4j 2.0.11, which conflicts with the version used by OpenSearch core.
    // Excluding slf4j here since json-path is only used for testing, and logging failures in this context are acceptable.
    runtimeOnly('com.jayway.jsonpath:json-path:2.9.0') {
        // OpenSearch core is using slf4j 1.7.36. Therefore, we cannot change the version here.
        exclude group: 'org.slf4j', module: 'slf4j-api'
        exclude group: 'net.minidev', module: 'json-smart'
    }
    runtimeOnly group: 'net.minidev', name:'json-smart', version: "${versions.json_smart}"
    implementation("com.fasterxml.jackson.core:jackson-annotations:${versions.jackson_annotations}")
    implementation("tools.jackson.core:jackson-databind:${versions.jackson3_databind}")
    testFixturesImplementation "org.opensearch:common-utils:${version}"
    testFixturesImplementation group: 'org.apache.commons', name: 'commons-lang3', version: "${versions.commonslang}"
    testFixturesCompileOnly (group: 'com.google.guava', name: 'guava', version:"${versions.guava}") {
        exclude group: 'com.google.errorprone', module: 'error_prone_annotations'
    }
    testFixturesImplementation fileTree(dir: knnJarDirectory, include: ["opensearch-knn-${opensearch_build}.jar", "remote-index-build-client-${opensearch_build}.jar"])
    testImplementation fileTree(dir: knnJarDirectory, include: ["opensearch-knn-${opensearch_build}.jar", "remote-index-build-client-${opensearch_build}.jar"])
    testImplementation "org.opensearch.plugin:parent-join-client:${opensearch_version}"
    testImplementation 'org.assertj:assertj-core:3.27.7'
    testImplementation group: 'net.bytebuddy', name: 'byte-buddy', version: "${versions.bytebuddy}"
    testImplementation group: 'org.objenesis', name: 'objenesis', version: "${versions.objenesis}"
    testImplementation group: 'net.bytebuddy', name: 'byte-buddy-agent', version: "${versions.bytebuddy}"
    // gRPC dependencies for integration tests
    testImplementation "org.opensearch:protobufs:${versions.opensearchprotobufs}"
    testImplementation "org.opensearch.plugin:transport-grpc-spi:${opensearch_version}"
    testImplementation "io.grpc:grpc-api:${versions.grpc}"
    testImplementation "io.grpc:grpc-core:${versions.grpc}"
    testImplementation "io.grpc:grpc-netty-shaded:${versions.grpc}"
    testImplementation "io.grpc:grpc-stub:${versions.grpc}"
    testImplementation "io.grpc:grpc-protobuf:${versions.grpc}"
    testImplementation "io.grpc:grpc-protobuf-lite:${versions.grpc}"
    testRuntimeOnly "io.perfmark:perfmark-api:0.27.0"
    testRuntimeOnly "com.google.guava:guava:${versions.guava}"
}

// In order to add the jar to the classpath, we need to unzip the
// k-NN zip and then copy it into a directory that we specify as a dependency.
task extractKnnJar(type: Copy) {
    mustRunAfter()
    from(zipTree(configurations.zipArchive.find { it.name.startsWith("opensearch-knn")}))
    into knnJarDirectory
}

project.tasks.delombok.dependsOn(extractKnnJar)

compileJava {
    dependsOn extractKnnJar
    dependsOn spotlessApply
    options.compilerArgs.addAll(["-processor", 'lombok.launch.AnnotationProcessorHider$AnnotationProcessor'])
}
compileTestJava {
    options.compilerArgs.addAll(["-processor", 'lombok.launch.AnnotationProcessorHider$AnnotationProcessor'])
}
compileTestFixturesJava {
    options.compilerArgs.addAll(["-processor", 'lombok.launch.AnnotationProcessorHider$AnnotationProcessor'])
}

def _numNodes = findProperty('numNodes') as Integer ?: 1
def _configureNodeRoles = findProperty('configureNodeRoles') as Boolean ?: false

test {
    include '**/*Tests.class'
    systemProperty 'tests.security.manager', 'false'
    systemProperty "jdk.attach.allowAttachSelf", true
}

// Setting up Integration Tests
task integTest(type: RestIntegTestTask) {
    description = "Run tests against a cluster"
    testClassesDirs = sourceSets.test.output.classesDirs
    classpath = sourceSets.test.runtimeClasspath
    boolean runCompleteAggsTestSuite = Boolean.parseBoolean(System.getProperty('test_aggs', "false"))
    if (!runCompleteAggsTestSuite) {
        filter {
            excludeTestsMatching "org.opensearch.neuralsearch.query.aggregation.*IT"
        }
    }
}
tasks.named("check").configure { dependsOn(integTest) }

integTest {
    systemProperty 'tests.security.manager', 'false'
    systemProperty 'java.io.tmpdir', opensearch_tmp_dir.absolutePath
    // allows integration test classes to access test resource from project root path
    systemProperty('project.root', project.rootDir.absolutePath)
    systemProperty 'security.enabled', System.getProperty('security.enabled')
    var is_https = System.getProperty("https")
    var user = System.getProperty("user")
    var password = System.getProperty("password")
    if (System.getProperty("security.enabled") != null) {
        // If security is enabled, set is_https/user/password defaults
        is_https = is_https == null ? "true" : is_https
        user = user == null ? "admin" : user
        password = password == null ? "admin" : password
    }
    systemProperty("https", is_https)
    systemProperty("user", user)
    systemProperty("password", password)

    doFirst {
        // Tell the test JVM if the cluster JVM is running under a debugger so that tests can
        // use longer timeouts for requests.
        def isDebuggingCluster = getDebug() || System.getProperty("test.debug") != null
        systemProperty 'cluster.debug', isDebuggingCluster
        // Set number of nodes system property to be used in tests
        systemProperty 'cluster.number_of_nodes', "${_numNodes}"

        // There seems to be an issue when running multi node run or integ tasks with unicast_hosts
        // not being written, the waitForAllConditions ensures it's written
        getClusters().forEach { cluster ->
            cluster.waitForAllConditions()
        }

        // Discover the actual gRPC port from the cluster log.
        // The gRPC transport binds to the first available port in 9400-9500,
        // which may differ from 9400 due to port conflicts or dual-stack binding.
        // Only discover when NOT using an external cluster (tests.rest.cluster),
        // since the Gradle-managed cluster's gRPC port is irrelevant for external clusters.
        // Note: For multi-node clusters, gRPC may run on any node; we parse the first node's log.
        if (System.getProperty("tests.rest.cluster") == null) {
            getClusters().forEach { cluster ->
                def logFile = new File(cluster.nodes.first().workingDir.toFile(), "logs/integTest.log")
                // Retry to handle race condition where gRPC transport hasn't written port to log yet
                def grpcPort = null
                for (int attempt = 0; attempt < 10 && grpcPort == null; attempt++) {
                    if (logFile.exists()) {
                        def matcher = (logFile.text =~ /Netty4GrpcServerTransport.*publish_address \{([^}]+):(\d+)\}/)
                        if (matcher.find()) {
                            grpcPort = matcher.group(2)
                        }
                    }
                    if (grpcPort == null) {
                        logger.lifecycle("Waiting for gRPC port in log file (attempt ${attempt + 1}/10)...")
                        Thread.sleep(1000)
                    }
                }
                if (grpcPort != null) {
                    systemProperty 'tests.grpc.port', grpcPort
                    logger.lifecycle("Discovered gRPC port: ${grpcPort}")
                } else {
                    logger.warn("Could not discover gRPC port from cluster log — gRPC tests will use default port ${9400}")
                }
            }
        }
    }

    // The --debug-jvm command-line option makes the cluster debuggable; this makes the tests debuggable
    if (System.getProperty("test.debug") != null) {
        jvmArgs '-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=*:5005'
    }
}

testClusters.integTest {
    testDistribution = "ARCHIVE"

    // Enable gRPC transport for integration tests
    // This is required for gRPC integration tests to work
    setting("aux.transport.types", "[transport-grpc]")
    // Bind gRPC to IPv4 loopback to avoid dual-stack port mismatch
    // where IPv6 takes 9400 and IPv4 gets 9401
    setting("grpc.bind_host", "127.0.0.1")

    // Optionally install security
    if (System.getProperty("security.enabled") != null && System.getProperty("security.enabled") == "true") {
        configureSecurityPlugin(testClusters.integTest)
    }

    // Install K-NN/ml-commons plugins on the integTest cluster nodes except security
    configurations.zipArchive.asFileTree.each {
        plugin(provider(new Callable<RegularFile>(){
                @Override
                RegularFile call() throws Exception {
                    return new RegularFile() {
                        @Override
                        File getAsFile() {
                            return it
                        }
                    }
                }
        }))
    }

    // This installs our neural-search plugin into the testClusters
    plugin(project.tasks.bundlePlugin.archiveFile)
    // Cluster shrink exception thrown if we try to set numberOfNodes to 1, so only apply if > 1
    if (_numNodes > 1) numberOfNodes = _numNodes

    // Configure node roles if enabled
    if (_configureNodeRoles && _numNodes > 1) {
        nodes.eachWithIndex { node, index ->
            if (index == 0) {
                node.setting('node.roles', '["cluster_manager"]')
            } else if (index == 1) {
                node.setting('node.roles', '["cluster_manager", "data", "ingest", "ml"]')
            } else {
                node.setting('node.roles', '["data", "ingest"]')
            }
        }
    }
    // When running integration tests it doesn't forward the --debug-jvm to the cluster anymore
    // i.e. we have to use a custom property to flag when we want to debug opensearch JVM
    // since we also support multi node integration tests we increase debugPort per node
    if (System.getProperty("cluster.debug") != null) {
        def debugPort = 5005
        nodes.forEach { node ->
            node.jvmArgs("-agentlib:jdwp=transport=dt_socket,server=n,suspend=y,address=${debugPort}")
            debugPort += 1
        }
    }

    // Increase heap size from default of 512mb to 1gb. When heap size is 512mb, our integ tests sporadically fail due
    // to ml-commons memory circuit breaker exception
    jvmArgs("-Xms1g", "-Xmx2g")
}

// Remote Integration Tests
task integTestRemote(type: RestIntegTestTask) {
    testClassesDirs = sourceSets.test.output.classesDirs
    classpath = sourceSets.test.runtimeClasspath

    systemProperty "https", System.getProperty("https")
    systemProperty "user", System.getProperty("user")
    systemProperty "password", System.getProperty("password")

    systemProperty 'cluster.number_of_nodes', "${_numNodes}"

    systemProperty 'tests.security.manager', 'false'

    // Run tests with remote cluster only if rest case is defined
    if (System.getProperty("tests.rest.cluster") != null) {
        filter {
            includeTestsMatching "org.opensearch.neuralsearch.*IT"
        }
    }
}

run {
    useCluster testClusters.integTest
}

jacocoTestReport {
    dependsOn integTest, test
    reports {
        xml.required = true
        html.required = true
    }
}

check.dependsOn spotlessCheck
check.dependsOn jacocoTestCoverageVerification
jacocoTestCoverageVerification.dependsOn jacocoTestReport

// updateVersion: Task to auto update version to the next development iteration
task updateVersion {
    onlyIf { System.getProperty('newVersion') }
    doLast {
        ext.newVersion = System.getProperty('newVersion')
        println "Setting version to ${newVersion}."

        // Parse version strings
        ext.newVersionWithoutSnapshot = newVersion.tokenize('-')[0]
        ext.currentVersionWithoutSnapshot = opensearch_version.tokenize('-')[0]

        // Step 1: Update build.gradle with new OpenSearch version
        ant.replaceregexp(
            file: 'build.gradle',
            match: '"opensearch.version", "\\d.*"',
            replace: "\"opensearch.version\", \"${newVersionWithoutSnapshot}-SNAPSHOT\"",
            flags: 'g',
            byline: true
        )

        // Step 2: Read current BWC versions from gradle.properties to determine what to update
        def gradlePropsFile = file('gradle.properties')
        def currentBwcVersion = null
        def currentBwcBundleVersion = null

        gradlePropsFile.eachLine { line ->
            if (line.startsWith('systemProp.bwc.version=')) {
                currentBwcVersion = line.split('=')[1].trim().tokenize('-')[0]
            }
            if (line.startsWith('systemProp.bwc.bundle.version=')) {
                currentBwcBundleVersion = line.split('=')[1].trim()
            }
        }

        if (currentBwcVersion == null || currentBwcBundleVersion == null) {
            throw new GradleException("Could not read BWC versions from gradle.properties")
        }

        println "Current BWC version: ${currentBwcVersion}"
        println "Current BWC bundle version: ${currentBwcBundleVersion}"
        println "New OpenSearch version: ${newVersionWithoutSnapshot}"
        println "Previous OpenSearch version: ${currentVersionWithoutSnapshot}"

        // Check version bump type
        def currentVersionParts = currentVersionWithoutSnapshot.tokenize('.')
        def newVersionParts = newVersionWithoutSnapshot.tokenize('.')
        def currentMajor = currentVersionParts[0] as Integer
        def currentMinor = currentVersionParts[1] as Integer
        def newMajor = newVersionParts[0] as Integer
        def newMinor = newVersionParts[1] as Integer
        def isPatchVersionBump = (currentMajor == newMajor && currentMinor == newMinor)
        def isMinorVersionBump = (currentMajor == newMajor && currentMinor != newMinor)
        def isMajorVersionBump = (currentMajor != newMajor)

        if (isPatchVersionBump) {
            println "Detected patch version bump (${currentVersionWithoutSnapshot} -> ${newVersionWithoutSnapshot}), will skip adding to BWC version arrays"
        } else if (isMinorVersionBump) {
            println "Detected minor version bump (${currentVersionWithoutSnapshot} -> ${newVersionWithoutSnapshot}), will skip updating bwc.bundle.version"
        } else if (isMajorVersionBump) {
            println "Detected major version bump (${currentVersionWithoutSnapshot} -> ${newVersionWithoutSnapshot})"
        }

        // Step 3: Update gradle.properties - systemProp.bwc.version
        ant.replaceregexp(
            file: 'gradle.properties',
            match: "systemProp.bwc.version=${currentBwcVersion}(-SNAPSHOT)?",
            replace: "systemProp.bwc.version=${newVersionWithoutSnapshot}-SNAPSHOT",
            flags: 'g',
            byline: true
        )

        // Step 4: Update gradle.properties - systemProp.bwc.bundle.version
        // Only update for major version bumps
        if (isMajorVersionBump) {
            ant.replaceregexp(
                file: 'gradle.properties',
                match: "systemProp.bwc.bundle.version=${currentBwcBundleVersion}",
                replace: "systemProp.bwc.bundle.version=${currentVersionWithoutSnapshot}",
                flags: 'g',
                byline: true
            )
            println "Updated systemProp.bwc.bundle.version to ${currentVersionWithoutSnapshot}"
        } else {
            println "Skipped updating systemProp.bwc.bundle.version (only updated for major version bumps)"
        }

        // Step 5: Update backwards_compatibility_tests_workflow.yml - opensearch_version
        ant.replaceregexp(
            file: '.github/workflows/backwards_compatibility_tests_workflow.yml',
            match: "opensearch_version\\s*:\\s*\\[\\s*\"${currentVersionWithoutSnapshot}-SNAPSHOT\"\\s*\\]",
            replace: "opensearch_version : [\"${newVersionWithoutSnapshot}-SNAPSHOT\"]",
            flags: 'g',
            byline: true
        )

        // Step 6: Update backwards_compatibility_tests_workflow.yml - bwc_version arrays
        // Add the current version to the BWC version list if not already present
        // Skip this step for patch version bumps (e.g., 3.7.0 -> 3.7.1)
        def workflowFile = file('.github/workflows/backwards_compatibility_tests_workflow.yml')
        def workflowContent = workflowFile.text

        if (!isPatchVersionBump && !workflowContent.contains("\"${currentVersionWithoutSnapshot}\"")) {
            // Read the workflow file and update it manually for more precise control
            def updatedContent = workflowContent

            // Helper closure to check if current version is later than the last version in a list
            def isCurrentVersionNewer = { String lastVersion ->
                def lastVersionParts = lastVersion.tokenize('.')
                def lastMajor = lastVersionParts[0] as Integer
                def lastMinor = lastVersionParts[1] as Integer
                def lastPatch = lastVersionParts.size() > 2 ? lastVersionParts[2] as Integer : 0

                def currVersionParts = currentVersionWithoutSnapshot.tokenize('.')
                def currMajor = currVersionParts[0] as Integer
                def currMinor = currVersionParts[1] as Integer
                def currPatch = currVersionParts.size() > 2 ? currVersionParts[2] as Integer : 0

                return (currMajor > lastMajor) ||
                       (currMajor == lastMajor && currMinor > lastMinor) ||
                       (currMajor == lastMajor && currMinor == lastMinor && currPatch > lastPatch)
            }

            // Helper closure to extract versions from bwc_version array
            def extractVersions = { String content ->
                def matcher = content =~ /bwc_version\s*:\s*\[([^\]]+)\]/
                if (matcher.find()) {
                    def versionListStr = matcher.group(1)
                    return versionListStr.findAll(/"([^"]+)"/).collect { it.replaceAll('"', '').trim() }
                }
                return []
            }

            // For Restart-Upgrade-BWCTests-NeuralSearch job
            def restartStart = updatedContent.indexOf('Restart-Upgrade-BWCTests-NeuralSearch:')
            def rollingStart = updatedContent.indexOf('Rolling-Upgrade-BWCTests-NeuralSearch:')

            if (restartStart != -1 && rollingStart != -1) {
                // Extract and update Restart-Upgrade section
                def restartSection = updatedContent.substring(restartStart, rollingStart)
                def restartVersions = extractVersions(restartSection)
                if (!restartVersions.isEmpty() && isCurrentVersionNewer(restartVersions.last())) {
                    def updatedRestartSection = restartSection.replaceFirst(
                        /(?m)(^\s*bwc_version\s*:\s*\[([^\]]+))\]/,
                        "\$1, \"${currentVersionWithoutSnapshot}\"]"
                    )
                    updatedContent = updatedContent.substring(0, restartStart) + updatedRestartSection + updatedContent.substring(rollingStart)
                    println "Added ${currentVersionWithoutSnapshot} to Restart-Upgrade BWC version list"
                }

                // Extract and update Rolling-Upgrade section
                def rollingSection = updatedContent.substring(rollingStart)
                def rollingVersions = extractVersions(rollingSection)
                if (!rollingVersions.isEmpty() && isCurrentVersionNewer(rollingVersions.last())) {
                    def updatedRollingSection = rollingSection.replaceFirst(
                        /(?m)(^\s*bwc_version:\s*\[([^\]]+))\]/,
                        "\$1, \"${currentVersionWithoutSnapshot}\"]"
                    )
                    // Need to recalculate rollingStart in case the restart section changed
                    def newRollingStart = updatedContent.indexOf('Rolling-Upgrade-BWCTests-NeuralSearch:')
                    updatedContent = updatedContent.substring(0, newRollingStart) + updatedRollingSection
                    println "Added ${currentVersionWithoutSnapshot} to Rolling-Upgrade BWC version list"
                }
            }

            workflowFile.text = updatedContent
        } else if (isPatchVersionBump) {
            println "Skipped adding ${currentVersionWithoutSnapshot} to BWC version list (patch version bump)"
        } else {
            println "Version ${currentVersionWithoutSnapshot} already exists in BWC version list"
        }
    }
}


// Remote Model Integration Tests
def isDockerAvailable() {
    try {
        def process = "docker version".execute()
        process.waitFor()
        return process.exitValue() == 0
    } catch (Exception e) {
        return false
    }
}

task remoteModelIntegTest(type: RestIntegTestTask) {
    description = "Run all remote model integration tests with auto-discovered models"
    testClassesDirs = sourceSets.test.output.classesDirs
    classpath = sourceSets.test.runtimeClasspath

    // Filter to only run remote model tests
    filter {
        includeTestsMatching "*RemoteModelIT*"
    }

    // Path to handlers directory for auto-discovery
    def handlersDir = file("src/test/resources/remote-models/torchserve/handlers")
    def discoveredModels = []

    // Auto-discover all models from handler files
    if (handlersDir.exists()) {
        handlersDir.eachFile { file ->
            if (file.name.endsWith("_handler.py")) {
                def modelName = file.name.replace("_handler.py", "")
                discoveredModels.add(modelName)
            }
        }
    }

    // Set endpoint for each discovered model
    discoveredModels.each { model ->
        def envVarName = "${model.toUpperCase().replace('_', '_')}_ENDPOINT"
        def endpoint = "http://localhost:8080/predictions/${model}"
        environment envVarName, endpoint
        systemProperty "tests.${model}.endpoint", endpoint
    }

    // Common endpoints for backward compatibility
    environment "TORCHSERVE_ENDPOINT", "http://localhost:8080"
    systemProperty "tests.torchserve.endpoint", "http://localhost:8080"

    // System properties from integTest
    systemProperty 'tests.security.manager', 'false'
    systemProperty 'java.io.tmpdir', opensearch_tmp_dir.absolutePath
    systemProperty('project.root', project.rootDir.absolutePath)

    // Path to the run.sh script
    def runScript = file("src/test/resources/remote-models/torchserve/scripts/run.sh")

    doFirst {
        // Check Docker availability
        if (!isDockerAvailable()) {
            throw new GradleException("Docker is not available. Please install Docker to run remote model tests.")
        }

        // Ensure script exists
        if (!runScript.exists()) {
            throw new GradleException("run.sh script not found at: ${runScript.absolutePath}")
        }

        // Make script executable
        runScript.setExecutable(true)

        // Display discovered models
        if (!discoveredModels.isEmpty()) {
            println "Auto-discovered models: ${discoveredModels.join(', ')}"
        } else {
            println "WARNING: No models discovered in ${handlersDir.absolutePath}"
        }

        // Start TorchServe and setup all discovered models
        println "Setting up TorchServe and deploying all models..."
        providers.exec {
            commandLine(runScript.absolutePath, "lifecycle", "setup")
        }.result.get().assertNormalExitValue()
    }

    doLast {
        // Teardown TorchServe
        println "Stopping TorchServe..."
        try {
            providers.exec {
                commandLine(runScript.absolutePath, "lifecycle", "teardown")
            }.result.get().assertNormalExitValue()
        } catch (Exception e) {
            // Do not fail if already stopped
        }
    }

    // Test logging configuration
    testLogging {
        events "passed", "skipped", "failed"
        showExceptions true
        showCauses true
        showStackTraces true
        exceptionFormat = "full"
    }

    // Generate separate test report
    reports {
        html.outputLocation = file("$buildDir/reports/remote-model-tests")
        junitXml.outputLocation = file("$buildDir/test-results/remote-model-tests")
    }
}

// Configure test cluster for remoteModelIntegTest
testClusters.remoteModelIntegTest {
    testDistribution = "ARCHIVE"

    // Install plugins
    configurations.zipArchive.asFileTree.each {
        plugin(provider(new Callable<RegularFile>(){
                @Override
                RegularFile call() throws Exception {
                    return new RegularFile() {
                        @Override
                        File getAsFile() {
                            return it
                        }
                    }
                }
        }))
    }

    // This installs our neural-search plugin into the testClusters
    plugin(project.tasks.bundlePlugin.archiveFile)

    // Increase heap size to avoid memory issues
    jvmArgs("-Xms1g", "-Xmx2g")
}

// Helper tasks for manual control (optional)
task startTorchServe(type: Exec) {
    group = "Remote Model Testing"
    description = "Start TorchServe container manually"

    def runScript = file("src/test/resources/remote-models/torchserve/scripts/run.sh")
    commandLine runScript.absolutePath, "start"
}

task stopTorchServe(type: Exec) {
    group = "Remote Model Testing"
    description = "Stop TorchServe container manually"

    def runScript = file("src/test/resources/remote-models/torchserve/scripts/run.sh")
    commandLine runScript.absolutePath, "stop"
}

task torchServeStatus(type: Exec) {
    group = "Remote Model Testing"
    description = "Check TorchServe status"

    def runScript = file("src/test/resources/remote-models/torchserve/scripts/run.sh")
    commandLine runScript.absolutePath, "status"
}

task listRemoteModels(type: Exec) {
    group = "Remote Model Testing"
    description = "List all discovered remote models"

    def runScript = file("src/test/resources/remote-models/torchserve/scripts/run.sh")
    commandLine runScript.absolutePath, "list-models"
}