import groovy.transform.Field

@Field def GITHUB_REPOSITORY = 'dbt-accounting'
@Field def ECR_ACCOUNT_ID = '437795906767'
List<String> AWS_REGIONS = ['us-east-1']
String QA_ACCOUNT_ID = '437795906767'
String QA_DEPLOYMENT_ROLE = 'qa-dbt-accounting-deploy-role'
String QA_SESSION_NAME = "qa-dbt-accounting-session"
String QA_SNOWFLAKE_USER = 'QA_ACCOUNTING_DBT_SCHEDULER'
String QA_SNOWFLAKE_ROLE = 'QA_ACCOUNTING_DBT_SCHEDULER'
String QA_SNOWFLAKE_WAREHOUSE = 'QA_ACCOUNTING_DBT_WAREHOUSE'
@Field def SNOWFLAKE_ACCOUNT = 'sme-delphi'
@Field def SNOWFLAKE_DATABASE = 'ROYALTY_ACCOUNTING'
String PROD_ACCOUNT_ID = '437795906767'
String PROD_DEPLOYMENT_ROLE = 'prod-dbt-accounting-deploy-role'
String PROD_SESSION_NAME = 'prod-dbt-accounting-session'
String SLACK_NOTIFICATIONS_CHANNEL = '#moneyhub-alerts'
String QA_SNOWFLAKE_PRIVATE_KEY_ID = 'qa/dbt-accounting/SNOWFLAKE_PRIVATE_KEY'
String PROD_SNOWFLAKE_USER = 'PROD_ACCOUNTING_DBT_SCHEDULER'
String PROD_SNOWFLAKE_ROLE = 'PROD_ACCOUNTING_DBT_SCHEDULER'
String PROD_SNOWFLAKE_SCHEMA = 'PROD'
String PROD_SNOWFLAKE_WAREHOUSE = 'PROD_ACCOUNTING_DBT_WAREHOUSE'
String PROD_SNOWFLAKE_PRIVATE_KEY_ID = 'prod/dbt-accounting/SNOWFLAKE_PRIVATE_KEY'

def getSecretValue(String deploymentRole, String accountId, String sessionName, String region, String secretId) {
    withAWS(role: deploymentRole, roleAccount: accountId, roleSessionName: sessionName, region: region, useNode: true) {
        return sh(
            script: "aws secretsmanager get-secret-value --secret-id ${secretId} --query SecretString --output text",
            returnStdout: true
        ).trim()
    }
}

def runLintAndTest(Map args = [:]) {
    def deploymentRole = args.deploymentRole
    def accountId = args.accountId
    def sessionName = args.sessionName
    def region = args.region
    def secretId = args.secretId

    // Get the Snowflake private key from Secrets Manager
    def snowflakePrivateKey = getSecretValue(
        deploymentRole = deploymentRole,
        accountId = accountId,
        sessionName = sessionName,
        region = region,
        secretId = secretId
    )
    def snowflakeUser = args.snowflakeUser
    def snowflakeRole = args.snowflakeRole
    def snowflakeWarehouse = args.snowflakeWarehouse
    def snowflakeSchema = args.snowflakeSchema

    def envVars = [
        "SNOWFLAKE_ACCOUNT=${SNOWFLAKE_ACCOUNT}",
        "SNOWFLAKE_DATABASE=${SNOWFLAKE_DATABASE}",
        "SNOWFLAKE_USER=${snowflakeUser}",
        "SNOWFLAKE_ROLE=${snowflakeRole}",
        "SNOWFLAKE_WAREHOUSE=${snowflakeWarehouse}",
        "SNOWFLAKE_SCHEMA=${snowflakeSchema}",
        "SNOWFLAKE_PRIVATE_KEY=${snowflakePrivateKey}",
        "DBT_SELECT_MODELS=${getModifiedModels()}"
    ]

    withEnv(envVars) {
        withEcr {
            sh "make ci_lint_and_test"
        }
    }
}

def isScheduler = env.JOB_NAME.contains("scheduler")

def getModifiedModels() {
    def modifiedPaths = getModifiedPaths(
        branchName: env.CHANGE_BRANCH,
    )
    def selects = []
    for(modified in modifiedPaths) {
        def fileName = modified.split("/").last()
        if (fileName.startsWith("test_")) {  // test
            fileName = fileName.minus("test_").minus(".sql")
        } else if (fileName.startsWith("expected_")) {  // fixture
            fileName = fileName.minus("expected_").minus(".csv") + "_dbt"
        } else if (fileName.endsWith("_dbt.sql")) {  // model
            fileName = fileName.minus(".sql")
        } else {  // skip it
            fileName = null
        }

        if (fileName != null) {
            if (fileName != "abacus_fact_sales_unified_dbt") {
                fileName += "+"
            }

            selects.push(fileName)
        }
    }
    return selects.join(" ")
}

def runDbt(Map args = [:]) {
    def deploymentRole = args.deploymentRole
    def accountId = args.accountId
    def sessionName = args.sessionName
    def region = args.region
    def secretId = args.secretId
    def ecrImageTag = args.ecrImageTag

    // Get the Snowflake private key from Secrets Manager
    def snowflakePrivateKey = getSecretValue(
        deploymentRole = deploymentRole,
        accountId = accountId,
        sessionName = sessionName,
        region = region,
        secretId = secretId
    )
    def isScheduler = env.JOB_NAME.contains("scheduler")
    def deployTag = isScheduler ? ecrImageTag : env.GIT_COMMIT
    def deployImage = "${ECR_ACCOUNT_ID}.dkr.ecr.us-east-1.amazonaws.com/${GITHUB_REPOSITORY}:${deployTag}"
    def dbtTarget = args.dbtTarget
    def dbtSelectModels = isScheduler ? args.dbtSelectModels ?: '' : getModifiedModels()
    def dbtExcludeModels = args.dbtExcludeModels ?: ''
    def dbtFullRefresh = args.containsKey('dbtFullRefresh') ? args.dbtFullRefresh : true
    def snowflakeUser = args.snowflakeUser
    def snowflakeRole = args.snowflakeRole
    def snowflakeWarehouse = args.snowflakeWarehouse
    def snowflakeSchema = args.snowflakeSchema

    def envVars = [
        "DOCKER_CONFIG=${env.WORKSPACE}/.docker",
        "SNOWFLAKE_ACCOUNT=${SNOWFLAKE_ACCOUNT}",
        "SNOWFLAKE_DATABASE=${SNOWFLAKE_DATABASE}",
        "SNOWFLAKE_USER=${snowflakeUser}",
        "SNOWFLAKE_ROLE=${snowflakeRole}",
        "SNOWFLAKE_WAREHOUSE=${snowflakeWarehouse}",
        "SNOWFLAKE_SCHEMA=${snowflakeSchema}",
        "SNOWFLAKE_PRIVATE_KEY=${snowflakePrivateKey}",
        "DBT_TARGET=${dbtTarget}",
    ]

    if (dbtFullRefresh) {
        envVars.add("DBT_FULL_REFRESH=true")
    }

    if (dbtSelectModels) {
        envVars.add("DBT_SELECT_MODELS=${dbtSelectModels}")
    }

    if (dbtExcludeModels) {
        envVars.add("DBT_EXCLUDE_MODELS=${dbtExcludeModels}")
    }

    // Override DOCKER_CONFIG to directory local to workspace to ensure authentication is isolated from other jobs
    withEnv(envVars) {
        echo 'Logging in to ECR registry'
        def ecrRegion = deployImage.split(/\./)[3]
        def ecrRegistry = deployImage.split(/\//)[0]
        sh "aws ecr get-login-password --region ${ecrRegion} | docker login --username AWS --password-stdin ${ecrRegistry}"

        echo 'Pulling image'
        sh "docker pull ${deployImage}"
        echo "Running with DBT_FULL_REFRESH=${dbtFullRefresh}"
        sh "docker run -e SNOWFLAKE_ACCOUNT -e SNOWFLAKE_DATABASE -e SNOWFLAKE_USER -e SNOWFLAKE_PRIVATE_KEY -e SNOWFLAKE_ROLE -e SNOWFLAKE_WAREHOUSE -e SNOWFLAKE_SCHEMA -e DBT_TARGET -e DBT_FULL_REFRESH -e DBT_BUILD -e DBT_SELECT_MODELS -e DBT_EXCLUDE_MODELS ${deployImage}"
    }
}

pipeline {
    agent {
        label 'aws'
    }

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

    parameters {
        booleanParam(name: 'RUN_ON_QA', defaultValue: true, description: 'Whether or not to run on qa.')
        booleanParam(name: 'RUN_ON_PROD', defaultValue: true, description: 'Whether or not to run on prod.')
        booleanParam(name: 'FULL_REFRESH', defaultValue: true, description: 'Whether or not to run a full refresh.')
        string(name: 'SELECT_MODELS', defaultValue: '', description: 'Specify which models should be run (--select).')
        string(name: 'EXCLUDE_MODELS', defaultValue: 'abacus_fact_sales_unified_dbt workstation_fact_sales_unified_dbt', description: 'Specify which models should be run (--exclude).')
        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.')
    }

    stages {
        stage('Cleanup Workspace') {
            steps {
                cleanWs()
            }
        }
        stage('Checkout Code') {
            steps {
                script {
                    def scmVars = checkout scm
                    env.GIT_COMMIT = scmVars.GIT_COMMIT
                    env.BRANCH_NAME = scmVars.GIT_BRANCH.tokenize("/").last()
                }
            }
        }
        stage('Load Shared Libraries') {
            steps {
                library "jenkins-global-libraries@${params.SHARED_LIBRARIES_VERSION}"
            }
        }
        stage('Compliance Checks') {
            when {
                 expression { !isScheduler }
            }
            steps {
                complianceChecks()
            }
        }
        stage('Validate Software Catalog Definition') {
            when {
                 expression { !isScheduler }
            }
            steps {
                datadogSoftwareCatalogValidate()
            }
        }
        stage('Style Checks And Unit Tests') {
            when {
                 expression { !isScheduler }
            }
            steps {
                withEcr {
                    script {
                        def snowflakeSchema = env.BRANCH_NAME == 'master' ? 'CI_TEST' : 'CI_PR'
                        runLintAndTest(
                            deploymentRole: QA_DEPLOYMENT_ROLE,
                            accountId: QA_ACCOUNT_ID,
                            sessionName: QA_SESSION_NAME,
                            region: AWS_REGIONS[0],
                            secretId: QA_SNOWFLAKE_PRIVATE_KEY_ID,
                            snowflakeUser: QA_SNOWFLAKE_USER,
                            snowflakeRole: QA_SNOWFLAKE_ROLE,
                            snowflakeWarehouse: QA_SNOWFLAKE_WAREHOUSE,
                            snowflakeSchema: snowflakeSchema,
                        )
                    }
                }
            }
        }
        stage('Create a Release') {
            when {
                branch 'master'
                expression { !isScheduler }
            }
            steps {
                dockerToEcr awsRegions: AWS_REGIONS, ecrAccountId: ECR_ACCOUNT_ID, imageName: GITHUB_REPOSITORY, imageTag: env.GIT_COMMIT,
                    dockerBuildTarget: 'deploy'
                datadogSoftwareCatalogPublish()
            }
        }
        stage('Run on QA') {
            when {
                allOf {
                    branch 'master'
                    expression { params.RUN_ON_QA }
                    expression { !env.JOB_NAME.contains("scheduler") || (env.JOB_NAME.contains("scheduler") && params.SELECT_MODELS) }
                }
            }
            steps {
                runDbt(
                    deploymentRole: QA_DEPLOYMENT_ROLE,
                    accountId: QA_ACCOUNT_ID,
                    sessionName: QA_SESSION_NAME,
                    region: AWS_REGIONS[0],
                    ecrImageTag: 'latest',
                    secretId: QA_SNOWFLAKE_PRIVATE_KEY_ID,
                    dbtFullRefresh: params.FULL_REFRESH,
                    dbtSelectModels: params.SELECT_MODELS,
                    dbtExcludeModels: params.EXCLUDE_MODELS,
                    dbtTarget: 'qa',
                    snowflakeUser: QA_SNOWFLAKE_USER,
                    snowflakeRole: QA_SNOWFLAKE_ROLE,
                    snowflakeWarehouse: QA_SNOWFLAKE_WAREHOUSE,
                    snowflakeSchema: 'QA',
                )
            }
        }
        stage('Run on PROD') {
            when {
                allOf {
                    branch 'master'
                    expression { params.RUN_ON_PROD }
                    expression { !env.JOB_NAME.contains("scheduler") || (env.JOB_NAME.contains("scheduler") && params.SELECT_MODELS) }
                }
            }
            steps {
                runDbt(
                    deploymentRole: PROD_DEPLOYMENT_ROLE,
                    accountId: PROD_ACCOUNT_ID,
                    sessionName: PROD_SESSION_NAME,
                    region: AWS_REGIONS[0],
                    ecrImageTag: 'latest',
                    secretId: PROD_SNOWFLAKE_PRIVATE_KEY_ID,
                    dbtFullRefresh: params.FULL_REFRESH,
                    dbtSelectModels: params.SELECT_MODELS,
                    dbtExcludeModels: params.EXCLUDE_MODELS,
                    dbtTarget: 'prod',
                    snowflakeUser: PROD_SNOWFLAKE_USER,
                    snowflakeRole: PROD_SNOWFLAKE_ROLE,
                    snowflakeWarehouse: PROD_SNOWFLAKE_WAREHOUSE,
                    snowflakeSchema: 'PROD',
                )
            }
        }
    }

    post {
        regression {
            script {
                if (env.BRANCH_NAME == 'master') {
                    slackNotify channel: SLACK_NOTIFICATIONS_CHANNEL
                }
            }
        }
        fixed {
            script {
                if (env.BRANCH_NAME == 'master') {
                    slackNotify channel: SLACK_NOTIFICATIONS_CHANNEL
                }
            }
        }
    }
}