#!/usr/bin/env groovy

docker_registry_server = targetImageTag.split(':')[0..1].join(':')
// This will be the "RUN" displayed on Blue Ocean
currentBuild.displayName = targetImageTag.split(':')[2]
// This will be the "MESSAGE" displayed on Blue Ocean
currentBuild.description = sourceBranch + ": " + commitHash

podTemplate(
cloud:'sc-ipp-blossom-prod',
//volumes: [persistentVolumeClaim(mountPath: '/mount_binaries', claimName: 'kaolin-pvc', readOnly: false)],
yaml:'''
apiVersion: v1
kind: Pod
spec:
  volumes:
  - name: pvc-mount
    persistentVolumeClaim:
      claimName: 'kaolin-pvc'
  containers:
  - name: docker
    image: docker:19.03.1
    command:
    - sleep
    args:
    - 1d
    env:
      - name: DOCKER_HOST
        value: tcp://localhost:2375
    volumeMounts:
      - mountPath: /mount_binaries
        name: pvc-mount
  - name: docker-daemon
    image: docker:19.03.1-dind
    securityContext:
      privileged: true
    env:
      - name: DOCKER_TLS_CERTDIR
        value: ""
    resources:
        requests:
          memory: 32Gi
          cpu: 12
        limits:
          memory: 32Gi
          cpu: 12
    volumeMounts:
      - mountPath: /mount_binaries
        name: pvc-mount
''') {
  node(POD_LABEL) {
    container("docker") {
      try {
        gitlabCommitStatus("build-${configName}") {
          stage("Checkout") {
            checkout([
                $class: 'GitSCM',
                branches: [[name: "${commitHash}"]],
                extensions: [[
                    $class: 'SubmoduleOption',
                    disableSubmodules: false,
                    parentCredentials: false,
                    recursiveSubmodules: true,
                    reference: '',
                    trackingSubmodules: false
                ]],
                userRemoteConfigs: [[
                    credentialsId: 'kaolin-gitlab-access-token-as-password',
                    url: "${repoUrl}"
                ]]
            ])
          }
          docker.withRegistry("https://${docker_registry_server}", 'kaolin-gitlab-access-token-as-password') {
            stage("Build") {
              baseImage = docker.build(
                  "${targetImageTag}-base",
                  """--no-cache --network host -f ./tools/linux/Dockerfile.base \
                     --build-arg CUDA_VERSION=${cudaVer} \
                     --build-arg CUDNN_VERSION=${cudnnVer} \
                     --build-arg PYTHON_VERSION=${pythonVer} \
                     --build-arg PYTORCH_VERSION=${torchVer} \
                     .
                  """)
              targetImage = docker.build(
                  "${targetImageTag}",
                  """--no-cache --network host -f ./tools/linux/Dockerfile.install \
                     --build-arg BASE_IMAGE=${targetImageTag}-base \
                     .
                  """)
            }
            if (buildWheel.toBoolean()) {
              stage("Build wheel") {
                cudaTag = cudaVer.split('\\.')[0..<2].join('')
                targetImage.inside() {
                  sh """
                  ls .
                  python setup.py bdist_wheel --dist-dir .
                  """
                }
                pythonVerTag = pythonVer.split('\\.').join('')
                Integer MinorVal = pythonVer.split('\\.')[1].toInteger()
                pythonVerAbiTag = (MinorVal < 8) ? pythonVerTag + 'm' : pythonVerTag
                kaolinVer = sh(script: "cat ./version.txt", returnStdout: true).trim()
                baseWheelName = "kaolin-${kaolinVer}-cp${pythonVerTag}-cp${pythonVerAbiTag}"
                wheelName = "${baseWheelName}-linux_x86_64.whl"
              }
              stage("Reinstall from wheel") {
                targetImage = docker.build(
                    "${targetImageTag}",
                    """--no-cache --network host -f ./tools/linux/Dockerfile.install_wheel \
                       --build-arg BASE_IMAGE=${targetImageTag}-base \
                       --build-arg WHEEL_NAME=${wheelName} \
                       .
                    """)
              }
              stage("Push wheel to volume") {
                sh """
                mkdir -p /mount_binaries/whl/torch-${torchVer}_cu${cudaTag}
                cp ./${wheelName} /mount_binaries/whl/torch-${torchVer}_cu${cudaTag}/${wheelName}
                """
              }
              stage("Push wheel to artifact") {
                archiveArtifacts artifacts: "${wheelName}"
              }
            }
            stage("Push") {
              targetImage.push()
            }
          }
        }
      } catch (e) {
        // In case of build failure, we need to update the following tests as we won't run them.
        for (arch in archsToTest.split(';')) {
          updateGitlabCommitStatus(name: "test-${configName}-${arch}", state: 'canceled')
        }
        throw e
      }
      stage("Launch tests") {
        jobMap = [:]
        for (arch in archsToTest.split(';')) {
          jobMap["${arch}"] = prepareUbuntuTestJob(arch)
        }
        parallel jobMap
      }
    }
  }
}

def prepareUbuntuTestJob(arch) {
  return {
    stage("Test ${arch}") {
      build job: "ubuntu_test_template_CI",
      parameters: [
        string(name: 'sourceBranch', value: "${sourceBranch}"),
        string(name: 'configName', value: "${configName}"),
        string(name: 'imageTag', value: "${targetImageTag}"),
        string(name: 'arch', value: "${arch}"),
        string(name: 'commitHash', value: "${commitHash}")
      ],
      // This node doesn't need to be held while tests run.
      wait: false,
      // Success of this script depend only on successful build
      // and launch of tests, not successful tests.
      propagate: false
    }
  }
}