from airflow.contrib.operators.bigquery_operator import BigQueryOperator
from airflow.operators.bash_operator import BashOperator
from airflow.contrib.operators.dataflow_operator import DataflowTemplateOperator
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.python_operator import PythonOperator
from airflow.operators.python_operator import BranchPythonOperator
import slack_util
from airflow.models import Variable
import dates as dt
import gcs_storage as gcs_strg

def set_next_run_date(**kwargs):
    Variable.set("apple_date_to_run", dt.get_next_day_date(kwargs['date_to_download']))

def is_ready_to_proceed_for_bq_bt(**kwargs):
    no_of_files = gcs_strg.list_blobs_with_prefix( Variable.get("gcs_bucket"), \
              Variable.get("apple_folder_in_bucket") + '/processed/' + kwargs['date_to_download'])

    print('no_of_files == ', no_of_files)
    apple_threshold_files_size = Variable.get("apple_threshold_files_size")

    if  no_of_files > int(apple_threshold_files_size):
        kwargs['ti'].xcom_push(key='next_task_id', value='apple_aggregate')
    else:
        kwargs['ti'].xcom_push(key='next_task_id', value='notify_not_enough_files')

def branch_for_bq_bt(**kwargs):
    return kwargs['ti'].xcom_pull(key='next_task_id')

def create_bq_bt_tasks(dag, date_to_download, wait_task, cleanup_task, is_manual_download):
    # This is the task that checks whether to proceed loading the data from files to BQ
    check_logic = PythonOperator(
        task_id='is_ready_to_proceed_for_bq_bt',
        python_callable=is_ready_to_proceed_for_bq_bt,
        op_kwargs={'date_to_download': date_to_download},
        provide_context=True,
        dag=dag
    )

    # Branch to the next task.
    fork = BranchPythonOperator(
        task_id='branch_for_bq_bt',
        python_callable=branch_for_bq_bt,
        provide_context=True,
        dag=dag
    )
    '''
    delegate_to_load_dwlds=DummyOperator(
                                task_id="delegate_to_load_dwlds",
                                dag=dag
    )
    '''
    aggregation_sql = '''
    SELECT
      CONCAT(isrc,'',FORMAT_DATE('%F', S.report_date),'',LOWER(S.Storefront_Name ),'~',S.report_licensor) AS newRowkey,
      isrc,
      S.report_date AS date,
      LOWER(S.Storefront_Name ) AS country_code,
      S.report_licensor AS licensor,
      SUM(CASE WHEN S.Source_of_Stream = 0 THEN 1 ELSE 0 END) AS source_0,
      SUM(CASE WHEN S.Source_of_Stream = 1 THEN 1 ELSE 0 END) AS source_1,
      SUM(CASE WHEN S.Source_of_Stream = 2 THEN 1 ELSE 0 END) AS source_2,
      SUM(CASE WHEN S.Source_of_Stream = 3 THEN 1 ELSE 0 END) AS source_3,
      SUM(CASE WHEN S.Source_of_Stream = 4 THEN 1 ELSE 0 END) AS source_4,
      SUM(CASE WHEN S.Source_of_Stream = 5 THEN 1 ELSE 0 END) AS source_5,
      SUM(CASE WHEN S.Source_of_Stream = 6 THEN 1 ELSE 0 END) AS source_6,
      SUM(CASE WHEN S.container_type = '0' THEN 1 ELSE 0 END) AS container_type_0,
      SUM(CASE WHEN S.container_type = '1' THEN 1 ELSE 0 END) AS container_type_1,
      SUM(CASE WHEN S.container_type = '2' THEN 1 ELSE 0 END) AS container_type_2,
      SUM(CASE WHEN S.container_type = '3' THEN 1 ELSE 0 END) AS container_type_3,
      COUNT(DISTINCT S.anonymized_person_id) AS listeners,
      COUNT(*) AS streams
    FROM
      apple.apple_amStreams S
    JOIN
      apple.apple_amContent C
    ON
      (C.report_date = S.report_date)
      AND (C.report_licensor = S.report_licensor)
      AND (C.apple_identifier = S.apple_identifier)
    WHERE
      S.report_date = DATE("''' + date_to_download + '''")
      AND C.report_date = DATE("''' + date_to_download + '''")
      AND C.isrc != ''
    GROUP BY
      C.isrc,
      S.report_date,
      S.storefront_name,
      S.report_licensor
    '''

    aggregation = BigQueryOperator(
        task_id='apple_aggregate',
        sql=aggregation_sql,
        write_disposition='WRITE_APPEND',
        destination_dataset_table='apple.aggregation',
        use_legacy_sql=False,
        dag=dag
    )

    # Change the number of nodes in BT before running Dataflow job.
    bt_before_cmd = 'gcloud beta bigtable clusters update ' + Variable.get('bigtable_cluster_id') + \
                    ' --instance=' + Variable.get('bigtable_instance_id') + ' --num-nodes=' + \
                    Variable.get('bigtable_num_nodes_before_running_dataflow')

    modify_bt_before = BashOperator(task_id='change_to_' + Variable.get('bigtable_num_nodes_before_running_dataflow'),
                                bash_command=bt_before_cmd,
                                dag=dag)

    # Run Dataflow job. Moves aggregate data from BQ to BT.
    bq_to_bt = DataflowTemplateOperator(
        task_id='sony_apple_bq_to_bt_template',
        template='gs://' + Variable.get('dataflow-template-bucket') + '/templates/sony-apple-bq-to-bt-template',
        gcp_conn_id='google_cloud_default',
        parameters={
               'bqQuery': 'select * from `' + Variable.get('apple_agg_dataset') + '` WHERE date = DATE("' + date_to_download + '")'
           },
        on_failure_callback=slack_util.dataflow_fail_for_bq_bt,
        dag=dag)

    # Revert back the number of nodes in BT after Dataflow job.
    bt_after_cmd = 'gcloud beta bigtable clusters update ' + Variable.get('bigtable_cluster_id') + \
                   ' --instance=' + Variable.get('bigtable_instance_id') + ' --num-nodes=' + \
                   Variable.get('bigtable_num_nodes')

    modify_bt_after = BashOperator(task_id='change_to_' + Variable.get('bigtable_num_nodes'),
                            bash_command=bt_after_cmd,
                            dag=dag)


    notify_not_enough_files = slack_util.get_notify_slack_hook_for_not_enough_files(dag, date_to_download);

    wait_task >> check_logic
    check_logic >> fork

    fork >> aggregation
    fork >> notify_not_enough_files


    aggregation >> modify_bt_before
    modify_bt_before >> bq_to_bt
    bq_to_bt >> modify_bt_after

    #modify_bt_after >> stop_vms_task
    modify_bt_after >> cleanup_task

    if not is_manual_download:
        set_next_run_date_task = PythonOperator(
                        task_id = 'set_next_run_date',
                        python_callable=set_next_run_date,
                        op_kwargs={'date_to_download': date_to_download},
                        dag= dag
        )
        cleanup_task >> set_next_run_date_task