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',
  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 {
  compile project(sparkRunnerProject)
  compile project(path: sparkRunnerProject, configuration: "provided")
  validatesPortableRunner project(path: sparkRunnerProject, configuration: "testRuntime")
  validatesPortableRunner project(path: sparkRunnerProject, configuration: "provided")
  validatesPortableRunner project(path: ":sdks:java:core", configuration: "shadowTest")
  validatesPortableRunner project(path: ":runners:core-java", configuration: "testRuntime")
  validatesPortableRunner project(path: ":runners:portability:java", configuration: "testRuntime")
  compile project(":sdks:java:extensions:google-cloud-platform-core")
//  TODO: Enable AWS and 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 portableValidatesRunnerTask(String name) {
  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,
    environment: BeamModulePlugin.PortableValidatesRunnerConfiguration.Environment.EMBEDDED,
    systemProperties: [
      "beam.spark.test.reuseSparkContext": "false",
      "spark.ui.enabled": "false",
      "spark.ui.showConsoleProgress": "false",
    ],
    testCategories: {
      includeCategories 'org.apache.beam.sdk.testing.ValidatesRunner'
      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.UsesParDoLifecycle'
      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'
      //SplitableDoFnTests
      excludeCategories 'org.apache.beam.sdk.testing.UsesBoundedSplittableParDo'
      excludeCategories 'org.apache.beam.sdk.testing.UsesSplittableParDoWithWindowedSideInputs'
      excludeCategories 'org.apache.beam.sdk.testing.UsesUnboundedSplittableParDo'
      excludeCategories 'org.apache.beam.sdk.testing.UsesStrictTimerOrdering'
      excludeCategories 'org.apache.beam.sdk.testing.UsesBundleFinalizer'
    },
    testFilter: {
      // TODO(BEAM-10094)
      excludeTestsMatching 'org.apache.beam.sdk.transforms.FlattenTest.testFlattenWithDifferentInputAndOutputCoders2'
    },
  )
}

project.ext.validatesPortableRunnerBatch = portableValidatesRunnerTask("Batch")

task validatesPortableRunner() {
  dependsOn validatesPortableRunnerBatch
}

def jobPort = BeamModulePlugin.startingExpansionPortNumber.getAndDecrement()
def artifactPort = BeamModulePlugin.startingExpansionPortNumber.getAndDecrement()

def setupTask = project.tasks.create(name: "sparkJobServerSetup", type: 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.create(name: "sparkJobServerCleanup", type: 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,
  jobEndpoint: "localhost:${jobPort}",
  classpath: configurations.validatesPortableRunner,
  numParallelTests: 1
)

def addTestJavaJarCreator(String pyVersion) {
  def pyBuildPath = pyVersion.startsWith("2") ? "2" : pyVersion.replaceAll("\\.", "")
  project.tasks.create(name: "testJavaJarCreatorPy${pyBuildPath}") {
    dependsOn shadowJar
    dependsOn ":sdks:python:container:py${pyBuildPath}:docker"
    doLast{
      exec {
        executable "sh"
        def options = [
                "--runner SparkRunner",
                "--job_server_jar ${shadowJar.archivePath}",
                "--env_dir ${project.rootProject.buildDir}/gradleenv/${project.path.hashCode()}",
                "--python_root_dir ${project.rootDir}/sdks/python",
                "--python_version ${pyVersion}",
                "--python_container_image ${project.docker_image_default_repo_root}/"
                        + "${project.docker_image_default_repo_prefix}"
                        + "python${pyVersion}_sdk:${project.sdk_version}",
        ]
        args "-c", "../../portability/test_pipeline_jar.sh ${options.join(' ')}"
      }
    }
  }
}

["2.7", "3.5", "3.6", "3.7"].each { pyVersion ->
  addTestJavaJarCreator(pyVersion)
}

task testPipelineJar() {
  dependsOn testJavaJarCreatorPy37
}