import logging
from datetime import datetime, timedelta

from airflow import DAG
from airflow.operators.dummy import DummyOperator
from airflow.operators.python import PythonOperator
from airflow.sensors.external_task import ExternalTaskSensor
from airflow.utils.task_group import TaskGroup

from common import common_config
from common.sql_tasks import SQLTemplateOperator
from flows.spotify import config

logger = logging.getLogger(__file__)

default_args = dict(
    owner=config.PROVIDER,
)


def create_load_staging_fact_table(dag):
    with TaskGroup(group_id='staging_fact_table') as task_group:
        staging_fact_table = config.staging_fact_table_template.format(
            licensor=licensor,
            date_nodash='{{ ds_nodash }}',
            feed_name=config.FLOW_NAME
        )

        create_staging_fact = SQLTemplateOperator(
            task_id='create',
            template='create_table.sql',
            template_dir=config.QUERIES_DIR,
            parameters=dict(
                table_name=staging_fact_table,
                or_replace=True,
                like='fact_analytics',
            )
        )

        load_staging_fact = SQLTemplateOperator(
            task_id='load',
            template='load_staging_fact.sql',
            template_dir=config.QUERIES_DIR,
            parameters=dict(
                staging_fact_analytics_table=staging_fact_table,
                staging_raw_table=config.reports['streams']['staging_raw'],
                storeid=config.storeid,
                feedid=config.feedid,
                reportdate='{{ ds }}',
                licensor=licensor,
            )
        )
        load_staging_fact << create_staging_fact
    return task_group


def create_fact_tables():
    with TaskGroup(group_id='fact_tables') as task_group:
        staging_fact_table = config.staging_fact_table_template.format(
            licensor=licensor,
            date_nodash='{{ ds_nodash }}',
            feed_name=config.FLOW_NAME
        )

        clean = SQLTemplateOperator(
            task_id='clean',
            template='delete_from_table.sql',
            template_dir=config.QUERIES_DIR,
            parameters=dict(
                table_name='fact_analytics',
                where=dict(
                    feedid=config.feedid,
                    storeid=config.storeid,
                    reportdate='{{ ds }}',
                    licensorid=config.licensors[licensor]['licensorid'],
                )
            )
        )

        load = SQLTemplateOperator(
            task_id='load',
            template='load_fact_analytics.sql',
            template_dir=config.QUERIES_DIR,
            parameters=dict(
                fact_table='fact_analytics',
                staging_fact_analytics_table='',
                staging_raw_table=staging_raw_table,
                reportdate='{{ ds }}',
            )
        )
        load << clean

    return task_group

for licensor in config.spotify_api_licensors:

    dag_args = dict(
        tags=[config.FLOW_NAME],
        default_args=default_args,
        schedule_interval='@daily',
        start_date=datetime(2022, 12, 24),
        catchup=True,
        dagrun_timeout=timedelta(hours=8),
    )
    dag_id = '_'.join([config.FLOW_NAME, licensor])

    with DAG(
            dag_id=dag_id,
            **dag_args
    ) as dag:

        combine_awaiter = DummyOperator(
            task_id='combine_awaiter',
        )

        for depend_report in ['streams', 'tracks', 'sub_30_sec_streams', 'aggregated_streams']:
            downstream_dag = '_'.join([config.FLOW_NAME, licensor, depend_report])

            wait_for_report = ExternalTaskSensor(
                task_id=f'wait_for_{depend_report}',
                external_dag_id=downstream_dag,
                external_task_id='load_temp_table',
                poke_interval=timedelta(hours=1).total_seconds(),
                timeout=config.DATA_THRESHOLD.total_seconds(),
                allowed_states=["success"],
                failed_states=["failed", "skipped"],
                mode="reschedule",
            )
            wait_for_report >> combine_awaiter

        # 1. load_common_tables
        with TaskGroup(group_id='load_common_tables') as task_group_load_common_tables:
            for report_name, report_config in config.reports.items():

                if report_name in config.common_reports and 'staging_raw' in report_config:
                    temp_staging_raw_table = config.temp_staging_raw_table.format(
                        date_nodash='{{ ds_nodash }}',
                        licensor=licensor,
                        report_name=report_name
                    )
                    transitional_temp_table = config.transitional_temp_table.format(
                        date_nodash='{{ ds_nodash }}',
                        licensor=licensor,
                        report_name=report_name
                    )
                    staging_raw_table = report_config['staging_raw']

                    create_temp_common_tables = SQLTemplateOperator(
                        task_id=f'create_temp_common_tables_{report_name}',
                        template=f'create_transitional_temp_staging_raw_{report_name}.sql',
                        template_dir=config.QUERIES_DIR,
                        parameters=dict(
                            transitional_temp_table=transitional_temp_table,
                        )
                    )
                    create_temp_common_tables << combine_awaiter

                    task_group_load_common_tables = SQLTemplateOperator(
                        task_id=f'load_common_tables_{report_name}',
                        template=f'load_transitional_temp_staging_raw_{report_name}.sql',
                        template_dir=config.QUERIES_DIR,
                        parameters=dict(
                            transitional_temp_table=transitional_temp_table,
                            date='{{ ds }}',
                            licensor=licensor,
                            temp_staging_raw_table=temp_staging_raw_table,
                            staging_raw_table=staging_raw_table
                        )
                    )
                    task_group_load_common_tables << create_temp_common_tables

        # 2. load_staging_raw for all reports
        with TaskGroup(group_id='load_staging_raw') as task_group_load_staging_raw:

            for report_name, report_config in config.reports.items():
                # WTF? flow.py:572   if report_name not in config.common_reports:
                if report_name in config.common_reports:
                # if 'staging_raw' not in report_config:
                    continue

                temp_streams_staging_raw_table = config.temp_staging_raw_table.format(
                    report_name='streams',
                    date_nodash='{{ ds_nodash }}',
                    licensor=licensor,
                )
                temp_users_staging_raw_table = config.temp_staging_raw_table.format(
                    report_name='users',
                    date_nodash='{{ ds_nodash }}',
                    licensor=licensor,
                )
                temp_tracks_staging_raw_table = config.temp_staging_raw_table.format(
                    report_name='tracks',
                    date_nodash='{{ ds_nodash }}',
                    licensor=licensor,
                )
                staging_raw_tracks_table = config.reports['tracks']['staging_raw']
                staging_raw_table = report_config['staging_raw']

                clean_staging_raw = SQLTemplateOperator(
                    task_id=f'clean_staging_raw_{report_name}',
                    template=f'delete_from_table.sql',
                    template_dir=config.QUERIES_DIR,
                    parameters=dict(
                        table_name=staging_raw_table,
                        where=dict(
                            download_date='{{ ds }}',
                            licensor=licensor,
                        ),
                    )
                )
                clean_staging_raw << combine_awaiter

                # << staging_raw_tracks_table
                task_group_load_staging_raw = SQLTemplateOperator(
                    task_id=f'load_staging_raw_{report_name}',
                    template=f'load_staging_raw_{report_name}.sql',
                    template_dir=config.QUERIES_DIR,
                    parameters=dict(
                        date='{{ ds }}',
                        licensor=licensor,
                        filename='HZ',  # TODO: ???
                        temp_staging_raw_table=temp_staging_raw_table,
                        staging_raw_table=staging_raw_table,
                        staging_raw_tracks_table=staging_raw_tracks_table,
                        temp_streams_staging_raw_table = temp_streams_staging_raw_table,
                        temp_users_staging_raw_table=temp_users_staging_raw_table,
                        temp_tracks_staging_raw_table=temp_tracks_staging_raw_table,
                    )
                )
                task_group_load_staging_raw << clean_staging_raw
                # load_staging_raw_aggregated_streams << load_staging_raw_tracks
                # if report_name == 'tracks':
                #     load_staging_raw_tracks = load_staging_raw
                # if report_name == 'aggregated_streams':
                #     load_staging_raw_aggregated_streams = load_staging_raw
            #
            # load_staging_raw_aggregated_streams << load_staging_raw_tracks


        with TaskGroup(group_id='update_dim_table') as task_group_update_dim_table:
            for dim_table in config.dimension_tables['tables_to_update']:
                task_group_update_dim_table = SQLTemplateOperator(
                    task_id=f'update_{dim_table}',
                    template=f'update_{dim_table}.sql',
                    template_dir=config.QUERIES_DIR,
                    parameters=dict(
                        date='{{ ds }}',
                        licensor=licensor,
                        feedid=config.feedid,
                    )
                )
        task_group_update_dim_table << task_group_load_staging_raw
        task_group_update_dim_table << task_group_load_common_tables

        task_group_load_staging_fact_table = create_load_staging_fact_table(dag=dag)
        task_group_load_staging_fact_table << task_group_update_dim_table
        task_group_load_fact_table = create_fact_tables()
        task_group_load_fact_table << task_group_load_staging_fact_table

        with TaskGroup(group_id='aggregated_skips_and_saves') as task_group_aggregated_skips_and_saves:
            clean_aggregated_skips_and_saves = SQLTemplateOperator(
                task_id='clean',
                template=f'delete_from_table.sql',
                template_dir=config.QUERIES_DIR,
                parameters=dict(
                    table_name='aggregated_skips_and_saves',
                    where=dict(
                        reportdate='{{ ds }}',
                        feedid=config.feedid,
                        vendor_name=licensor,
                    ),
                )
            )

            load_aggregated_skips_and_saves = SQLTemplateOperator(
                task_id='load',
                template='load_aggregated_skips_and_saves.sql',
                template_dir=config.QUERIES_DIR,
                parameters=dict(
                    storeid=config.storeid,
                    feedid=config.feedid,
                    reportdate='{{ ds }}',
                    licensor=licensor,
                )
            )
            load_aggregated_skips_and_saves << clean_aggregated_skips_and_saves

        task_group_aggregated_skips_and_saves << task_group_update_dim_table
        task_group_aggregated_skips_and_saves << task_group_load_staging_raw