import logging
from datetime import datetime, timedelta

from airflow import DAG
from airflow.operators.python import PythonOperator

from flows.awal_salesforce import config
from flows.awal_salesforce import tasks

logger = logging.getLogger(__file__)

default_args = dict(
    owner="theorcahrd",
    email_on_failure=False,
    email_on_retry=False
)

with DAG(
        dag_id="awal_salesforce",
        default_args=default_args,
        schedule_interval='@daily',
        start_date=datetime(2022, 12, 1),
        catchup=False,
        max_active_runs=1,
        dagrun_timeout=timedelta(minutes=15),
) as dag:
    bootstrap = PythonOperator(
        task_id='bootstrap',
        python_callable=tasks.bootstrap,
    )

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

        fetch = PythonOperator(
            task_id=f'fetch_{suffix}',
            # retries=3,
            # retry_delay=timedelta(minutes=5),

            python_callable=tasks.fetch,
            op_kwargs=dict(
                salesforce_entity=report_config['salesforce_entity'],
                s3_bucket='{{ ti.xcom_pull("bootstrap").archive_s3_bucket }}',
                s3_key='{{ ti.xcom_pull("bootstrap").archive_s3_key }}',
            )
        )
        bootstrap >> fetch

        create_table = PythonOperator(
            task_id=f'create_table_{suffix}',

            python_callable=tasks.create_table,
            op_kwargs=dict(
                s3_bucket='{{ ti.xcom_pull("bootstrap").archive_s3_bucket }}',
                s3_key='{{ ti.xcom_pull("bootstrap").archive_s3_key }}',
                structure_file=f"{{{{ ti.xcom_pull('{fetch.task_id}').structure_file }}}}",
                table_name=report_config['table_name'],
            )
        )
        fetch >> create_table

        copy_data = PythonOperator(
            task_id=f'copy_data_{suffix}',

            python_callable=tasks.copy_data,
            op_kwargs=dict(
                data_file=f"{{{{ ti.xcom_pull('{fetch.task_id}').data_file }}}}",
                table_name=report_config['table_name'],
                s3_path='{{ ti.xcom_pull("bootstrap").archive_s3_path }}',
            )
        )
        create_table >> copy_data