import groovy.json.JsonOutput
import org.apache.beam.gradle.BeamModulePlugin

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * License); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an AS IS BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 * Spark Runner JobServer build file
 */

apply plugin: 'org.apache.beam.module'
apply plugin: 'application'
// we need to set mainClassName before applying shadow plugin
mainClassName = "org.apache.beam.runners.spark.SparkJobServerDriver"

applyJavaNature(
  automaticModuleName: 'org.apache.beam.runners.spark.jobserver',
  archivesBaseName: project.hasProperty('archives_base_name') ? archives_base_name : archivesBaseName,
  validateShadowJar: false,
  exportJavadoc: false,
  shadowClosure: {
    append "reference.conf"
  },
)

def sparkRunnerProject = project.parent.path

description = "Apache Beam :: Runners :: Spark :: Job Server"

configurations {
  validatesPortableRunner
}

configurations.all {
  exclude group: "org.slf4j", module: "slf4j-jdk14"
}

dependencies {
  implementation project(sparkRunnerProject)
  permitUnusedDeclared project(sparkRunnerProject)
  validatesPortableRunner project(path: sparkRunnerProject, configuration: "testRuntimeMigration")
  validatesPortableRunner project(path: ":sdks:java:core", configuration: "shadowTest")
  validatesPortableRunner project(path: ":runners:core-java", configuration: "testRuntimeMigration")
  validatesPortableRunner project(path: ":runners:portability:java", configuration: "testRuntimeMigration")
  runtimeOnly project(":sdks:java:extensions:google-cloud-platform-core")
  runtimeOnly project(":sdks:java:io:amazon-web-services2")
//  TODO: Enable HDFS file system.
}

// NOTE: runShadow must be used in order to run the job server. The standard run
// task will not work because the Spark runner classes only exist in the shadow
// jar.
runShadow {
  args = []
  if (project.hasProperty('jobHost'))
    args += ["--job-host=${project.property('jobHost')}"]
  if (project.hasProperty('artifactsDir'))
    args += ["--artifacts-dir=${project.property('artifactsDir')}"]
  if (project.hasProperty('cleanArtifactsPerJob'))
    args += ["--clean-artifacts-per-job=${project.property('cleanArtifactsPerJob')}"]
  if (project.hasProperty('sparkMasterUrl'))
    args += ["--spark-master-url=${project.property('sparkMasterUrl')}"]

  systemProperties System.properties

  // Enable remote debugging.
  jvmArgs = ["-Xdebug", "-Xrunjdwp:transport=dt_socket,server=y,suspend=n,address=5005"]
  if (project.hasProperty("logLevel"))
    jvmArgs += ["-Dorg.slf4j.simpleLogger.defaultLogLevel=${project.property('logLevel')}"]
}

def sickbayTests = [
        // TODO(BEAM-13498)
        'org.apache.beam.sdk.transforms.ParDoTest$TimestampTests.testProcessElementSkew',
        // TODO(https://github.com/apache/beam/issues/29973)
        'org.apache.beam.sdk.transforms.ReshuffleTest.testReshufflePreservesMetadata',
]

def portableValidatesRunnerTask(String name, boolean streaming, boolean docker, ArrayList<String> sickbayTests) {
  def pipelineOptions = []
  def testCategories
  def testFilter

  if (docker) {
    // Run the limited set of tests that need to validate the environment
    // that contains the SDK is configured properly.
    testCategories = { includeCategories 'org.apache.beam.sdk.testing.UsesSdkHarnessEnvironment' }
    testFilter = { }
  } else {
    if (streaming) {
      pipelineOptions += "--streaming"
      pipelineOptions += "--streamingTimeoutMs=120000"

      testCategories = {
        includeCategories 'org.apache.beam.sdk.testing.ValidatesRunner'
        excludeCategories 'org.apache.beam.sdk.testing.UsesExternalService'
        // Should be run only in a properly configured SDK harness environment
        excludeCategories 'org.apache.beam.sdk.testing.UsesSdkHarnessEnvironment'
        excludeCategories 'org.apache.beam.sdk.testing.FlattenWithHeterogeneousCoders'
        excludeCategories 'org.apache.beam.sdk.testing.LargeKeys$Above100MB'
        excludeCategories 'org.apache.beam.sdk.testing.UsesCommittedMetrics'
        excludeCategories 'org.apache.beam.sdk.testing.UsesCustomWindowMerging'
        excludeCategories 'org.apache.beam.sdk.testing.UsesFailureMessage'
        excludeCategories 'org.apache.beam.sdk.testing.UsesGaugeMetrics'
        excludeCategories 'org.apache.beam.sdk.testing.UsesPerKeyOrderedDelivery'
        excludeCategories 'org.apache.beam.sdk.testing.UsesMapState'
        excludeCategories 'org.apache.beam.sdk.testing.UsesSetState'
        excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState'
        excludeCategories 'org.apache.beam.sdk.testing.UsesTimerMap'
        excludeCategories 'org.apache.beam.sdk.testing.UsesKeyInParDo'
        excludeCategories 'org.apache.beam.sdk.testing.UsesOnWindowExpiration'
        excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream'
        // TODO (https://github.com/apache/beam/issues/19468) SplittableDoFnTests
        excludeCategories 'org.apache.beam.sdk.testing.UsesBoundedSplittableParDo'
        excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo'
        excludeCategories 'org.apache.beam.sdk.testing.UsesStrictTimerOrdering'
        excludeCategories 'org.apache.beam.sdk.testing.UsesBundleFinalizer'
        // Currently unsupported in portable streaming:
        // TODO (https://github.com/apache/beam/issues/20395)
        excludeCategories 'org.apache.beam.sdk.testing.UsesSideInputs'
        // TODO (https://github.com/apache/beam/issues/20396)
        excludeCategories 'org.apache.beam.sdk.testing.UsesStatefulParDo'
        // TODO (https://github.com/apache/beam/issues/20397)
        excludeCategories 'org.apache.beam.sdk.testing.UsesTimersInParDo'
        excludeCategories 'org.apache.beam.sdk.testing.UsesTriggeredSideInputs'
      }

      testFilter = {
        // TODO (https://github.com/apache/beam/issues/20189)
        excludeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenWithDifferentInputAndOutputCoders2'
        // TODO (https://github.com/apache/beam/issues/20429) Currently unsupported in portable streaming:
        // // Timeout error
        excludeTestsMatching 'org.apache.beam.sdk.testing.PAssertTest.testWindowedContainsInAnyOrder'
        excludeTestsMatching 'org.apache.beam.sdk.testing.PAssertTest.testWindowedSerializablePredicate'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.windowing.WindowTest.testNoWindowFnDoesNotReassignWindows'
        // // Assertion error: empty iterable output
        excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest$WindowingTests.testFixedWindowsCombine'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.CombineTest$WindowingTests.testSessionsCombine'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.GroupByKeyTest$WindowTests'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.ReshuffleTest.testReshuffleAfterFixedWindowsAndGroupByKey'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.ReshuffleTest.testReshuffleAfterSessionsAndGroupByKey'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.ReshuffleTest.testReshuffleAfterSlidingWindowsAndGroupByKey'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.join.CoGroupByKeyTest.testCoGroupByKeyWithWindowing'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.windowing.WindowingTest'
        // // Assertion error: incorrect output
        excludeTestsMatching 'CombineTest$BasicTests.testHotKeyCombining'
        // TODO(https://github.com/apache/beam/issues/29973)
        excludeTestsMatching 'org.apache.beam.sdk.transforms.ReshuffleTest.testReshufflePreservesMetadata'
        // TODO(https://github.com/apache/beam/issues/31231)
        excludeTestsMatching 'org.apache.beam.sdk.transforms.RedistributeTest.testRedistributePreservesMetadata'
        // TODO(https://github.com/apache/beam/issues/31234) same reason as GroupByKeyTest and ReshuffleTest above
        excludeTestsMatching 'org.apache.beam.sdk.transforms.RedistributeTest.testRedistributeAfterFixedWindows'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.RedistributeTest.testRedistributeAfterSlidingWindows'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.RedistributeTest.testRedistributeAfterFixedWindowsAndGroupByKey'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.RedistributeTest.testRedistributeAfterSessionsAndGroupByKey'
        excludeTestsMatching 'org.apache.beam.sdk.transforms.RedistributeTest.testRedistributeAfterSlidingWindowsAndGroupByKey'
      }
    }
    else {
      // Batch
      testCategories = {
        includeCategories 'org.apache.beam.sdk.testing.ValidatesRunner'
        excludeCategories 'org.apache.beam.sdk.testing.UsesExternalService'
        // Should be run only in a properly configured SDK harness environment
        excludeCategories 'org.apache.beam.sdk.testing.UsesSdkHarnessEnvironment'
        excludeCategories 'org.apache.beam.sdk.testing.FlattenWithHeterogeneousCoders'
        excludeCategories 'org.apache.beam.sdk.testing.LargeKeys$Above100MB'
        excludeCategories 'org.apache.beam.sdk.testing.UsesCommittedMetrics'
        excludeCategories 'org.apache.beam.sdk.testing.UsesCustomWindowMerging'
        excludeCategories 'org.apache.beam.sdk.testing.UsesFailureMessage'
        excludeCategories 'org.apache.beam.sdk.testing.UsesGaugeMetrics'
        excludeCategories 'org.apache.beam.sdk.testing.UsesPerKeyOrderedDelivery'
        excludeCategories 'org.apache.beam.sdk.testing.UsesPerKeyOrderInBundle'
        excludeCategories 'org.apache.beam.sdk.testing.UsesMapState'
        excludeCategories 'org.apache.beam.sdk.testing.UsesMultimapState'
        excludeCategories 'org.apache.beam.sdk.testing.UsesSetState'
        excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState'
        excludeCategories 'org.apache.beam.sdk.testing.UsesTimerMap'
        excludeCategories 'org.apache.beam.sdk.testing.UsesLoopingTimer'
        excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedPCollections'
        excludeCategories 'org.apache.beam.sdk.testing.UsesKeyInParDo'
        excludeCategories 'org.apache.beam.sdk.testing.UsesOnWindowExpiration'
        excludeCategories 'org.apache.beam.sdk.testing.UsesTestStream'
        // TODO (https://github.com/apache/beam/issues/19468) SplittableDoFnTests
        excludeCategories 'org.apache.beam.sdk.testing.UsesBoundedSplittableParDo'
        excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo'
        excludeCategories 'org.apache.beam.sdk.testing.UsesStrictTimerOrdering'
        excludeCategories 'org.apache.beam.sdk.testing.UsesBundleFinalizer'
      }
      testFilter = {
        // TODO (https://github.com/apache/beam/issues/20189)
        excludeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenWithDifferentInputAndOutputCoders2'
        // TODO(https://github.com/apache/beam/issues/31231)
        excludeTestsMatching 'org.apache.beam.sdk.transforms.RedistributeTest.testRedistributePreservesMetadata'
        for (String test : sickbayTests) {
          excludeTestsMatching test
        }
      }
    }
  }

  createPortableValidatesRunnerTask(
          name: "validatesPortableRunner${name}",
          jobServerDriver: "org.apache.beam.runners.spark.SparkJobServerDriver",
          jobServerConfig: "--job-host=localhost,--job-port=0,--artifact-port=0,--expansion-port=0",
          testClasspathConfiguration: configurations.validatesPortableRunner,
          numParallelTests: 4,
          pipelineOpts: pipelineOptions,
          environment: docker ? BeamModulePlugin.PortableValidatesRunnerConfiguration.Environment.DOCKER : BeamModulePlugin.PortableValidatesRunnerConfiguration.Environment.EMBEDDED,
          systemProperties: [
                  "beam.spark.test.reuseSparkContext": "false",
                  "spark.ui.enabled": "false",
                  "spark.ui.showConsoleProgress": "false",
          ],
          testCategories: testCategories,
          testFilter: testFilter,
    )
}

project.ext.validatesPortableRunnerDocker= portableValidatesRunnerTask("Docker", false, true, sickbayTests)
project.ext.validatesPortableRunnerBatch = portableValidatesRunnerTask("Batch", false, false, sickbayTests)
project.ext.validatesPortableRunnerStreaming = portableValidatesRunnerTask("Streaming", true, false, sickbayTests)

tasks.register("validatesPortableRunner") {
  dependsOn validatesPortableRunnerDocker
  dependsOn validatesPortableRunnerBatch
  dependsOn validatesPortableRunnerStreaming
}

tasks.register("validatesRunnerSickbay", Test) {
  group = "Verification"
  description "Validates Spark runner (Sickbay Tests)"
  systemProperty "beamTestPipelineOptions", JsonOutput.toJson([
          "--runner=TestSparkRunner",
  ])

  classpath = configurations.validatesPortableRunner
  testClassesDirs = files(project(":sdks:java:core").sourceSets.test.output.classesDirs)

  filter {
    for (String test : sickbayTests) {
      includeTestsMatching test
    }
  }
}

def jobPort = BeamModulePlugin.getRandomPort()
def artifactPort = BeamModulePlugin.getRandomPort()

def setupTask = project.tasks.register("sparkJobServerSetup", Exec) {
  dependsOn shadowJar
  def pythonDir = project.project(":sdks:python").projectDir
  def sparkJobServerJar = shadowJar.archivePath

  executable 'sh'
  args '-c', "$pythonDir/scripts/run_job_server.sh stop --group_id ${project.name} && $pythonDir/scripts/run_job_server.sh start --group_id ${project.name} --job_port ${jobPort} --artifact_port ${artifactPort} --job_server_jar ${sparkJobServerJar}"
}

def cleanupTask = project.tasks.register("sparkJobServerCleanup", Exec) {
  def pythonDir = project.project(":sdks:python").projectDir

  executable 'sh'
  args '-c', "$pythonDir/scripts/run_job_server.sh stop --group_id ${project.name}"
}

createCrossLanguageValidatesRunnerTask(
  startJobServer: setupTask,
  cleanupJobServer: cleanupTask,
  classpath: configurations.validatesPortableRunner,
  numParallelTests: 1,
  pythonPipelineOptions: [
    "--runner=PortableRunner",
    "--job_endpoint=localhost:${jobPort}",
    "--environment_cache_millis=10000",
    "--experiments=beam_fn_api",
  ],
  javaPipelineOptions: [
    "--runner=PortableRunner",
    "--jobEndpoint=localhost:${jobPort}",
    "--environmentCacheMillis=10000",
    "--experiments=beam_fn_api",
    "--customBeamRequirement=${project.project(":sdks:python").projectDir}/build/apache-beam.tar.gz",
  ],
  goScriptOptions: [
    "--runner spark",
    "--tests \"./test/integration/xlang ./test/integration/io/xlang/...\"",
    "--endpoint localhost:${jobPort}",
  ],
)

shadowJar {
  outputs.upToDateWhen { false }
}
