import json
import requests

from airflow.models import Variable
from datetime import timedelta, datetime, date
from airflow import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.python_operator import PythonOperator

from airflow.operators.python_operator import BranchPythonOperator
from airflow.operators.bash_operator import BashOperator


import spotify_etl_dwnld_dag_util as sptfy_dwnld_dag_utl
import spotify_etl_bq_bt_util as sptfy_bq_bt_dag_utl
import dates as dt

# Dag level settings
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': dt.get_start_date_for_dag(),
    'email': [Variable.get('error_notification_email')],
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 1,
    'dataflow_default_options': {
        'project': 'dsp-automation',
        'tempLocation': 'gs://dataflow-dsp-demo/staging/'
    },
    'retry_delay': timedelta(minutes=5)
}

# Create the dag
dag_name = 'spotify_etl_v1_0'
#schedule_interval = '*/5 * * * *'
schedule_interval = None
dag = DAG(dag_name, default_args=default_args, max_active_runs=1, catchup=False, \
          schedule_interval=schedule_interval)

# Wait task is created for letting all the extracted file downloads to finish
wait_task = DummyOperator(
    task_id='Wait',
    dag=dag
)

# Dummy Downloads task for letting all the gzip downloads to finish
downloads_completed_task = DummyOperator(
    task_id='DownloadsCompleted',
    dag=dag
)

# If there are no downloads available for a date then this will get executed
no_dwnlds = DummyOperator(
    task_id='no_dwnlds',
    dag=dag
)

# This dummy dag is used for delegating the control to downloads
delegate_to_dwlds=DummyOperator(
    task_id="delegate_to_dwlds",
    dag=dag
)

date_to_download = Variable.get('spotify_date_to_run')[0:10]
# This is for checking if the downloads are available for a download date
def is_ready_to_proceed_for_dwnlds(**kwargs):
    are_downloads_available = sptfy_dwnld_dag_utl.are_downloads_available('sonybmgmusicentertainment', date_to_download)
    if are_downloads_available:
        kwargs['ti'].xcom_push(key='next_task_id', value='delegate_to_dwlds')
    else:
        kwargs['ti'].xcom_push(key='next_task_id', value='no_dwnlds')

#This is operator that calls the function to check if the downloads are available for a given date
check_dwnlds_logic = PythonOperator(
    task_id='is_ready_to_proceed_for_dwnlds',
    python_callable=is_ready_to_proceed_for_dwnlds,
    provide_context=True,
    dag=dag
)

#This is the branching function whether to proceed for downloads or not
def branch_for_dwnlds(**kwargs):
    return kwargs['ti'].xcom_pull(key='next_task_id')

#This is the branching python operator whether to proceed for downloads or not
fork_dwnlds = BranchPythonOperator(
    task_id='branch_for_dwlds',
    python_callable=branch_for_dwnlds,
    provide_context=True,
    dag=dag
)

# fork dependencies for download branching
check_dwnlds_logic >> fork_dwnlds
fork_dwnlds >> delegate_to_dwlds
fork_dwnlds >> no_dwnlds

# This is for creating the actual download tasks
sptfy_dwnld_dag_utl.create_download_tasks(date_to_download, dag, delegate_to_dwlds, wait_task, downloads_completed_task)
#This is for creating tasks for loading data from files to BQ, aggregation and transferring the aggregation data from BQ to BT 
sptfy_bq_bt_dag_utl.create_bq_bt_tasks(date_to_download, dag, wait_task, False)