import datetime as dt

import airflow
from airflow.models import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.latest_only_operator import LatestOnlyOperator
from airflow.operators.python_operator import PythonOperator

from airflow.operators.docker_operator import DockerOperator
from airflow.operators.bash_operator import BashOperator

from airflow.utils.dates import days_ago
from pprint import pprint

DAG_ID = "MA_scraper"  # must change before uploading to airflow !!!
START_DATE = days_ago(30)
# SCHEDULE_INTERVAL = dt.timedelta(hours=12)
# SCHEDULE_INTERVAL = "@daily"
SCHEDULE_INTERVAL = "5 0 * * 2-6"

MAX_ACTIVE_RUN = 1
MAX_CONCURRENT_TASKS = 3
RETRY_COUNT = 3
RETRY_DELAY = dt.timedelta(minutes=5)

EMAIL = "mobinalhassan1@gmail.com"
OWNER = "datascope"

default_args = {
    'owner': OWNER,
    'depends_on_past': False,
    'start_date': START_DATE,
    'email': [EMAIL],
    'email_on_failure': True,
    'email_on_retry': True,
    'retries': RETRY_COUNT,
    'retry_delay': RETRY_DELAY,
    # 'execution_timeout': timedelta(seconds=300),
    # 'on_failure_callback': some_function,
    # 'on_success_callback': some_other_function,
    # 'on_retry_callback': another_function,
    'trigger_rule': 'all_done'
}

dag = DAG(
    dag_id=DAG_ID,
    schedule_interval=SCHEDULE_INTERVAL,
    start_date=START_DATE,
    max_active_runs=MAX_ACTIVE_RUN,
    catchup=False,
    default_args=default_args,
    concurrency=MAX_CONCURRENT_TASKS

)

latest_only = LatestOnlyOperator(task_id='my_latest_only', dag=dag)
# Pre-requisite
prerequisite_scheduler = DockerOperator(
    task_id='Prerequisite-scheduler',
    image='registry.gitlab.com/aintelproject/prerequisite:latest',
    auto_remove=True,
    force_pull=True,
    dag=dag
)


# N2
norfolk_scheduler = DockerOperator(
    task_id='N2-T1-Norfolk-County-scheduler',
    image='registry.gitlab.com/aintelproject/n2-t1-norfolk-county:latest',
    auto_remove=True,
    force_pull=True,
    environment={
        "AWS_ACCESS_KEY_ID": "{{var.value.AWS_ACCESS_KEY_ID_BOSTON}}",
        "AWS_SECRET_ACCESS_KEY": "{{var.value.AWS_SECRET_ACCESS_KEY_BOSTON}}",
        "BUCKET_NAME": "{{var.value.BUCKET_NAME_BOSTON}}",
        "REGION": "{{var.value.REGION_BOSTON}}",
    },
    dag=dag
)
# N3
north_worcester_scheduler = DockerOperator(
    task_id='N3-T1-North-Worcester-County-scheduler',
    image='registry.gitlab.com/aintelproject/n3-t1-north-worcester-county:latest',
    auto_remove=True,
    force_pull=True,
    environment={
        "AWS_ACCESS_KEY_ID": "{{var.value.AWS_ACCESS_KEY_ID_BOSTON}}",
        "AWS_SECRET_ACCESS_KEY": "{{var.value.AWS_SECRET_ACCESS_KEY_BOSTON}}",
        "BUCKET_NAME": "{{var.value.BUCKET_NAME_BOSTON}}",
        "REGION": "{{var.value.REGION_BOSTON}}",
    },
    dag=dag
)

# N11
middlesex_north_country_scheduler = DockerOperator(
    task_id='N11-T2-MiddlesexNorth-Country-scheduler',
    image='registry.gitlab.com/aintelproject/n11-t2-middlesexnorth-county:latest',
    auto_remove=True,
    force_pull=True,
    environment={
        "AWS_ACCESS_KEY_ID": "{{var.value.AWS_ACCESS_KEY_ID_BOSTON}}",
        "AWS_SECRET_ACCESS_KEY": "{{var.value.AWS_SECRET_ACCESS_KEY_BOSTON}}",
        "BUCKET_NAME": "{{var.value.BUCKET_NAME_BOSTON}}",
        "REGION": "{{var.value.REGION_BOSTON}}",
    },
    dag=dag
)

# N12
middlesex_south_country_scheduler = DockerOperator(
    task_id='N12-T2-MiddlesexSouth-Country-scheduler',
    image='registry.gitlab.com/aintelproject/n12-t2-middlesexsouth-county:latest',
    auto_remove=True,
    force_pull=True,
    environment={
        "AWS_ACCESS_KEY_ID": "{{var.value.AWS_ACCESS_KEY_ID_BOSTON}}",
        "AWS_SECRET_ACCESS_KEY": "{{var.value.AWS_SECRET_ACCESS_KEY_BOSTON}}",
        "BUCKET_NAME": "{{var.value.BUCKET_NAME_BOSTON}}",
        "REGION": "{{var.value.REGION_BOSTON}}",
    },
    dag=dag
)

# N14
plymouthdeeds_county_scheduler = DockerOperator(
    task_id='N14-T2-plymouthdeeds-County-scheduler',
    image='registry.gitlab.com/aintelproject/n14-t2-plymouthdeeds-county:latest',
    auto_remove=True,
    force_pull=True,
    environment={
        "AWS_ACCESS_KEY_ID": "{{var.value.AWS_ACCESS_KEY_ID_BOSTON}}",
        "AWS_SECRET_ACCESS_KEY": "{{var.value.AWS_SECRET_ACCESS_KEY_BOSTON}}",
        "BUCKET_NAME": "{{var.value.BUCKET_NAME_BOSTON}}",
        "REGION": "{{var.value.REGION_BOSTON}}",
    },
    dag=dag
)
# N16
worcester_county_scheduler = DockerOperator(
    task_id='N16-T2-Worcester-County-scheduler',
    image='registry.gitlab.com/aintelproject/n16-t2-worcester-county:latest',
    auto_remove=True,
    force_pull=True,
    environment={
        "AWS_ACCESS_KEY_ID": "{{var.value.AWS_ACCESS_KEY_ID_BOSTON}}",
        "AWS_SECRET_ACCESS_KEY": "{{var.value.AWS_SECRET_ACCESS_KEY_BOSTON}}",
        "BUCKET_NAME": "{{var.value.BUCKET_NAME_BOSTON}}",
        "REGION": "{{var.value.REGION_BOSTON}}",
    },
    dag=dag
)
# Post-requirement
post_requisite_scheduler = DockerOperator(
    task_id='Post-Requirement-scheduler',
    image='registry.gitlab.com/aintelproject/postrequisite:latest',
    auto_remove=True,
    force_pull=True,
    environment={
        "AWS_ACCESS_KEY_ID": "{{var.value.AWS_ACCESS_KEY_ID_BOSTON}}",
        "AWS_SECRET_ACCESS_KEY": "{{var.value.AWS_SECRET_ACCESS_KEY_BOSTON}}",
        "BUCKET_NAME": "{{var.value.BUCKET_NAME_BOSTON}}",
        "REGION": "{{var.value.REGION_BOSTON}}",
    },
    dag=dag
)

print('We are in N2 N11 N12 N14 N16 Dag')

prerequisite_scheduler >> [norfolk_scheduler >> plymouthdeeds_county_scheduler >> middlesex_north_country_scheduler, middlesex_south_country_scheduler >> worcester_county_scheduler >> north_worcester_scheduler] >> post_requisite_scheduler