from airflow import DAG
from template.utils import set_config_variable_in_emr_steps
from airflow.utils.dates import days_ago
from template.dag_template import build_dag
from airflow.models import Variable

# DAG specific parameters
config = Variable.get("primary_audience_global_conf", deserialize_json=True)

ENV = Variable.get("env")
config['ENV'] = ENV.lower()
COUNTRY = config["COUNTRY"]
MEDIA_OWNER = config["MEDIA_OWNER"]
AUD_VERSION = config["AUD_VERSION"]

# primary-audience-global job checks
cluster_name = 'primary-audience-global-' + MEDIA_OWNER + '-' + AUD_VERSION
dag_id = 'primary-audience-global_v003'

BUCKET_OVERRIDE = config["BUCKET_OVERRIDE"]
CATEGORY_GROUP_NAME = config["CATEGORY_GROUP_NAME"]
CATEGORY_GROUP_ID = config["CATEGORY_GROUP_ID"]
BUCKET_NAME = "s3a://viooh-datashare-jcd-" + COUNTRY.lower()
CATEGORY_NAME = COUNTRY
CATEGORY_ID = "001"

if len(BUCKET_OVERRIDE) > 0:
    BUCKET_NAME = BUCKET_OVERRIDE
else:
    print("Bucket not overrided using - " + BUCKET_NAME)
config['BUCKET_NAME'] = BUCKET_NAME

if len(CATEGORY_GROUP_NAME) > 0:
    CATEGORY_NAME = CATEGORY_GROUP_NAME
else:
    print("Category group name not overrided using isocode - " + COUNTRY)
config['CATEGORY_NAME'] = CATEGORY_NAME

if len(CATEGORY_GROUP_ID) > 0:
    CATEGORY_ID = CATEGORY_GROUP_ID
else:
    print("Category group id not overrided - " + CATEGORY_ID)
config['CATEGORY_ID'] = CATEGORY_ID

global_job_steps = """{
        "step-name": "PrimaryAudienceGlobal-integrity",
        "py-script": "jobs/integrity_job.py",
        "config-file": "configs/primary-audience-global/1.0.0/job.json",
        "config-json": [
          {"spark.app.audience.country":"$COUNTRY"},
          {"spark.app.audience.media.owner":"$MEDIA_OWNER"},
          {"spark.app.audience.version":"$AUD_VERSION"},
          {"spark.app.audience.process.dynamic":"$PROCESS_DYNAMIC"},
          {"spark.app.audience.process.static":"$PROCESS_STATIC"},
          {"spark.app.audience.bucket.name":"$BUCKET_NAME"},
          {"spark.app.audience.digital.file":"$DIGITAL_FILE_TO_PROCESS"},
          {"spark.app.audience.static.file":"$STATIC_FILE_TO_PROCESS"},
          {"spark.app.audience.demographic.file":"$DEMOGRAPHIC_FILE_TO_PROCESS"},
          {"spark.app.audience.check.demo":"$DEMOGRAPHIC_CHECK"}
        ],
        "artifact": "AudiencesPipeline"
      },
      {
        "step-name": "PrimaryAudienceGlobal-load",
        "py-script": "jobs/etl_job.py",
        "config-file": "configs/primary-audience-global/1.0.0/job.json",
        "config-json": [
          {"spark.app.audience.country":"$COUNTRY"},
          {"spark.app.audience.media.owner":"$MEDIA_OWNER"},
          {"spark.app.audience.version":"$AUD_VERSION"},
          {"spark.app.audience.process.dynamic":"$PROCESS_DYNAMIC"},
          {"spark.app.audience.process.static":"$PROCESS_STATIC"},
          {"spark.app.audience.category.group.name":"$CATEGORY_NAME"},
          {"spark.app.audience.category.group.id":"$CATEGORY_ID"}
        ],
        "artifact": "AudiencesPipeline"
}"""

# primary-audience-global-cassandra job checks

if COUNTRY.lower() in ('be', 'es', 'dk', 'nl'):
    KEYSPACE = COUNTRY.lower()
else:
    KEYSPACE = MEDIA_OWNER.lower()
config['KEYSPACE'] = KEYSPACE

if COUNTRY.lower() == 'uk':
    cassandra_steps = """{
        "step-name": "RegenerateRoute",
        "config-file": "configs/primary-audience-uk-cassandra/1.0.0/job.json",
        "config-json": [
                        {"spark.app.primary.audience.cassandra.keyspace": "uk_audience"},
                        {"spark.app.route.regenerate.date": "$AUD_UPLOAD_DATE"},
                        {"spark.app.route.regenerate.version": "$AUD_VERSION"}
                    ],
        "artifact": "RegenerateRoute",
        "packages": "com.datastax.spark:spark-cassandra-connector_2.11:2.3.0",
        "py-script": "com/viooh/routegen/route_regenerate_v1.py"
    }"""
elif COUNTRY.lower() == 'hk' and MEDIA_OWNER.lower() == 'jcdecaux':
    cassandra_steps = """{
        "step-name": "RegenerateRoute",
        "config-file": "configs/primary-audience-global-cassandra/1.0.0/job.json",
        "config-json": [
                        {"spark.app.primary.audience.base.path":"s3a://$AUDIENCE_BUCKET/$COUNTRY/$MEDIA_OWNER/$AUD_VERSION/raw/%date%/dynamic/"},
                        {"spark.app.primary.audience.date":"$AUD_UPLOAD_DATE"},
                        {"spark.app.primary.audience.cassandra.keyspace":"jcdecaux_hk_airport_audience"},
                        {"spark.app.primary.audience.version":"$AUD_VERSION"}
                    ],
        "artifact": "RegenerateRoute",
        "packages": "com.datastax.spark:spark-cassandra-connector_2.11:2.3.0",
        "py-script": "com/viooh/routegen/primary_audience_loader_cassandra.py"
      }"""
else:
    cassandra_steps = """{
        "step-name": "RegenerateRoute",
        "config-file": "configs/primary-audience-global-cassandra/1.0.0/job.json",
        "config-json": [
                        {"spark.app.primary.audience.base.path": "s3a://$AUDIENCE_BUCKET/$COUNTRY/$MEDIA_OWNER/$AUD_VERSION/raw/%date%/dynamic/"},
                        {"spark.app.primary.audience.date": "$AUD_UPLOAD_DATE"},
                        {"spark.app.primary.audience.cassandra.keyspace": "$KEYSPACE_audience"},
                        {"spark.app.primary.audience.version": "$AUD_VERSION"}
                    ],
        "artifact": "RegenerateRoute",
        "packages": "com.datastax.spark:spark-cassandra-connector_2.11:2.3.0",
        "py-script": "com/viooh/routegen/primary_audience_loader_cassandra.py"
    }"""

# primary-audience-global-masterdb job checks
MODELING_STEPS = """{
        "step-name": "MathModelV1",
        "config-file": "configs/primary-audience-uk-masterdb/1.0.0/job.json",
        "config-json": [
                        {"spark.route.version": "$AUD_VERSION"},
                        {"spark.date.daytouplaodtomysql": "$AUD_UPLOAD_DATE"}
                      ],
        "main-class": "com.viooh.restoremasterdb.MathModelV1",
        "artifact": "restoremasterdb_2.11"
      },
    {
        "step-name": "RestoreMasterDBAuto",
        "config-file": "configs/primary-audience-uk-masterdb/1.0.0/job.json",
        "config-json": [
                        {"spark.route.version": "$AUD_VERSION"},
                        {"spark.external.db.database": "$AUD_DATABASE"},
                        {"spark.date.daytouplaodtomysql": "$AUD_UPLOAD_DATE"}
                      ],
        "main-class": "com.viooh.restoremasterdb.RestoreMasterDBAuto",
        "artifact": "restoremasterdb_2.11"
}"""

LOADING_STEP_MYSQL = """{
        "step-name": "SaveToDataBase",
        "config-file": "configs/primary-audience-uk-masterdb/1.0.0/job.json",
        "config-json": [
                        {"spark.route.version": "$AUD_VERSION"},
                        {"spark.external.db.database": "$AUD_DATABASE"},
                        {"spark.date.daytouplaodtomysql": "$AUD_UPLOAD_DATE"}
                      ],
        "main-class": "com.viooh.restoremasterdb.SaveToDataBase",
        "artifact": "restoremasterdb_2.11",
        "packages": "mysql:mysql-connector-java:5.1.46"
}"""

LOADING_STEP_CASSANDRA_IS_ESTIMATED = """{
        "step-name": "SaveToCassandra",
        "config-file": "configs/primary-audience-uk-masterdb/1.0.0/job.json",
        "config-json": [
                        {"spark.route.version": "$AUD_VERSION"},
                        {"spark.external.cassandra.keyspace": "uk_audience"},
                        {"spark.date.daytouplaodtomysql": "$AUD_UPLOAD_DATE"}
                      ],
        "main-class": "com.viooh.restoremasterdb.SaveToCassandra",
        "artifact": "restoremasterdb_2.11",
        "packages": "datastax:spark-cassandra-connector:2.3.0-s_2.11"
}"""

if COUNTRY.lower() == 'uk':
    # For uk only
    if ENV.lower() == 'prd':
        masterdb_steps = LOADING_STEP_MYSQL + "," + LOADING_STEP_CASSANDRA_IS_ESTIMATED
    else:
        # Run modeling + loading on UAT
        masterdb_steps = MODELING_STEPS + "," + LOADING_STEP_MYSQL + "," + LOADING_STEP_CASSANDRA_IS_ESTIMATED
else:
    # For all other countries
    masterdb_steps = """{
        "step-name": "SaveToDataBase",
        "config-file": "configs/primary-audience-global-masterdb-$ENV/1.0.0/job.json",
        "config-json": [
                        {"spark.route.version": "$AUD_VERSION"},
                        {"spark.location.outputFoldermdb":"s3a://$AUDIENCE_BUCKET/$COUNTRY/$MEDIA_OWNER/$AUD_VERSION/"},
                        {"spark.external.db.database":"$AUD_DATABASE"},
                        {"spark.date.daytouplaodtomysql":"$AUD_UPLOAD_DATE"}
                    ],
        "main-class": "com.viooh.restoremasterdb.SaveToDataBase",
        "artifact": "restoremasterdb_2.11",
        "packages": "mysql:mysql-connector-java:5.1.46"
    }"""


# emr_step containing all audience job steps
emr_steps = "[" + global_job_steps + "," + cassandra_steps + "," + masterdb_steps + "]"

# cluster level parameters
cluster_args = {
    "cluster-name": cluster_name,
    "audience-config-file": "primary_audience_global_conf",
    "master-instance-types": "m5.2xlarge,m5.4xlarge",
    "core-instance-types": "m5.2xlarge,m5.4xlarge",
    "task-instance-types": "m5.2xlarge,m5.4xlarge",
    "core-instance-capacity": 5,
    "task-instance-capacity": 0,
    "job-type": "batch",
    "emr-version": "emr-5.21.0"
}

# set config variables in emr-steps
emr_steps = set_config_variable_in_emr_steps(emr_steps, config)

# dag parameter
dag_args = {
    'owner': 'data.engineers@viooh.com',
    'start_date': days_ago(1)
}

dag = DAG(
    dag_id,
    schedule_interval=None,  # cron expression
    default_args=dag_args)

build_dag(emr_steps=emr_steps, dag=dag, cluster_args=cluster_args)