/*
 * 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.
 */

def pythonVersionSuffix = project.ext.pythonVersion.replace('.', '')

description = "Apache Beam :: SDKs :: Python :: Container :: ML :: Python ${pythonVersionSuffix} Container"

configurations {
  sdkSourceTarball
  pythonHarnessLauncher
}

dependencies {
  sdkSourceTarball project(path: ":sdks:python", configuration: "distTarBall")
  pythonHarnessLauncher project(path: ":sdks:python:container", configuration: "pythonHarnessLauncher")
}

def generatePythonRequirements = tasks.register("generatePythonRequirements") {
  dependsOn ':sdks:python:sdist'
  def pipExtraOptions = project.hasProperty("testRCDependencies") ? "--pre" : ""
  def runScriptsPath = "${rootDir}/sdks/python/container/run_generate_requirements.sh"
  doLast {
    exec {
      executable 'sh'
      args '-c', "cd ${rootDir} && ${runScriptsPath} " +
              "${project.ext.pythonVersion} " +
              "${files(configurations.sdkSourceTarball.files).singleFile} " +
              "base_image_requirements.txt " +
              "[gcp,dataframe,test] " +
              "${pipExtraOptions}"
    }
    // Generate versions for ML dependencies
    exec {
      executable 'sh'
      args '-c', "cd ${rootDir} && ${runScriptsPath} " +
              "${project.ext.pythonVersion} " +
              "${files(configurations.sdkSourceTarball.files).singleFile} " +
              "ml_image_requirements.txt " +
              "[gcp,dataframe,test,tensorflow,torch,transformers] " +
              "${pipExtraOptions}"
    }
  }
}

def copyDockerfileDependencies = tasks.register("copyDockerfileDependencies", Copy) {
  from configurations.sdkSourceTarball
  from file("base_image_requirements.txt")
  into "build/target"
  if(configurations.sdkSourceTarball.isEmpty()) {
      throw new StopExecutionException();
  }
}

def copyLicenseScripts = tasks.register("copyLicenseScripts", Copy){
  from ("../../license_scripts")
  into "build/target/license_scripts"
}

def copyLauncherDependencies = tasks.register("copyLauncherDependencies", Copy) {
  from configurations.pythonHarnessLauncher
  into "build/target/launcher"

  // Avoid seemingly gradle bug stated in https://github.com/apache/beam/issues/29220
  mustRunAfter "copyLicenses"

  if(configurations.pythonHarnessLauncher.isEmpty()) {
      throw new StopExecutionException();
  }
}

def pushContainers = project.rootProject.hasProperty(["isRelease"]) || project.rootProject.hasProperty("push-containers")

docker {
  name containerImageName(
          name: project.docker_image_default_repo_prefix + "python${project.ext.pythonVersion}_sdk_ml",
          root: project.rootProject.hasProperty(["docker-repository-root"]) ?
                  project.rootProject["docker-repository-root"] :
                  project.docker_image_default_repo_root,
          tag: project.rootProject.hasProperty(["docker-tag"]) ?
                  project.rootProject["docker-tag"] : project.sdk_version)
  // tags used by dockerTag task
  tags containerImageTags()
  files "../../Dockerfile", "./build"
  buildArgs(['py_version': "${project.ext.pythonVersion}",
             'pull_licenses': project.rootProject.hasProperty(["docker-pull-licenses"]) ||
                     project.rootProject.hasProperty(["isRelease"])])
  buildx project.useBuildx()
  platform(*project.containerPlatforms())
  load project.useBuildx() && !pushContainers
  push pushContainers
}

dockerPrepare.dependsOn copyLauncherDependencies
dockerPrepare.dependsOn copyDockerfileDependencies
dockerPrepare.dependsOn copyLicenseScripts

if (project.rootProject.hasProperty("docker-pull-licenses")) {
  def copyGolangLicenses = tasks.register("copyGolangLicenses", Copy) {
    from "${project(':release:go-licenses:py').buildDir}/output"
    into "build/target/go-licenses"
    dependsOn ':release:go-licenses:py:createLicenses'
  }
  dockerPrepare.dependsOn copyGolangLicenses
} else {
  def skipPullLicenses = tasks.register("skipPullLicenses", Exec) {
    executable "sh"
    // Touch a dummy file to ensure the directory exists.
    args "-c", "mkdir -p build/target/go-licenses && touch build/target/go-licenses/skip"
  }
  dockerPrepare.dependsOn skipPullLicenses
}