# By: Riasat Ullah
# This file contains all constants and functions related to the DataDog integration.

from dbqueries import db_integrations
from dbqueries.integrations import db_datadog
from utils import constants, integration_type_names as intt, logging, s3, var_names
import json
import requests


# datadog s3 credential file variables
datadog_s3_bucket = 'taskcall-prod-data'
datadog_s3_key = 'credentials/datadog_credentials.json'

# url paths
url_dd_incidents = '{0}/api/v2/incidents/{1}'
url_dd_refresh_token = '{0}/oauth2/v1/token'

# DataDog alert statuses
no_data_status = 'No data'
null_status = 'null'
recovered_status = 'Recovered'
re_no_data_status = 'Re-No data'
re_notify_status = 'Renotify'
re_triggered_status = 'Re-Triggered'
re_warn_status = 'Re-Warn'
triggered_status = 'Triggered'
warn_status = 'Warn'

# DataDog priorities
low_priority = 'low'
normal_priority = 'normal'

# DataDog variables
var_aggreg_key = 'aggreg_key'
var_alert = 'alert'
var_date = 'date'
var_id = 'id'
var_incident_public_id = 'incident_public_id'
var_incident_severity = 'incident_severity'
var_incident_title = 'incident_title'
var_incident_url = 'incident_url'
var_incident_uuid = 'incident_uuid'
var_link = 'link'
var_priority = 'priority'
var_scope = 'scope'
var_snapshot = 'snapshot'
var_tags = 'tags'
var_text_msg = 'text_msg'
var_title = 'title'
var_transition = 'transition'

# possible starting statuses of an alert
starting_statuses = [no_data_status, null_status, triggered_status, warn_status]

# Datadog incident severity map
incident_severity_map = {
    constants.critical_urgency: 'SEV-1',
    constants.high_urgency: 'SEV-2',
    constants.medium_urgency: 'SEV-3',
    constants.low_urgency: 'SEV-4',
    constants.minor_urgency: 'SEV-5'
}

# DataDog incident status map
incident_status_map = {
    constants.open_state: 'active',
    constants.acknowledged_state: 'stable',
    constants.resolved_state: 'resolved'
}


def get_alert_scope_set(scope_str):
    '''
    Get the str $ALERT_SCOPE variable as a set.
    :param scope_str: comma separated string value of $ALERT_SCOPE
    :return: (set) of $ALERT_SCOPE values
    '''
    return set([y.lstrip().rstrip() for y in scope_str.split(',')])


def get_standardized_alert_priority(alert_priority):
    '''
    Standardizes a DataDog alert priority to the its equivalent TaskCall urgency level.
    :param alert_priority: priority designated to the alert by DataDog
    :return: (str) TaskCall urgency level
    '''
    if alert_priority == low_priority:
        return constants.low_urgency
    else:
        return constants.high_urgency


def get_standardized_severity(severity):
    '''
    Get the TaskCall urgency of a Datadog severity.
    :param severity: Datadog severity
    :return: TaskCall urgency
    '''
    std_sev = constants.critical_urgency
    for key in incident_severity_map:
        if incident_severity_map[key] == severity:
            std_sev = key
            break
    return std_sev


def get_datadog_credentials():
    '''
    Get the credentials needed for handling and making API calls to Datadog.
    :return: (dict) of credentials
    '''
    creds = s3.read_json(datadog_s3_bucket, datadog_s3_key)
    return creds


def datadog_patch_request(url, access_token, body):
    '''
    Makes a PATCH request to the api and returns the response status.
    :param url: api url
    :param access_token: access token received from Datadog
    :param body: the JSON body to send with the request
    :return: status code
    '''
    try:
        standard_headers = {
            'Authorization': 'Bearer ' + access_token,
            'Content-Type': 'application/json'
        }
        response = requests.patch(url, headers=standard_headers, data=json.dumps(body))
        if response.status_code not in [200, 201]:
            logging.error('Datadog patch request failed with code: ' + str(response.status_code))
            logging.error(response.json())
        return response.status_code
    except (ConnectionRefusedError, ConnectionError) as e:
        logging.exception('Failed to connect with Datadog')
        logging.exception(str(e))
        raise ConnectionRefusedError
    except Exception as e:
        logging.exception('Datadog API request failed...')
        logging.exception(str(e))
        raise ConnectionRefusedError


def datadog_refresh_token(conn, timestamp, org_id, integ_key):
    '''
    Refresh a Datadog token.
    :param conn: db connection
    :param timestamp: timestamp when this request is being made
    :param org_id: organization ID
    :param integ_key: the integration key of this Datadog integration
    :return: [status code, response]
    '''
    try:
        new_acc_tok = None
        dd_creds = get_datadog_credentials()
        dd_integ_det = db_datadog.get_datadog_account_tokens(conn, timestamp, org_id, integration_key=integ_key)
        dd_site = dd_integ_det[var_names.vendor_endpoint]
        dd_ref_tok = dd_integ_det[var_names.refresh_token]

        body = {
            'grant_type': 'refresh_token',
            'client_id': dd_creds[var_names.client_id],
            'client_secret': dd_creds[var_names.client_secret],
            'refresh_token': dd_ref_tok,
        }
        response = requests.post(url_dd_refresh_token.format(dd_site), data=body)
        status = response.status_code
        output = response.json()
        if status in [200, 201]:
            new_acc_tok = output[var_names.access_token]
            external_info = {
                var_names.vendor_endpoint: dd_site,
                var_names.access_token: new_acc_tok,
                var_names.refresh_token: output[var_names.refresh_token]
            }
            db_integrations.check_and_update_organization_integration_details(
                conn, timestamp, org_id, intt.datadog, dd_integ_det[var_names.vendor_endpoint_name], external_info
            )
        return new_acc_tok
    except (ConnectionRefusedError, ConnectionError) as e:
        logging.exception('Failed to connect with Datadog')
        logging.exception(str(e))
        raise ConnectionRefusedError
    except Exception as e:
        logging.exception('Datadog API request failed...')
        logging.exception(str(e))
        raise ConnectionRefusedError


def update_datadog_status(conn, timestamp, org_id, integ_key, dd_inc_id, instance_state):
    '''
    Update the status of a Datadog incident.
    :param conn: db connection
    :param timestamp: timestamp when this request is being made
    :param org_id: organization ID
    :param integ_key: the integration key of this Datadog integration
    :param dd_inc_id: Datadog incident ID
    :param instance_state: the current state of the instance
    '''
    dd_integ_det = db_datadog.get_datadog_account_tokens(conn, timestamp, org_id, integration_key=integ_key)
    dd_site = dd_integ_det[var_names.vendor_endpoint]
    dd_acc_tok = dd_integ_det[var_names.access_token]
    dd_status = incident_status_map[instance_state]

    body = {
        'data': {
            'type': 'incidents',
            'attributes': {
                'fields': {
                    'state': {'value': dd_status}
                }
            }
        }
    }
    resp_status = datadog_patch_request(url_dd_incidents.format(dd_site, str(dd_inc_id)), dd_acc_tok, body)
    if resp_status not in [200, 201]:
        new_acc_tok = datadog_refresh_token(conn, timestamp, org_id, integ_key)
        if new_acc_tok is not None:
            datadog_patch_request(url_dd_incidents.format(dd_site, str(dd_inc_id)), new_acc_tok, body)


def update_datadog_severity(conn, timestamp, org_id, integ_key, dd_inc_id, urgency):
    '''
    Update the severity of a Datadog incident.
    :param conn: db connection
    :param timestamp: timestamp when this request is being made
    :param org_id: organization ID
    :param integ_key: the integration key of this Datadog integration
    :param dd_inc_id: Datadog incident ID
    :param urgency: TaskCall urgency
    '''
    dd_integ_det = db_datadog.get_datadog_account_tokens(conn, timestamp, org_id, integration_key=integ_key)
    dd_site = dd_integ_det[var_names.vendor_endpoint]
    dd_acc_tok = dd_integ_det[var_names.access_token]
    dd_severity = incident_severity_map[urgency]

    body = {
        'data': {
            'type': 'incidents',
            'attributes': {
                'fields': {
                    'severity': {'value': dd_severity}
                }
            }
        }
    }
    resp_status = datadog_patch_request(url_dd_incidents.format(dd_site, str(dd_inc_id)), dd_acc_tok, body)
    if resp_status not in [200, 201]:
        new_acc_tok = datadog_refresh_token(conn, timestamp, org_id, integ_key)
        if new_acc_tok is not None:
            datadog_patch_request(url_dd_incidents.format(dd_site, str(dd_inc_id)), new_acc_tok, body)