from airflow.operators.python_operator import PythonOperator
from airflow.contrib.operators.gcs_to_bq import GoogleCloudStorageToBigQueryOperator
from airflow.contrib.operators.bigquery_operator import BigQueryOperator
from airflow.contrib.operators.dataflow_operator import DataflowTemplateOperator
import dates as dt
from airflow.models import Variable
from airflow.operators.dummy_operator import DummyOperator
from airflow.operators.bash_operator import BashOperator
from airflow.operators.python_operator import BranchPythonOperator
import gcs_storage as gcs_strg
import slack_util
# This function checks no of files available in incoming folder and if it is
# greater than the thresh hold then passes the control to load the data from files
# to BQ
def is_ready_to_proceed_for_bq_bt(**kwargs):
    #no_of_files = gcs_strg.list_blobs_with_prefix(Variable.get("gcs_bucket"), \
                  #Variable.get("spotify_folder_in_bucket") + '/incoming')
    no_of_files_for_sony_bmg = gcs_strg.list_blobs_with_partial_file_name( Variable.get("gcs_bucket"), \
              Variable.get("spotify_folder_in_bucket") + '/processed/' + kwargs['date_to_download'], \
              'spotify_streams_v2_sonybmgmusicentertainment_' + kwargs['date_to_download'])
    print('no_of_files_for_sony_bmg == ', no_of_files_for_sony_bmg)
    no_of_files_for_orchard = gcs_strg.list_blobs_with_partial_file_name( Variable.get("gcs_bucket"), \
              Variable.get("spotify_folder_in_bucket") + '/processed/' + kwargs['date_to_download'], \
              'spotify_streams_v2_theorchard_' + kwargs['date_to_download'])
    print('no_of_files_for_orchard == ', no_of_files_for_orchard)
    no_of_files_for_marketing = gcs_strg.list_blobs_with_partial_file_name( Variable.get("gcs_bucket"), \
              Variable.get("spotify_folder_in_bucket") + '/processed/' + kwargs['date_to_download'], \
              'spotify_streams_v2_sonymusicmarketing_' + kwargs['date_to_download'])
    print('no_of_files_for_marketing == ', no_of_files_for_marketing)
    no_of_files_for_marketing_intl = gcs_strg.list_blobs_with_partial_file_name( Variable.get("gcs_bucket"), \
              Variable.get("spotify_folder_in_bucket") + '/processed/' + kwargs['date_to_download'], \
              'spotify_streams_v2_sonymusicmarketingintlcontent_' + kwargs['date_to_download'])
    print('no_of_files_for_marketing_intl == ', no_of_files_for_marketing_intl)
    threshold_for_sony_bmg = Variable.get("spotify_sonybmg_threshold_files_size")
    threshold_for_orchard = Variable.get("spotify_orchard_threshold_files_size")
    threshold_for_marketing = Variable.get("spotify_sonymarketing_threshold_files_size")
    threshold_for_marketing_intl = Variable.get("spotify_sonymarketing_intl_threshold_files_size")

    if  no_of_files_for_sony_bmg > int(threshold_for_sony_bmg) and \
        no_of_files_for_orchard > int(threshold_for_orchard) and \
        no_of_files_for_marketing > int(threshold_for_marketing) and \
        no_of_files_for_marketing_intl > int(threshold_for_marketing_intl):
        kwargs['ti'].xcom_push(key='next_task_id', value='aggregate')
    else:
        kwargs['ti'].xcom_push(key='next_task_id', value='notify_not_enough_files')

# This is the branch that determines whether to proceed to load the data from
# files to BQ
def branch_for_bq_bt(**kwargs):
    return kwargs['ti'].xcom_pull(key='next_task_id')

# Next run date task is created to set the next download date and this is called
# when all the files including archive files are downloaded
def set_next_run_date(**kwargs):
    Variable.set("spotify_date_to_run", dt.get_next_day_date(kwargs['date_to_download']))

def get_cleanup_task(dag, download_date):
    # Remove all the files in the unprocessed folder.

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

# This is the main function for creating the tasks like loading to BQ, aggregation and
# transferring data from BQ to BT
def create_bq_bt_tasks(date_to_download, dag, wait_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
    )

    # Notify. Need to implement.
    '''
    case_false = DummyOperator(task_id='case_false',
                            dag=dag)
    '''
    notify_not_enough_files = slack_util.get_notify_slack_hook_for_not_enough_files(dag, date_to_download);


    new_sql_query = '''
    SELECT
      CONCAT(T.isrc,'~',FORMAT_DATE('%F', S.report_date),'~',LOWER(U.country),'~',S.report_licensor) AS newRowkey,
      T.isrc,
      S.report_date AS date,
      LOWER(U.country) AS country_code,
      S.report_licensor AS licensor,
      SUM(CASE WHEN S.source = 'album' THEN 1 ELSE 0 END) AS source_album,
      SUM(CASE WHEN S.source = 'artist' THEN 1 ELSE 0 END) AS source_artist,
      SUM(CASE WHEN S.source = 'chart' THEN 1 ELSE 0 END) AS source_chart,
      SUM(CASE WHEN S.source = 'collection' THEN 1 ELSE 0 END) AS source_collection,
      SUM(CASE WHEN S.source = 'radio' AND s.source_uri = 'dailyMix' THEN 1 ELSE 0 END) AS source_daily_mix,
      SUM(CASE WHEN S.source = 'others_playlist' AND s.source_uri = 'discoverWeekly' THEN 1 ELSE 0 END) AS source_discover_weekly,
      SUM(CASE WHEN S.source = 'other' THEN 1 ELSE 0 END) AS source_other,
      SUM(CASE WHEN S.source = 'others_playlist' AND s.source_uri NOT IN ('discoverWeekly',  'releaseRadar') THEN 1 ELSE 0 END) AS source_others_playlist,
      SUM(CASE WHEN S.source = 'play_queue' THEN 1 ELSE 0 END) AS source_play_queue,
      SUM(CASE WHEN S.source = 'radio' AND s.source_uri != 'dailyMix' THEN 1 ELSE 0 END) AS source_radio,
      SUM(CASE WHEN S.source = 'others_playlist' AND s.source_uri = 'releaseRadar' THEN 1 ELSE 0 END) AS source_release_radar,
      SUM(CASE WHEN S.source = 'search' THEN 1 ELSE 0 END) AS source_search,
      SUM(CASE WHEN S.source NOT IN ('album',  'artist',  'chart',  'collection',  'other',  'others_playlist', 'play_queue',  'radio', 'search') THEN 1 ELSE 0 END) AS source_unknown,
      SUM(CASE WHEN S.source = 'album' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_album,
      SUM(CASE WHEN S.source = 'artist' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_artist,
      SUM(CASE WHEN S.source = 'chart' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_chart,
      SUM(CASE WHEN S.source = 'collection' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_collection,
      SUM(CASE WHEN S.source = 'radio' AND s.source_uri = 'dailyMix' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_daily_mix,
      SUM(CASE WHEN S.source = 'others_playlist' AND s.source_uri = 'discoverWeekly' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_discover_weekly,
      SUM(CASE WHEN S.source = 'other' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_other,
      SUM(CASE WHEN S.source = 'others_playlist' AND U.access = 'free' AND s.source_uri NOT IN ('discoverWeekly',  'releaseRadar') THEN 1 ELSE 0 END) AS free_source_others_playlist,
      SUM(CASE WHEN S.source = 'play_queue' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_play_queue,
      SUM(CASE WHEN S.source = 'radio' AND s.source_uri != 'dailyMix' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_radio,
      SUM(CASE WHEN S.source = 'others_playlist' AND s.source_uri = 'releaseRadar' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_release_radar,
      SUM(CASE WHEN S.source = 'search' AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_search,
      SUM(CASE WHEN S.source NOT IN ('album',  'artist',  'chart',  'collection',  'other',  'others_playlist',  'play_queue',  'radio',  'search') AND U.access = 'free' THEN 1 ELSE 0 END) AS free_source_unknown,
      SUM(CASE WHEN S.source = 'album' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_album,
      SUM(CASE WHEN S.source = 'artist' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_artist,
      SUM(CASE WHEN S.source = 'chart' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_chart,
      SUM(CASE WHEN S.source = 'collection' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_collection,
      SUM(CASE WHEN S.source = 'radio' AND s.source_uri = 'dailyMix' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_daily_mix,
      SUM(CASE WHEN S.source = 'others_playlist' AND s.source_uri = 'discoverWeekly' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_discover_weekly,
      SUM(CASE WHEN S.source = 'other' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_other,
      SUM(CASE WHEN S.source = 'others_playlist' AND U.access != 'free' AND s.source_uri NOT IN ('discoverWeekly',  'releaseRadar') THEN 1 ELSE 0 END) AS paid_source_others_playlist,
      SUM(CASE WHEN S.source = 'play_queue' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_play_queue,
      SUM(CASE WHEN S.source = 'radio' AND s.source_uri != 'dailyMix' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_radio,
      SUM(CASE WHEN S.source = 'others_playlist' AND s.source_uri = 'releaseRadar' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_release_radar,
      SUM(CASE WHEN S.source = 'search' AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_search,
      SUM(CASE WHEN S.source NOT IN ('album',  'artist',  'chart',  'collection',  'other',  'others_playlist',  'play_queue',  'radio',  'search') AND U.access != 'free' THEN 1 ELSE 0 END) AS paid_source_unknown,
      SUM(CASE WHEN S.repeat_play = TRUE THEN 1 ELSE 0 END) AS repeat_play,
      SUM(CASE WHEN S.shuffle = TRUE THEN 1 ELSE 0 END) AS shuffle,
      COUNT(DISTINCT (CASE WHEN U.access = 'free' THEN S.user_id END)) AS free_listeners,
      COUNT(DISTINCT (CASE WHEN U.access != 'free' THEN S.user_id END)) AS paid_listeners,
      COUNT(DISTINCT S.user_id) AS listeners,
      COUNT(CASE WHEN U.access = 'free' THEN 1 ELSE 0 END) AS free_streams,
      COUNT(CASE WHEN U.access != 'free' THEN 1 ELSE 0 END) AS paid_streams,
      COUNT(*) AS streams
    FROM
      spotify.spotify_streams S
    JOIN
      `spotify.spotify_users`  U
    ON
      (U.report_date = S.report_date)
      AND (U.report_licensor = S.report_licensor)
      AND (U.user_id = S.user_id)
    JOIN
      `spotify.spotify_tracks`  T
    ON
      (T.report_date = S.report_date)
      AND (T.report_licensor  = S.report_licensor)
      AND (T.track_id = S.track_id)
    WHERE
      S.report_date = DATE("''' + date_to_download + '''")
      AND T.report_date = DATE("''' + date_to_download + '''")
      AND U.report_date = DATE("''' + date_to_download + '''")
      AND T.isrc != ''
    GROUP BY
      T.isrc,
      S.report_date,
      U.country,
      S.report_licensor
    '''

    aggregation = BigQueryOperator(
        task_id='aggregate',
        sql=new_sql_query,
        write_disposition='WRITE_APPEND',
        destination_dataset_table='spotify.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_spotify_bq_to_bt_template',
        template='gs://' + Variable.get('dataflow-template-bucket') + '/templates/sony-spotify-bq-to-bt-template',
        gcp_conn_id='google_cloud_default',
        parameters={
               'bqQuery': 'select * from `' + Variable.get('spotify_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)

    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

    cleanup_task = get_cleanup_task(dag, date_to_download)
    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