/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: MIT-0
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this
 * software and associated documentation files (the "Software"), to deal in the Software
 * without restriction, including without limitation the rights to use, copy, modify,
 * merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
 * permit persons to whom the Software is furnished to do so.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
 * PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
 * OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

pipeline {

    agent any

    parameters {
        string(name: "ACCOUNT_ID", description: "", defaultValue: "")
        string(name: "REGION", description: "", defaultValue: "")
        string(name: "ROLE_NAME", description: "", defaultValue: "")
        string(name: "SAGEMAKER_MODEL_PACKAGE_GROUP", description: "", defaultValue: "")
        string(name: "SAGEMAKER_INSTANCE_COUNT", description: "", defaultValue: "1")
        string(name: "SAGEMAKER_INSTANCE_TYPE", description: "", defaultValue: "ml.m5.xlarge")
    }

    stages {
        stage("Get last Approved model") {
            steps {
                withAWS(credentials: "jenkins-aws-profile") {
                    script {
                        def model_arn = sh returnStdout: true, script: """
                            aws sagemaker list-model-packages \
                            --model-package-group-name ${params.SAGEMAKER_MODEL_PACKAGE_GROUP} \
                            --model-approval-status Approved \
                            --sort-order Descending \
                            --sort-by CreationTime \
                            --max-items 1 \
                            --region ${params.REGION} \
                            | grep -o '"ModelPackageArn": "[^"]*' \
                            | grep -o '[^"]*\$'
                        """

                        model_arn = model_arn.trim()
                        model_arn = model_arn.replaceAll("\\s","")

                        def model_package_name = model_arn.tokenize("/")

                        env.model_package_name = model_package_name[-2] + "-" + model_package_name[-1]

                        echo "Model Package Name ${env.model_package_name}"

                        if(model_arn != "") {
                            env.model_package_arn = model_arn
                        } else {
                            error "No SageMaker Model trained available in the status Approved"
                        }
                    }
                }
            }
        }

        stage("Deploy SageMaker Model") {
            steps {
                withAWS(credentials: "jenkins-aws-profile") {
                    script {
                        def image_uri = sh returnStdout: true, script: """
                            aws sagemaker describe-model-package \
                            --model-package-name ${env.model_package_arn} \
                            --region ${params.REGION} \
                            | grep -o '"Image": "[^"]*' \
                            | grep -o '[^"]*\$'
                        """

                        image_uri = image_uri.trim()
                        env.image_uri = image_uri.replaceAll("\\s","")

                        echo "Image URI ${env.image_uri}"

                        def model_url = sh returnStdout: true, script: """
                            aws sagemaker describe-model-package \
                            --model-package-name ${env.model_package_arn} \
                            --region ${params.REGION} \
                            | grep -o '"ModelDataUrl": "[^"]*' \
                            | grep -o '[^"]*\$'
                        """

                        echo "Model URL ${model_url}"

                        model_url = model_url.trim()
                        env.model_url = model_url.replaceAll("\\s","")

                        echo "Model URL ${env.model_url}"
                    }
                }
            }
        }
        stage("Check SageMaker Model") {
            steps {
                withAWS(credentials: "jenkins-aws-profile") {
                    script {
                        def status = sh returnStatus: true, script: """
                           aws sagemaker describe-model \
                           --model-name ${env.model_package_name} \
                           --region ${params.REGION}
                        """

                        if(status != 0) {
                            // If not exists, create model and endpoint
                            env.created = "false"
                            echo "SageMaker Model ${env.model_package_name} not found. Creating it"
                        } else {
                            env.created = "true"
                            echo "SageMaker Model ${env.model_package_name} found"
                        }
                    }
                }
             }
        }
        stage("Create SageMaker Model") {
            steps {
                withAWS(credentials: "jenkins-aws-profile") {
                    script {
                        if(env.created == "false") {
                            sh """
                               aws sagemaker create-model \
                               --model-name ${env.model_package_name} \
                               --primary-container ContainerHostname=${env.BUILD_ID},Image=""" + env.image_uri + """,ModelDataUrl=""" + env.model_url + """,Mode='SingleModel' \
                               --execution-role-arn arn:aws:iam::${params.ACCOUNT_ID}:role/${params.ROLE_NAME} \
                               --region ${params.REGION}
                            """
                        } else {
                            echo "SageMaker Model ${env.model_package_name} already created"
                        }
                    }
                }
             }
        }
        stage("Deploy SageMaker Endpoint config") {
            steps {
                withAWS(credentials: "jenkins-aws-profile") {
                    script {
                        if(env.created == "false") {
                            sh """
                               aws sagemaker create-endpoint-config \
                               --endpoint-config-name ${env.model_package_name} \
                               --production-variants VariantName='single-model',ModelName=${env.model_package_name},InstanceType='${params.SAGEMAKER_INSTANCE_TYPE}',InitialVariantWeight=1,InitialInstanceCount=${params.SAGEMAKER_INSTANCE_COUNT} \
                               --region ${params.REGION}
                            """
                        } else {
                            echo "SageMaker Endpoint Config ${env.model_package_name} already created"
                        }
                    }
                }
             }
        }
        stage("Check SageMaker Endpoint") {
            steps {
                withAWS(credentials: "jenkins-aws-profile") {
                    script {
                        def status = sh returnStatus: true, script: """
                           aws sagemaker describe-endpoint \
                           --endpoint-name ${params.SAGEMAKER_MODEL_PACKAGE_GROUP} \
                           --region ${params.REGION}
                        """

                        if(status != 0) {
                            // If not exists, create model and endpoint
                            env.created_endpoint = "false"
                            echo "SageMaker Endpoint ${params.SAGEMAKER_MODEL_PACKAGE_GROUP} not found. Creating it"
                        } else {
                            env.created_endpoint = "true"
                            echo "SageMaker Endpoint ${params.SAGEMAKER_MODEL_PACKAGE_GROUP} found"
                        }
                    }
                }
             }
        }
        stage("Deploy SageMaker Endpoint") {
            steps {
                withAWS(credentials: "jenkins-aws-profile") {
                    script {
                        if(env.created_endpoint == "false") {
                            echo "Creating SageMaker Endpoint ${params.SAGEMAKER_MODEL_PACKAGE_GROUP}"

                            sh """
                               aws sagemaker create-endpoint \
                               --endpoint-name ${params.SAGEMAKER_MODEL_PACKAGE_GROUP} \
                               --endpoint-config-name ${env.model_package_name} \
                               --region ${params.REGION}
                            """
                        } else {
                            if(env.created == "false") {
                                echo "Updating SageMaker Endpoint ${params.SAGEMAKER_MODEL_PACKAGE_GROUP}"

                                sh """
                                   aws sagemaker update-endpoint \
                                   --endpoint-name ${params.SAGEMAKER_MODEL_PACKAGE_GROUP} \
                                   --endpoint-config-name ${env.model_package_name} \
                                   --region ${params.REGION}
                                """
                            } else {
                                echo "SageMaker Endpoint ${params.SAGEMAKER_MODEL_PACKAGE_GROUP} doesn't changed"
                            }
                        }
                    }
                }
            }
        }
    }
}