import logging
from datetime import datetime, timedelta

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.sensors.time_delta import TimeDeltaSensor
from airflow.utils.trigger_rule import TriggerRule

from common.common_tasks import BackfillDetectorBranchOperator
from common.sql_tasks import SQLTemplateOperator
from flows.youtube_monthly import config
from flows.youtube_monthly import tasks

logger = logging.getLogger(__file__)

default_args = dict(
    owner="youtube",
)

for report_name, report_config in config.reports.items():
    for account, account_config in report_config['accounts'].items():

        report_dag_args = account_config.get('dag_args', {})
        common_dag_args = dict(
            tags=[config.FLOW_NAME],
            default_args=default_args,
            schedule_interval='@monthly',
            start_date=datetime(2022, 6, 1),
            catchup=True,
            max_active_runs=1,
            dagrun_timeout=timedelta(minutes=120),
        )
        dag_args = {**common_dag_args, **report_dag_args}

        dag_id = '_'.join([config.FLOW_NAME, account.lower(), report_name])

        with DAG(
                dag_id=dag_id,
                **dag_args
        ) as dag:
            # need to register as multiple DAGs generated
            globals()[dag_id] = dag

            wait_for_data = TimeDeltaSensor(
                task_id='wait_for_data',
                delta=config.DATA_DELAY,
                mode='reschedule',
                poke_interval=60,
            )

            bootstrap = PythonOperator(
                task_id='bootstrap',
                python_callable=tasks.bootstrap,
                op_kwargs=dict(
                    report_name=report_name,
                    report_config=report_config,
                    mcn_account=account,
                )
            )
            bootstrap << wait_for_data

            detect_backfill = BackfillDetectorBranchOperator(
                task_id="detect_backfill",
                go_to_task_if_true=["move_and_extract_files"],
                go_to_task_if_false=["source_sensor"],
                threshold=config.DATA_THRESHOLD,
            )
            bootstrap >> detect_backfill

            source_sensor = S3KeySensor(
                task_id='source_sensor',
                bucket_key='{{ ti.xcom_pull("bootstrap").drop_s3_url }}',
                aws_conn_id=config.AWS_DROP_BUCKET_CONN_ID,

                poke_interval=timedelta(hours=6).total_seconds(),
                timeout=config.DATA_THRESHOLD.total_seconds(),
                mode='reschedule',
            )
            detect_backfill >> source_sensor

            move_and_extract_files = PythonOperator(
                task_id='move_and_extract_files',
                trigger_rule=TriggerRule.NONE_FAILED,
                python_callable=tasks.move_and_extract_files,
                op_kwargs=dict(
                    report_name=report_name,
                    filename='{{ ti.xcom_pull("bootstrap").drop_filename }}',
                    drop_s3_url='{{ ti.xcom_pull("bootstrap").drop_s3_url }}',
                    archive_s3_path_url='{{ ti.xcom_pull("bootstrap").archive_s3_path_url }}',
                    subreports=report_config.get('subreports'),
                )
            )
            detect_backfill >> move_and_extract_files
            source_sensor >> move_and_extract_files

            if 'subreports' in report_config:
                all_reports = report_config['subreports'].keys()
            else:
                all_reports = [None]

            for subreport_name in all_reports:
                staging_raw_table = config.stagin_raw_table_name(report_name, subreport_name)
                temp_staging_raw_table = f'temp_{staging_raw_table}'
                suffix = f'__{subreport_name}' if subreport_name else ''

                create_temp_staging_raw_table = SQLTemplateOperator(
                    task_id=f'create_temp_staging_raw_table{suffix}',
                    parameters=dict(
                        table_name=temp_staging_raw_table,
                    ),
                    template=f'create_temp_staging_raw_{config.report_id(report_name, subreport_name)}.sql',
                    template_dir=config.QUERIES_DIR,
                )
                move_and_extract_files >> create_temp_staging_raw_table

                if subreport_name:
                    filename = config.subreport_csv_filename(report_name, subreport_name)
                else:
                    # for single report reports we use original filename
                    filename = '{{ ti.xcom_pull("bootstrap").drop_filename }}'

                load_temp_staging_raw_table = SQLTemplateOperator(
                    task_id=f'load_temp_staging_raw_table{suffix}',
                    template='load_temp_staging_raw.sql',
                    template_dir=config.QUERIES_DIR,
                    parameters=dict(
                        table_name=temp_staging_raw_table,
                        s3_dir_path='{{ ti.xcom_pull("bootstrap").archive_s3_path_url }}',
                        file_name=filename,
                    ),
                    aws_conn_id=config.AWS_ARCHIVE_BUCKET_CONN_ID,
                )
                create_temp_staging_raw_table >> load_temp_staging_raw_table

                load_staging_raw_table = SQLTemplateOperator(
                    task_id=f'load_staging_raw_table{suffix}',
                    template='load_staging_raw_reports.sql',
                    template_dir=config.QUERIES_DIR,
                    parameters=dict(
                        temp_staging_raw_table_name=temp_staging_raw_table,
                        staging_raw_table_name=staging_raw_table,
                        account=account,
                        date='{{ ti.xcom_pull("bootstrap").date }}',
                    ),
                    aws_conn_id=config.AWS_ARCHIVE_BUCKET_CONN_ID,
                )
                load_temp_staging_raw_table >> load_staging_raw_table

                drop_temp_staging_raw_table = SQLTemplateOperator(
                    task_id=f'drop_temp_staging_raw_table{suffix}',
                    template='drop_table.sql',
                    template_dir=config.QUERIES_DIR,
                    parameters=dict(
                        table_name=temp_staging_raw_table,
                    ),
                )
                load_staging_raw_table >> drop_temp_staging_raw_table