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_java_dag_util as sptfy_dwnld_dag_utl
import spotify_etl_bq_bt_util as sptfy_bq_bt_dag_utl
import slack_util
import dates as dt

dag_name = 'spotify_etl_java_v1_0'



# Dag level settings
default_args = {
    'owner': 'airflow',
    'depends_on_past': True,
    'start_date': dt.get_start_date_for_dag(),
    'on_failure_callback':slack_util.dag_failure,
    'dataflow_default_options': {
        'project': Variable.get('dataflow-project'),
        'tempLocation': 'gs://' + Variable.get('dataflow-template-bucket') + '/staging/'
    }
}

# Create the dag

#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
)



# 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='start_vms')
    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
)



start_vms_cmd = 'gcloud compute instances start --zone us-east1-b ' + \
                'spotify-upload-test  spotify-upload-test-1  spotify-upload-test-2 ' + \
                'spotify-upload-test-3'

start_vms_task = BashOperator(task_id='start_vms',
                            bash_command=start_vms_cmd,
                            dag=dag)


stop_vms_cmd = 'gcloud compute instances stop --zone us-east1-b ' + \
                'spotify-upload-test  spotify-upload-test-1  spotify-upload-test-2 ' + \
                'spotify-upload-test-3'

stop_vms_task = BashOperator(task_id='stop_vms',
                            bash_command=stop_vms_cmd,
                            dag=dag)

# fork dependencies for download branching
check_dwnlds_logic >> fork_dwnlds
fork_dwnlds >> start_vms_task
start_vms_task>>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, stop_vms_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)
# this is for creating the tasks for loading the aggregated streams to BQ
#spotify_etl_agg_streams_util.create_agg_strms_to_bq_tasks(date_to_download, dag, wait_task, cleanup_task)