import groovy.transform.Field

@Field def GITHUB_REPOSITORY = 'dbt-luminate'
@Field def AWS_REGIONS = ['us-east-1']
@Field def AWS_ACCOUNT_ID = '437795906767'
@Field def ECR_ACCOUNT_ID = '086679231553'
@Field def ECR_IMAGE_REPO = "${ECR_ACCOUNT_ID}.dkr.ecr.${AWS_REGIONS[0]}.amazonaws.com/${GITHUB_REPOSITORY}"
@Field def ECR_IMAGE_TAG_LATEST = 'latest' // Pin to a specific ECR tag to 'freeze' the code executed by scheduler
String SLACK_NOTIFICATIONS_CHANNEL = '#data-alerts'
String SNOWFLAKE_ACCOUNT = 'delphi.us-east-1'
String SNOWFLAKE_DATABASE = 'LUMINATE_MODELS'
String DBT_TARGET = 'default'

// prod
String PROD_AWS_DEPLOYMENT_ROLE = 'prod-dbt-luminate-scheduler-deploy-role'
def prodSecrets = [
    [
        id: 'prod/dbt-luminate-scheduler/SNOWFLAKE_PRIVATE_KEY',
        environmentVariable: 'SNOWFLAKE_PRIVATE_KEY',
        region: AWS_REGIONS[0],
    ],
    [
        id: 'prod/dbt-luminate-scheduler/SNOWFLAKE_KEY_PASSPHRASE',
        environmentVariable: 'SNOWFLAKE_KEY_PASSPHRASE',
        region: AWS_REGIONS[0],
    ]
]
def prodBaseEnvVars = [
    SNOWFLAKE_ACCOUNT: SNOWFLAKE_ACCOUNT,
    SNOWFLAKE_DATABASE: SNOWFLAKE_DATABASE,
    SNOWFLAKE_USER: 'PROD_LUMINATE_DBT_SCHEDULER',
    SNOWFLAKE_ROLE: 'PROD_LUMINATE_DBT_SCHEDULER_ROLE',
    SNOWFLAKE_WAREHOUSE: 'PROD_LUMINATE_DBT_SCHEDULER_WH',
    SNOWFLAKE_LIGHT_WAREHOUSE: 'PROD_ETL_WH',
    SNOWFLAKE_HEAVY_WAREHOUSE: 'PROD_LUMINATE_DBT_SCHEDULER_HEAVY_WH',
    SNOWFLAKE_SCHEMA: 'PROD',
    DBT_TARGET: DBT_TARGET,
]

// qa
String QA_AWS_DEPLOYMENT_ROLE = 'qa-dbt-luminate-scheduler-deploy-role'
def qaSecrets = [
    [
        id: 'qa/dbt-luminate-scheduler/SNOWFLAKE_PRIVATE_KEY',
        environmentVariable: 'SNOWFLAKE_PRIVATE_KEY',
        region: AWS_REGIONS[0],
    ],
    [
        id: 'qa/dbt-luminate-scheduler/SNOWFLAKE_KEY_PASSPHRASE',
        environmentVariable: 'SNOWFLAKE_KEY_PASSPHRASE',
        region: AWS_REGIONS[0],
    ]
]
def qaBaseEnvVars = [
    SNOWFLAKE_ACCOUNT: SNOWFLAKE_ACCOUNT,
    SNOWFLAKE_DATABASE: SNOWFLAKE_DATABASE,
    SNOWFLAKE_USER: 'QA_LUMINATE_DBT_SCHEDULER',
    SNOWFLAKE_ROLE: 'QA_LUMINATE_DBT_SCHEDULER_ROLE',
    SNOWFLAKE_WAREHOUSE: 'QA_LUMINATE_DBT_SCHEDULER_WH',
    SNOWFLAKE_LIGHT_WAREHOUSE: 'QA_ETL_WH',
    SNOWFLAKE_HEAVY_WAREHOUSE: 'QA_LUMINATE_DBT_SCHEDULER_HEAVY_WH',
    SNOWFLAKE_SCHEMA: 'QA',
    DBT_TARGET: DBT_TARGET,
]

def checkParams() {
    if (!params.DBT_DAILY_RUN && params.DBT_SELECT_MODELS == '') {
        error('DBT_SELECT_MODELS is required when DBT_DAILY_RUN is not checked.')
    }
}

def runDBT(String awsRole, List<Map> secrets, Map baseEnvVars = [:]) {
    withSecrets(
        awsAccountId: AWS_ACCOUNT_ID,
        awsRole: awsRole,
        secrets: secrets
    ) {
        withEcr(registries: [[accountId: ECR_ACCOUNT_ID, region: AWS_REGIONS[0]]]) {
            def dockerRunArgs = [
                '--rm',
                '--pull=always',
            ]
            def envVars = baseEnvVars + [
                SNOWFLAKE_PRIVATE_KEY: env.SNOWFLAKE_PRIVATE_KEY,
                SNOWFLAKE_KEY_PASSPHRASE: env.SNOWFLAKE_KEY_PASSPHRASE,
            ]
            if (params.DBT_DAILY_RUN) envVars['DBT_DAILY_RUN'] = 'true'
            if (params.DBT_FULL_REFRESH) envVars['DBT_FULL_REFRESH'] = 'true'
            if (params.SNOWFLAKE_WAREHOUSE) envVars['SNOWFLAKE_WAREHOUSE'] = params.SNOWFLAKE_WAREHOUSE
            if (params.DBT_SELECT_MODELS) envVars['DBT_SELECT_MODELS'] = params.DBT_SELECT_MODELS

            envVars.each {
                dockerRunArgs.add("--env ${it.key}=${dockerRun.escapeEnvValue(it.value)}")
            }

            sh """
                docker run ${dockerRunArgs.join(' ')} ${ECR_IMAGE_REPO}:${ECR_IMAGE_TAG_LATEST}
            """
        }
    }
}

pipeline {
    agent {
        label 'aws'
    }

    options {
        ansiColor('xterm')
        disableConcurrentBuilds()
        timestamps()
    }

    parameters {
        booleanParam(name: 'RUN_ON_PROD', defaultValue: false, description: 'Whether or not to run on prod.')
        booleanParam(name: 'RUN_ON_QA', defaultValue: false, description: 'Whether or not to run on qa.')
        booleanParam(name: 'DBT_FULL_REFRESH', defaultValue: false, description: 'Whether or not to run a full refresh.')
        booleanParam(name: 'DBT_DAILY_RUN', defaultValue: false, description: 'Run as a daily incremental update.')
        string(name: 'DBT_SELECT_MODELS', defaultValue: '', description: 'Specify which models should be run (--select).')
        choice(
            name: 'SNOWFLAKE_WAREHOUSE',
            choices: [
                '',
                'QA_ETL_WH',
                'QA_LUMINATE_DBT_SCHEDULER_WH',
                'QA_LUMINATE_DBT_SCHEDULER_LIGHT_WH',
                'QA_LUMINATE_DBT_SCHEDULER_HEAVY_WH',
                'PROD_ETL_WH',
                'PROD_LUMINATE_DBT_SCHEDULER_WH',
                'PROD_LUMINATE_DBT_SCHEDULER_LIGHT_WH',
                'PROD_LUMINATE_DBT_SCHEDULER_HEAVY_WH',
            ],
            description: 'Specify which warehouse to run on.'
        )
        string(name: 'SHARED_LIBRARIES_VERSION', defaultValue: 'master', description: 'The version of the Jenkins shared libraries to use. Can be a branch, tag or Git revision.')
    }

    triggers {
        parameterizedCron(
            '''
            TZ=America/New_York
            30 * * * * % RUN_ON_PROD=true;DBT_DAILY_RUN=true
            '''
        )
    }

    stages {
        stage('Load Shared Libraries') {
            steps {
                library "jenkins-global-libraries@${params.SHARED_LIBRARIES_VERSION}"
            }
        }
        stage('Run on QA') {
            when {
                expression { params.RUN_ON_QA }
            }
            steps {
                script {
                    checkParams()
                    runDBT(QA_AWS_DEPLOYMENT_ROLE, qaSecrets, qaBaseEnvVars)
                }
            }
        }
        stage('Run on PROD') {
            when {
                expression { params.RUN_ON_PROD }
            }
            steps {
                script {
                    checkParams()
                    runDBT(PROD_AWS_DEPLOYMENT_ROLE, prodSecrets, prodBaseEnvVars)
                }
            }
        }
    }
    post {
        failure {
            script {
                slackNotify channel: SLACK_NOTIFICATIONS_CHANNEL
            }
        }
        fixed {
            script {
                slackNotify channel: SLACK_NOTIFICATIONS_CHANNEL
            }
        }
    }
}