from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime

from lib import dynamodb, s3, snowflake

dag = DAG(
    'collaborator_report',
    description='Generate Collaborator Report',
    start_date=datetime(2019, 1, 1),
    schedule_interval=None
)


def bootstrap(dag_run, task_instance, **kwargs):
    config = dag_run.conf
    period_ids = config['period_ids']
    account_type = config['account_type']
    account_id = config['account_id']
    collaborator_params = config['collaborator_params']

    user_id_type = account_type[0].upper() + str(account_id)

    report = dynamodb.get_report(user_id_type, collaborator_params, period_ids)

    task_instance.xcom_push(key='report', value=report)


def get_unload_query(dag_run, task_instance, **kwargs):
    config = dag_run.conf
    period_ids = config['period_ids']
    account_type = config['account_type']
    account_id = config['account_id']

    report = task_instance.xcom_pull(key='report', task_ids='bootstrap')
    collaborator_id = report['exclude_transaction_types']
    exclude_transaction_types = report['exclude_transaction_types']

    sql = snowflake.get_sql(
        period_ids, account_type, account_id, collaborator_id,
        exclude_transaction_types)

    task_instance.xcom_push(key='sql', value=sql)


def unload_data_to_s3(dag_run, task_instance, **kwargs):
    config = dag_run.conf
    period_ids = config['period_ids']
    collaborator_params = config['collaborator_params']

    unload_query = task_instance.xcom_pull(
        key='sql', task_ids='get_unload_query')

    s3_url = f's3://collaborator_report/{period_ids}/{collaborator_params}'
    snowflake.unload_to_s3(s3_url, unload_query)

    task_instance.xcom_push(key='file_path', value=s3_url)


def generate_report_file(dag_run, task_instance, **kwargs):
    config = dag_run.conf
    period_ids = config['period_ids']

    source_file_path = task_instance.xcom_pull(
        key='file_path', task_ids='unload_data_to_s3')

    report = task_instance.xcom_pull(key='report', task_ids='bootstrap')
    filename = report['filename']
    file_format = report['file_format']

    s3_path = f's3://collaborator_report/{period_ids}/{filename}.{file_format}'

    s3.upload_file(source_file_path, s3_path)

    task_instance.xcom_push(key='file_path', value=s3_path)


def set_report_status_generated(dag_run, task_instance, **kwargs):
    config = dag_run.conf
    period_ids = config['period_ids']
    account_type = config['account_type']
    account_id = config['account_id']
    collaborator_params = config['collaborator_params']

    user_id_type = account_type[0].upper() + str(account_id)

    s3_path = task_instance.xcom_pull(key='file_path', task_ids='generate_report_file')

    dynamodb.update_report(user_id_type, collaborator_params, s3_path)


bootstrap = PythonOperator(
    task_id='bootstrap',
    provide_context=True,
    python_callable=bootstrap,
    dag=dag,
)

get_unload_query = PythonOperator(
    task_id='get_unload_query',
    provide_context=True,
    python_callable=get_unload_query,
    dag=dag,
)

unload_data_to_s3 = PythonOperator(
    task_id='unload_data_to_s3',
    provide_context=True,
    python_callable=unload_data_to_s3,
    dag=dag,
)

generate_report_file = PythonOperator(
    task_id='generate_report_file',
    provide_context=True,
    python_callable=generate_report_file,
    dag=dag,
)

set_report_status_generated = PythonOperator(
    task_id='set_report_status_generated',
    provide_context=True,
    python_callable=set_report_status_generated,
    dag=dag,
)


(bootstrap
    >> get_unload_query
    >> unload_data_to_s3
    >> generate_report_file
    >> set_report_status_generated)