from airflow.models import Variable
from datetime import timedelta, datetime, date
from airflow import DAG
import dates as dt
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import PythonOperator
from airflow.operators.python_operator import BranchPythonOperator
import apple_etl_dwld_dag_util as apple_etl_dwld_dag_util
import apple_etl_bq_bt_util
import slack_util


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

remove_cmd = 'gsutil -m rm -r gs://' + Variable.get("gcs_bucket") + '/' + \
             Variable.get("apple_folder_in_bucket") + '/incoming/* | echo'
cleanup_task = BashOperator(task_id='cleanup',
                        bash_command=remove_cmd,
                        dag=dag)

date_to_download = Variable.get('apple_date_to_run')[0:10]

start_vms_cmd = 'gcloud compute instances start --zone us-east1-b ' + \
                'apple-vm-1'

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 ' + \
                'apple-vm-1'

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

# This is for checking if the downloads are available for a download date
def is_ready_to_proceed_for_dwnlds(**kwargs):
    are_downloads_available = apple_etl_dwld_dag_util.are_downloads_available(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
)

# 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
apple_etl_dwld_dag_util.create_download_tasks(date_to_download, dag, delegate_to_dwlds, wait_task, downloads_completed_task)
apple_etl_bq_bt_util.create_bq_bt_tasks(dag, date_to_download, wait_task, stop_vms_task, cleanup_task, True)