from airflow.hooks.base_hook import BaseHook
from airflow.contrib.operators.slack_webhook_operator import SlackWebhookOperator
from airflow.operators.bash_operator import BashOperator

def dag_failure(context):
    SLACK_CONN_ID = 'slack'
    slack_webhook_token = BaseHook.get_connection(SLACK_CONN_ID).password
    message = '\nDAG: ' + str(context['dag']) + '\nTask: ' + str(context['task'])
              #'\nDescription: ' + str(context['exception'])
    failed_alert = SlackWebhookOperator(
            task_id='generic_error_notification',
            http_conn_id='slack',
            webhook_token=slack_webhook_token,
            message="There is a error in running the dag == " + message,
            username='airflow',
            dag=context['dag'])

    return failed_alert.execute(context=context)

def get_notify_slack_hook_for_not_enough_files(dag, date_to_download):
    SLACK_CONN_ID = 'slack'
    slack_webhook_token = BaseHook.get_connection(SLACK_CONN_ID).password

    notify_not_enough_files = SlackWebhookOperator(
            task_id='notify_not_enough_files',
            http_conn_id='slack',
            webhook_token=slack_webhook_token,
            message="There are not enough files to process for spotify and the download date is == " + date_to_download,
            username='airflow',
            dag=dag)

    return notify_not_enough_files

def dataflow_fail_for_bq_bt(context):
    message = str(context['dag']) + '\n' + str(context['task']) + '\nDescription: ' + str(context['exception'])
    change_back = "gcloud beta bigtable clusters update bigtable-dsp-automation-dem-c1 --instance=bigtable-dsp-automation-demo --num-nodes=3"
    change = BashOperator(task_id='changeback_to_3',
                            bash_command=change_back,
                            dag=context['dag'])
    SLACK_CONN_ID = 'slack'
    slack_webhook_token = BaseHook.get_connection(SLACK_CONN_ID).password
    failed_alert = SlackWebhookOperator(
            task_id='notify_bq_bt_fail',
            http_conn_id='slack',
            webhook_token=slack_webhook_token,
            message=message,
            username='airflow',
            dag=context['dag'])

    change.execute(context=context)
    return failed_alert.execute(context=context)