import logging
from datetime import datetime, timedelta

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.sensors.time_delta import TimeDeltaSensor
from airflow.utils.task_group import TaskGroup

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

logger = logging.getLogger(__file__)

default_args = dict(
    owner=config.PROVIDER,
)

for licensor in config.spotify_api_licensors:

    for report_name, report_config in config.reports.items():

        report_dag_args = {}
        common_dag_args = dict(
            tags=[config.FLOW_NAME],
            default_args=default_args,
            schedule_interval='@daily',
            start_date=datetime(2022, 12, 24),
            catchup=True,
            dagrun_timeout=config.DATA_THRESHOLD,
        )
        dag_args = {**common_dag_args, **report_dag_args}

        dag_id = '_'.join([config.FLOW_NAME, licensor, 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,
                    licensor=licensor,
                )
            )
            bootstrap << wait_for_data

            create_temp_table = SQLTemplateOperator(
                task_id='create_temp_table',
                template='create_table.sql',
                template_dir=config.QUERIES_DIR,
                parameters=dict(
                    table_name='{{ ti.xcom_pull("bootstrap").temp_table }}',
                    kind='TRANSIENT',
                    or_replace=True,
                    columns=[
                        dict(
                            name='content',
                            type='variant',
                        )
                    ],
                )
            )

            fetch_sensor_poke_interval = timedelta(hours=3).total_seconds()
            fetch_sensor_timeout = config.DATA_THRESHOLD.total_seconds()

            if report_config.get('use_countries'):
                countries = config.expected_countries
                splits = tasks.split_by_prefixes(prefixes=['A', 'E', 'I', 'O', ], data=config.expected_countries)
                for prefix, countries in splits.items():
                    with TaskGroup(group_id=f'expected_{prefix}') as tg1:
                        for number, country in enumerate(countries):
                            fetch = tasks.SpotifyAPISensor(
                                task_id=f'{country}',
                                report_name=report_name,
                                licensor=licensor,
                                country=country,
                                archive_s3_path='{{ ti.xcom_pull("bootstrap").archive_s3_path }}',
                                aws_conn_id=config.AWS_CONN_ID,
                                poke_interval=fetch_sensor_poke_interval,
                                timeout=fetch_sensor_timeout,
                                dev_mode=config.DEV_MODE,
                            )
                            fetch << bootstrap
                            fetch >> create_temp_table

                expected_countries_set = set(config.expected_countries)
                optional_countries = [c for c in config.countries if c not in expected_countries_set]

                countries = set(config.expected_countries)
                splits = tasks.split_by_prefixes(prefixes=['A', 'E', 'I', 'O', ], data=config.expected_countries)
                for prefix, countries in splits.items():
                    with TaskGroup(group_id=f'optional_{prefix}') as tg1:
                        for number, country in enumerate(countries):
                            fetch = tasks.SpotifyAPISensor(
                                task_id=f'{country}',
                                report_name=report_name,
                                licensor=licensor,
                                country=country,
                                archive_s3_path='{{ ti.xcom_pull("bootstrap").archive_s3_path }}',
                                aws_conn_id=config.AWS_CONN_ID,
                                poke_interval=fetch_sensor_poke_interval,
                                timeout=fetch_sensor_timeout,
                                dev_mode=config.DEV_MODE,
                            )
                            fetch << bootstrap
                            fetch >> create_temp_table
            else:
                fetch = tasks.SpotifyAPISensor(
                    task_id=f'fetch',
                    report_name=report_name,
                    licensor=licensor,
                    country=None,
                    archive_s3_path='{{ ti.xcom_pull("bootstrap").archive_s3_path }}',
                    aws_conn_id=config.AWS_CONN_ID,
                    poke_interval=fetch_sensor_poke_interval,
                    timeout=fetch_sensor_timeout,
                    dev_mode=config.DEV_MODE,
                )
                fetch << bootstrap
                fetch >> create_temp_table

            load_temp_table = SQLTemplateOperator(
                task_id='load_temp_table',
                template='copy_into.sql',
                template_dir=config.QUERIES_DIR,
                parameters=dict(
                    table_name='{{ ti.xcom_pull("bootstrap").temp_table }}',
                    from_s3='{{ ti.xcom_pull("bootstrap").archive_s3_path }}',
                    file_format=config.FILE_FORMAT,
                    on_error=f'SKIP_FILE_{config.snowflake_error_limit}',
                    # pattern not required because all files in the path applies
                    pattern=f'.*/{report_name}.*.gz'
                ),
                aws_conn_id=config.AWS_CONN_ID,
            )
            load_temp_table << create_temp_table