@Library('test-shared-library@dai-updated') _
import ai.h2o.ci.buildsummary.StagesSummary

DEFAULT_LABEL = 'h2o3_linux'

H2O_ECR='353750902984.dkr.ecr.us-east-1.amazonaws.com'

// initialize build summary
buildSummary('https://github.com/h2oai/h2o-3', true)
// use default StagesSummary implementation
buildSummary.get().addStagesSummary(this, new StagesSummary())

properties([
    parameters([
        booleanParam(defaultValue: false, description: 'Build even if the image is already present in the repository', name: 'forceBuild'),
        booleanParam(defaultValue: false, description: 'Publish built images', name: 'publish'),
        string(defaultValue: 'master', description: 'H2O-3 Branch used to load dependencies.', name: 'h2oBranch'),
        booleanParam(name: 'noCache', defaultValue: true, description: 'If set to true, the docker image is built from scratch, with no stages used from the cache.')
    ])
])

// IMAGE_NAME_PREFIX = 'harbor.h2o.ai/opsh2oai/h2o-3'

JDK_VERSIONS = ['8', '11', '17']
JDK_VERSIONS_PY_R = ['8', '11', '17'] // stable, last-supported, latest

def stashDockerImage(image) {
	sh "docker save -o ${image}.tar ${H2O_ECR}/h2o-3/${image}"
	stash includes: "${image}.tar", name: "docker-image-${image}"
	sh "rm ${image}.tar"
}

def unstashDockerImage(image) {
	unstash "docker-image-${image}"
	sh "docker load -i ${image}.tar"
	sh "rm ${image}.tar"
}

def stashDockerImages(target){
	if (!params.publish) {  // when publish this is noop since we can retrieve the image from registry
		toPersist = ["dev-base", "dev-jdk-8-base", "dev-jdk-others-base", "dev-python-3.7", "dev-r-4.5.2",
					 "dev-build-base", "dev-build-hadoop", "dev-python-base", "dev-r-base"]
		if (toPersist.contains(target)) {
			stashDockerImage(target)
		}
	}
}

def unstashDockerImages(target){
	if (!params.publish) {  // when publish this is noop since we can retrieve the image from registry
		switch (target) {
			case 'dev-jdk-8-base':
			case 'dev-jdk-others-base':
				unstashDockerImage("dev-base")
				break
			case 'dev-python-base':
				unstashDockerImage("dev-jdk-8-base")
				break
			case 'dev-r-base':
			case 'dev-mojocompat':
				unstashDockerImage("dev-python-3.7")
				break
			case 'dev-build-base':
				unstashDockerImage("dev-r-4.5.2")
				break
			case 'dev-build-hadoop':
				unstashDockerImage("dev-build-base")
				break
			case 'dev-release':
				unstashDockerImage("dev-build-hadoop")
				break
			default:
				if (target.startsWith('dev-python-')) {
					unstashDockerImage("dev-python-base")
				} else if (target.startsWith('dev-r-')) {
					unstashDockerImage("dev-r-base")
				} else if (target ==~ /dev-jdk-\d+/) { 
					unstashDockerImage("dev-jdk-others-base")
				}
				break
		}
	}
}


def pipelineContext

ansiColor('xterm') {
    timestamps {
        def checkoutStageName = 'Checkout and Init'
        buildSummary.stageWithSummary(checkoutStageName) {
            node (DEFAULT_LABEL) {
                buildSummary.refreshStage(checkoutStageName)
                cleanWs()
                def scmEnv = checkout scm

                def final pipelineContextFactory = load('scripts/jenkins/groovy/pipelineContext.groovy')
                pipelineContext = pipelineContextFactory('.', 'MODE_BUILD_DOCKER', scmEnv, true)
                def currentVersion = pipelineContext.getBuildConfig().getDefaultImageVersion()
                currentBuild.displayName = "v${currentVersion} #${currentBuild.number}"

                // if (params.publish && !params.forceBuild) {
				// 	dockerLoginH2OECR()
                //     if (imageExistsH2OECR("h2o-3/dev-base", "${currentVersion}")) {
                //         error "Image already exists. Use forceBuild if you wish to overwrite."
                //     }
                // } else {
                //     manager.addBadge("warning.gif", "forceBuild enabled")
                // }

                dir('docker') {
                    withCredentials([[$class: 'AmazonWebServicesCredentialsBinding', accessKeyVariable: 'AWS_ACCESS_KEY_ID', credentialsId: 'H2O-AWS_CT-JENKINS-SHARED-SERVICE-PROD', secretKeyVariable: 'AWS_SECRET_ACCESS_KEY']]) {
                        docker.image('d3fk/s3cmd').inside("--entrypoint='' -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}") {
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-8/x64-linux/jdk-8u311-linux-x64.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-10/x64-linux/jdk-10.0.2_linux-x64_bin.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-11/x64-linux/jdk-11.0.3_linux-x64_bin.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-12/x64-linux/jdk-12.0.2_linux-x64_bin.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-13/x64-linux/jdk-13.0.2_linux-x64_bin.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-14/x64-linux/jdk-14.0.2_linux-x64_bin.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-15/x64-linux/jdk-15.0.2_linux-x64_bin.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-16/x64-linux/jdk-16.0.2_linux-x64_bin.tar.gz"
                            sh "s3cmd get s3://artifacts.h2o.ai/releases/oracle/jdk-17/x64-linux/jdk-17.0.0_linux-x64_bin.tar.gz"
                        }
                    }
                    stash name: 'docker-sources'
                }
            }
        }

		PYTHON_VERSIONS = pipelineContext.getBuildConfig().PYTHON_VERSIONS
		R_VERSIONS = pipelineContext.getBuildConfig().R_VERSIONS
		
        buildInParallel(pipelineContext, ['dev-base'])
        buildInParallel(pipelineContext, ['dev-jdk-8-base', 'dev-jdk-others-base'])

        buildInParallel(pipelineContext, ['dev-python-base'])
        def pythonImages = []
        PYTHON_VERSIONS.each {version ->
            pythonImages += "dev-python-${version}"
        }
        buildInParallel(pipelineContext, pythonImages)

        buildInParallel(pipelineContext, ['dev-r-base'])
        def componentImages = []
        R_VERSIONS.each {version ->
            componentImages += "dev-r-${version}"
        }
        JDK_VERSIONS.each {version ->
            componentImages += "dev-jdk-${version}"
        }
        componentImages += "dev-mojocompat"
        buildInParallel(pipelineContext, componentImages)

        buildInParallel(pipelineContext, ["dev-build-base"])
        buildInParallel(pipelineContext, ["dev-build-hadoop"],)
        buildInParallel(pipelineContext, ["dev-release"])

        def componentJdkImages = []
        PYTHON_VERSIONS.each {version ->
            JDK_VERSIONS_PY_R.each {jdkVersion ->
                componentJdkImages += "dev-python-${version}-jdk-${jdkVersion}"
            }
        }
        R_VERSIONS.each {version ->
            JDK_VERSIONS_PY_R.each {jdkVersion ->
                componentJdkImages += "dev-r-${version}-jdk-${jdkVersion}"
            }
        }
        buildInParallel(pipelineContext, componentJdkImages)
    }
}

private buildInParallel(final pipelineContext, final targets) {
    parallel(targets.collectEntries { target ->
        [
                (target): makeStage(pipelineContext, target),
        ]
    })
}

private Closure makeStage(final pipelineContext, String target) {
    return {
        def buildImageStageName = "Build $target"
        buildSummary.stageWithSummary(buildImageStageName) {
            final String noCache = params.noCache ? "--no-cache" : ""
            node (DEFAULT_LABEL) {
                buildSummary.refreshStage(buildImageStageName)
                dir (target) {
                    deleteDir()
                    dockerLoginH2OECR()
					unstashDockerImages(target)
                    if(!imageExistsH2OECR("h2o-3/${target}", "${pipelineContext.getBuildConfig().getDefaultImageVersion()}") || params.forceBuild){
                        unstash 'docker-sources'
                        sh """
                            make \
                                CI=1 \
                                H2O_BRANCH=${params.h2oBranch} \
                                VERSION=${pipelineContext.getBuildConfig().getDefaultImageVersion()} \
                                PUSH=${params.publish} \
                                NO_CACHE=${noCache} \
                                ${target}
                        """
                    } else {
                        println "${H2O_ECR}/h2o-3/${target} already exists"
                    }  
                    stashDockerImages(target)
                }
            }
        }
    }
}

def dockerLoginH2OECR(){
    withCredentials([[$class: 'AmazonWebServicesCredentialsBinding', credentialsId: "H2O-AWS_CT-JENKINS-H2O-ECR"]]) {
        H2O_ECR_PWD = sh(script: 'aws ecr get-login-password --region us-east-1', returnStdout: true).trim()
        sh """
            docker login --username AWS --password $H2O_ECR_PWD $H2O_ECR
        """
    }
}

def imageExistsH2OECR(String repository, String tag) {
	try {
		def result = "-1"
		withCredentials([[$class: 'AmazonWebServicesCredentialsBinding', credentialsId: "H2O-AWS_CT-JENKINS-H2O-ECR"]]) {
			result = sh(
				script: "aws ecr batch-get-image --registry-id=353750902984 --repository-name='${repository}' --image-ids=imageTag='${tag}' --query 'images[].imageId.imageTag' --output text",
				returnStdout: true
			).trim()
		}
    
		return result == tag
	} catch (Exception e) {
		// We didn't find the image. Exception can be thrown due to status code 254 that corresponds to "The repository with name 'h2o-3/xyz' does not exist in the registry with id 'xxx'"
		return false
	}
}

