// 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.
{% include "utils/base.groovy.j2" with context %}
{% import 'utils/macros.j2' as m with context -%}

def ecr_push(full_name) {
  aws_account_id = sh(
    returnStdout: true,
    script: 'aws sts get-caller-identity | grep Account | cut -f4 -d\\"',
    label: 'Get AWS ID'
  ).trim()

  def ecr_name = "${aws_account_id}.{{ aws_ecr_url }}/${full_name}"
  try {
    withEnv([
      "AWS_ACCOUNT_ID=${aws_account_id}",
      'AWS_DEFAULT_REGION={{ aws_default_region }}',
      "AWS_ECR_REPO=${aws_account_id}.{{ aws_ecr_url }}"]) {
      sh(
        script: '''
          set -eux
          aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $AWS_ECR_REPO
        ''',
        label: 'Log in to ECR'
      )
      sh(
        script: """
          set -x
          . ${jenkins_scripts_root}/retry.sh
          docker tag ${full_name} \$AWS_ECR_REPO/${full_name}
          retry 5 docker push \$AWS_ECR_REPO/${full_name}
        """,
        label: 'Upload image to ECR'
      )
    }
  } finally {
    withEnv([
      "AWS_ACCOUNT_ID=${aws_account_id}",
      'AWS_DEFAULT_REGION={{ aws_default_region }}',
      "AWS_ECR_REPO=${aws_account_id}.{{ aws_ecr_url }}"]) {
      sh(
        script: 'docker logout $AWS_ECR_REPO',
        label: 'Clean up login credentials'
      )
    }
  }
  return ecr_name
}

def build_image(image_name) {
  hash = sh(
    returnStdout: true,
    script: 'git log -1 --format=\'%h\''
  ).trim()
  def full_name = "${image_name}:${env.BRANCH_NAME}-${hash}-${env.BUILD_NUMBER}".replace('/', '_')
  sh(
    script: "${docker_build} ${image_name} --spec ${full_name}",
    label: 'Build docker image'
  )
  return ecr_push(full_name)
}

def update_docker(ecr_image, hub_image) {
  if (ecr_image == null) {
    sh("image was not rebuilt, skipping")
    return
  }
  if (!ecr_image.contains("amazonaws.com")) {
    sh("echo \"Skipping '${ecr_image}' -> '${hub_image}' since it doesn\'t look like an ECR image\"")
    return
  }
  docker_init(ecr_image)
  sh(
    script: """
    set -eux
    . ${jenkins_scripts_root}/retry.sh
    docker tag \
      ${ecr_image} \
      ${hub_image}
    retry 5 docker push ${hub_image}
    """,
    label: "Update ${hub_image} on Docker Hub",
  )
}

def deploy() {
  stage('Deploy') {
    if (env.BRANCH_NAME == 'main') {
      parallel(
        {% call m.deploy_step(
          name="Upload built Docker images",
          feature_flag="env.DEPLOY_DOCKER_IMAGES == 'yes' && rebuild_docker_images && upstream_revision != null",
          ws="tvm/deploy-docker",
        ) %}
          init_git()
          try {
            withCredentials([string(
              credentialsId: 'dockerhub-tlcpackstaging-key',
              variable: 'DOCKERHUB_KEY',
            )]) {
              sh(
                script: 'docker login -u tlcpackstaging -p ${DOCKERHUB_KEY}',
                label: 'Log in to Docker Hub',
              )
            }
            def date_Ymd_HMS = sh(
              script: 'python3 -c \'import datetime; print(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))\'',
              label: 'Determine date',
              returnStdout: true,
            ).trim()
            def tag = "${date_Ymd_HMS}-${upstream_revision.substring(0, 8)}"
            {% for image in images %}
            update_docker(built_{{ image.name }}, "tlcpackstaging/{{ image.name }}:${tag}")
            {% endfor %}
          } finally {
            sh(
              script: 'docker logout',
              label: 'Clean up login credentials'
            )
          }
        {% endcall %}
        {% call m.deploy_step(
          name="Tag tlcpackstaging to tlcpack",
          feature_flag="env.DEPLOY_DOCKER_IMAGES == 'yes'",
          ws="tvm/tag-images",
        ) %}
          init_git()
          withCredentials([string(
            credentialsId: 'dockerhub-tlcpack-key',
            variable: 'TLCPACK_TOKEN',
          )]) {
            try {
              sh(
                script: 'echo $TLCPACK_TOKEN | docker login --username octomldriazati --password-stdin',
                label: 'Log in to Docker Hub'
              )
              {% for image in images %}
              if ({{ image.name }}.contains("tlcpackstaging")) {
                // Push image to tlcpack
                def tag = {{ image.name }}.split(":")[1]
                sh(
                  script: """
                    set -eux
                    . ${jenkins_scripts_root}/retry.sh
                    docker pull tlcpackstaging/{{ image.name }}:${tag}
                    docker tag tlcpackstaging/{{ image.name }}:${tag} tlcpack/{{ image.name.replace("_", "-") }}:${tag}
                    retry 5 docker push tlcpack/{{ image.name.replace("_", "-") }}:${tag}
                  """,
                  label: 'Tag tlcpackstaging/{{ image.name }} image to tlcpack',
                )
              }
              {% endfor %}
            } finally {
              sh(
                script: 'docker logout',
                label: 'Clean up login credentials'
              )
            }
          }
        {% endcall %}
      )
    }
  }
}



if (rebuild_docker_images) {
  stage('Docker Image Build') {
    parallel(
    {% for image in images %}
      '{{ image.name }}': {
        node('{{ image.platform }}') {
          timeout(time: max_time, unit: 'MINUTES') {
            init_git()
            // We're purposefully not setting the built image here since they
            // are not yet being uploaded to tlcpack
            // {{ image.name }} = build_image('{{ image.name }}')
            built_{{ image.name }} = build_image('{{ image.name }}');
          }
        }
      },
    {% endfor %}
    )
  }

  deploy()
}