from airflow import DAG

from template.dag_template import build_dag
from template.utils import hours_ago
from airflow.models import Variable
from datetime import datetime

# <cluster name>_<version id>
dag_id = 'exchange-trade-summary_v002'
SSPUI_JOB_EXECUTION_MODE = Variable.get("sspui_job_execution_mode")
AWS_REGION=Variable.get("aws_region")

# SSPUI_JOB_EXECUTION_MODE=0, db full load
# SSPUI_JOB_EXECUTION_MODE=1, load data for airflow execution time and execution time + 1 (default behavior)
# SSPUI_JOB_EXECUTION_MODE=2, load data for airflow execution time (backfill behavior)

ExportVioohHourLevelReportToDbV2_script_args = f"\"$ENV#0.0.2#db-import-viooh-hourly-ui-report-v2#{SSPUI_JOB_EXECUTION_MODE}#$EXECUTION_DATETIME#$ARTIFACT_BUCKET#scripts/exchange-trade-summary/sspui_reporting#deal#{AWS_REGION}\""
ExportVioohHourlyCreativeLevelReportToDb_script_args = f"\"$ENV#0.0.2#db-import-viooh-hourly-creative-report#{SSPUI_JOB_EXECUTION_MODE}#$EXECUTION_DATETIME#$ARTIFACT_BUCKET#scripts/exchange-trade-summary/sspui_reporting#creative#{AWS_REGION}\""

emr_steps = """[
      {
        "step-name": "SupplyDenormalizationJob",
        "config-json": [
          {"spark.driver.memory":"6g"}
        ],
        "main-class": "com.viooh.smex.supply.denormalization.SupplyDataDenormalization",
        "artifact": "supply-denormalization"
      },
      {
        "step-name": "DemandDenormalizationJob",
        "config-json": [
          {"spark.driver.memory":"6g"}
        ],
        "main-class": "com.viooh.smex.demand.denormalization.DemandDataDenormalization",
        "artifact": "demand-denormalization"
      },
      {
        "step-name": "DealSyncDenormalizationJob",
        "config-json": [
          {"spark.driver.memory":"6g"}
        ],
        "main-class": "com.viooh.smex.dealsync.denormalization.DealSyncDenormalization",
        "artifact": "dealsync-denormalization"
      },
      {
        "step-name": "ModerationDenormalizationJob",
        "config-json": [
          {"spark.driver.memory":"6g"}
        ],
        "main-class": "com.viooh.smex.moderation.denormalization.ModerationDenormalization",
        "artifact": "moderation-denormalization",
        "jars": "/usr/lib/spark/external/lib/spark-avro.jar"
      },
      {
        "step-name": "DealDenormalizationJob",
        "config-json": [
          {"spark.driver.memory":"6g"}
        ],
        "main-class": "com.viooh.smex.deal.denormalization.DealDenormalization",
        "artifact": "deal-denormalization",
        "jars": "/usr/lib/spark/external/lib/spark-avro.jar"
      },
      {
        "step-name": "ExchangeTradeSummaryJob",
        "config-json": [
          {"spark.driver.memory":"6g"},
          {"spark.serializer":"org.apache.spark.serializer.KryoSerializer"},
          {"spark.task.maxFailures":"20"},
          {"spark.yarn.maxAppAttempts":"10"},
          {"spark.stage.maxConsecutiveAttempts":"20"},
          {"spark.app.config.key":"exchange-trade-summary"}
        ],
        "main-class": "com.viooh.smex.trade.summary.ExchangeTradeSummary",
        "artifact": "exchange-trade-summary"
      },
      {
        "step-name": "CopySmartExchangeMediaOwnerReport",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/1.3.0/job.sh",
        "script-args": "$ENV#$NEXT_EXECUTION_DATE#$NEXT_EXECUTION_HOUR#mediaowner-trade-report#trade-summary-reports#v1#csv"
      },
      {
        "step-name": "CopyExchangeLiteMediaOwnerReport",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/1.3.0/job.sh",
        "script-args": "$ENV#$EXECUTION_DATE#$EXECUTION_HOUR#exchange-lite-mediaowner-trade-report#trade-summary-reports#v1#csv"
      },
      {
        "step-name": "CopySmartExchangeFrameLevelMediaOwnerReport",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/1.3.0/job.sh",
        "script-args": "$ENV#$NEXT_EXECUTION_DATE#$NEXT_EXECUTION_HOUR#frame-level-mediaowner-trade-report#frame-level-media-owner-reports#v1#csv"
      },
      {
        "step-name": "CopyExchangeLiteFrameLevelMediaOwnerReport",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/1.3.0/job.sh",
        "script-args": "$ENV#$EXECUTION_DATE#$EXECUTION_HOUR#frame-level-exchange-lite-mediaowner-trade-report#frame-level-media-owner-reports#v1#csv"
      },
            {
        "step-name": "CopySmartExchangeMediaOwnerReportV2",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/1.3.0/job.sh",
        "script-args": "$ENV#$NEXT_EXECUTION_DATE#$NEXT_EXECUTION_HOUR#mediaowner-trade-report#trade-summary-reports#v2#json"
      },
      {
        "step-name": "CopyExchangeLiteMediaOwnerReportV2",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/1.3.0/job.sh",
        "script-args": "$ENV#$EXECUTION_DATE#$EXECUTION_HOUR#exchange-lite-mediaowner-trade-report#trade-summary-reports#v2#json"
      },
      {
        "step-name": "ExportVioohHourLevelReportToDbV2",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/sspui_reporting/viooh_sspui_hourly_aggr_report_v2_installer.sh",
        "script-args": """ + ExportVioohHourLevelReportToDbV2_script_args + """
      },
      {
        "step-name": "ExportVioohHourlyCreativeLevelReportToDb",
        "jar-location":"s3://elasticmapreduce/libs/script-runner/script-runner.jar",
        "script-file":"scripts/exchange-trade-summary/sspui_reporting/viooh_sspui_hourly_aggr_report_v2_installer.sh",
        "script-args": """ + ExportVioohHourlyCreativeLevelReportToDb_script_args + """
      }
]"""

# cluster level parameters (optional)
cluster_args = {
    "master-instance-types": "m5.xlarge,m5.2xlarge,m5.4xlarge",
    "core-instance-types": "m5.xlarge,m5.2xlarge,m5.4xlarge",
    "task-instance-types": "m5.xlarge,m5.2xlarge,m5.4xlarge",
    "core-instance-capacity": 5,
    "task-instance-capacity": 0,
    "emr-version": "emr-5.35.0",
    "cluster-configurations": [
        {
            'Classification': 'yarn-site',
            'Properties': {
                'yarn.resourcemanager.am.max-attempts': '2'
            }
        }
    ]
}

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

dag = DAG(
    dag_id,
    schedule_interval='@hourly',  # cron expression
    default_args=dag_args)

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