#! /usr/bin/env python3

from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from contextlib import closing
from logging import info
import snowflake.connector
import psycopg2
import smtplib
import re
import argparse
import sys
import json
import datetime
import boto3

# connection handles
sf_conn = None
pg_conn = None

with open('config.json', 'r') as f:
    config = json.load(f)

sf_user = config['snowflake']["user"]
sf_password = config['snowflake']["password"]
sf_account = config['snowflake']["account"]
sf_warehouse = config['snowflake']['warehouse']

local_db_name = config['failure_counter']["database"]
local_db_user = config['failure_counter']["user"]
local_db_password = config['failure_counter']["password"]
local_db_host = config['failure_counter']["host"]
local_db_port = config['failure_counter']["port"]

failure_count_increase = "UPDATE failed_count SET failed_count = failed_count + 1, last = 'FAILED', last_updated = NOW() WHERE count_id = 'aegdprw_friday';"
failure_count_to_zero = "UPDATE failed_count SET failed_count = 0, last = 'OK', last_updated = NOW() WHERE count_id = 'aegdprw_friday';"

def postgres_update_failure_count(sql_query):
    with closing(psycopg2.connect(database=local_db_name, user=local_db_user, password=local_db_password, host=local_db_host, port=local_db_port)) as connection:
        with connection.cursor() as cursor:
            cursor.execute(sql_query)
            connection.commit()

def create_arg_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--mail',
                        nargs='+',
                        action='store',
                        dest='recipients',
                        type=str)
    parser.add_argument('-p', '--print', action="store_const", const=True)
    parser.add_argument('-afc', '--acceptable_fails_count', nargs=1, required=True, type=int)
    return parser


def send_mail(recipients: str, message: str, status: str) -> None:
    # Create a multipart message
    msg = MIMEMultipart()
    body_part = MIMEText(message, 'plain')
    msg['Subject'] = f"[{status}][PROD] Appreciation Engine GDPR request to forget ids fetching"
    msg['From'] = "support.sme@dataart.com"
    msg['To'] = recipients
    # Add body to email
    msg.attach(body_part)
    # Create SMTP object
    server = smtplib.SMTP('relay1.dataart.com', 25)
    server.sendmail(msg['From'], msg['To'].split(','), msg.as_string())
    server.quit()


def snowflake_fetch_data(sql_query) -> dict:
    # fetching data from snowflake
    global sf_conn
    if not sf_conn:
        try:
            sf_conn = snowflake.connector.connect(
                user=sf_user,
                password=sf_password,
                account=sf_account,
                warehouse=sf_warehouse,
                network_timeout=2400
            )
        except:
            return False

    cs = sf_conn.cursor(snowflake.connector.DictCursor)
    try:
        cs.execute(sql_query)
        result_dict = cs.fetchall()
    finally:
        cs.close()

    return result_dict

def check_1(acceptable_fails_count: int):

    # 1 Step-function status
    # Step-function failed less than 1 times since 02:00 pm UTC last Friday to 02:00 pm UTC this Friday (last 1w)
    # AWS credentials are stored in /root/.aws/config for the root user

    client = boto3.client('stepfunctions')
    data = client.list_executions(
            stateMachineArn='arn:aws:states:us-east-1:323555055331:stateMachine:prod-delphi-slzGdprAeEmailFlow',
        statusFilter='FAILED'
    )

    sfn_total_list = []
    sfn_prefix_names_list = []
    re_pattern_sfn_prefix = r'\w+-\d+-\w+-\w+-\w\d'

    # find the date of last Friday
    prev_day = datetime.date.fromordinal(datetime.date.today().toordinal()-1).strftime("%F")
    check_date = datetime.datetime.strptime(prev_day, '%Y-%m-%d')

    while (check_date.weekday() != 4):
        prev_day = datetime.date.fromordinal(check_date.toordinal()-1).strftime("%Y-%m-%d") + " 12:00:00"
        check_date = datetime.datetime.strptime(prev_day, '%Y-%m-%d %H:%M:%S')

    # find FAILED step-functions since last friday
    for sfn in data["executions"]:
        sfn.pop('executionArn')
        sfn.pop('stateMachineArn')
        sfn['startDate'] = str(sfn['startDate'])
        sfn['stopDate'] = str(sfn['stopDate'])
        # add {sfn} to the list if the date is appropriate
        if sfn['startDate'] >= str(check_date):
            sfn_total_list.append(sfn)
    # prepare step functions prefix list for further actions
    for sfn in sfn_total_list:
        sfn_prefix_name = re.findall(re_pattern_sfn_prefix, sfn['name'])[0]
        sfn_prefix_names_list.append(sfn_prefix_name)
    # prepare result structure
    sfn_failed_list = [{'sfn_name': sfn,
                        'fails_count': 0,
                        'details': []}
                       for sfn in set(sfn_prefix_names_list)]
    # fill the list with FAILED Step Functions
    for sfn in sfn_failed_list:
        sfn['details'] = [x for x in sfn_total_list
                          if re.findall(re_pattern_sfn_prefix, x['name'])[0] == sfn['sfn_name']
                          ]
        sfn['fails_count'] = len(sfn['details'])
    # prepare eventual message with Step Functions that have been FAILED more then {acceptable_fails_count} times
    message_data = [x for x in sfn_failed_list if x['fails_count'] > acceptable_fails_count]
    if message_data:
        print("Check 1: found " + str(len(sfn_total_list)) + " failed step-functions")
        return message_data
    else:
        print("Check 1: no failed step-functions found")
        return False

def check_2():

    # 2 New ‘request to forget’ file’s arrival into Snowflake
    # expected result: count > 0
    # "гипотетически возможно что rows == 0, если за неделю ни один юзер не попросил его удалить"

    sql_query = "select count(*) from DELPHI_EXPLORATION.RAW.AE_MEMBERS_GDPR_UNLISTED where REPORT_DATE = '" + str(datetime.date.today()) + "';"
    result_raw = snowflake_fetch_data(sql_query)
    if len(result_raw) > 0:
        # we have got some rows from snowflake
        result_count = int(result_raw[0]['COUNT(*)'])
        print("Check 2: count = " + str(result_count))
        return result_count
    else:
        return False

def check_3() -> dict:

    # 3 Check scope of data to be wiped out as a result of weekly stored procedure execution
    # expected result: none, just save the obtained list of IDs for future use

    #sql_query = 'select distinct ml.ID from "DELPHI_EXPLORATION"."RAW"."AE_MEMBERS_LOGINS" ml Limit 100'
    sql_query = 'select distinct ml.ID from "DELPHI_EXPLORATION"."RAW"."AE_MEMBERS_LOGINS" ml join "DELPHI_EXPLORATION"."RAW"."AE_MEMBERS_GDPR_UNLISTED" mgu on ml.EMAIL = mgu.USER_EMAIL union select distinct mp.MEMBER_ID from "DELPHI_EXPLORATION"."MAIN"."AE_MEMBER_PERSONAL" mp join "DELPHI_EXPLORATION"."RAW"."AE_MEMBERS_GDPR_UNLISTED" mgu on mp.EMAIL = mgu.USER_EMAIL union select distinct afp.MEMBER_ID from "DELPHI_EXPLORATION"."MAIN"."AE_ACTIVITIES_FEED_PERSONAL" afp join "DELPHI_EXPLORATION"."RAW"."AE_MEMBERS_GDPR_UNLISTED" mgu on afp.MEMBER_EMAIL = mgu.USER_EMAIL'
    result = snowflake_fetch_data(sql_query)
    print("Check 3: fetched " + str(len(result)) + " rows from snowflake")
    if result:
        #print(json.dumps(result))
        result_list = []
        for entry in result:
            result_list.append(entry["ID"])
        return result_list
    else:
        return False


################# run steps ######################

def run_steps() -> dict:
    message_dict = {'name':"[PROD] Appreciation Engine GDPR (weekly friday)",'status':"",'date':"", 'info':""}
    # set date according to sql searching date
    message_dict["date"] = (datetime.datetime.now() - datetime.timedelta(days=0)).strftime("%Y-%m-%d")

    full_log = {}
    failed_1 = False
    result_1 = check_1(0)
    if (result_1):
        failed_1 = True
        full_log["Step 1 failed SF"] = result_1

    failed_2 = False
    result_2 = check_2()
    full_log["Step 2 count"] = result_2
    if result_2 == 0:
        failed_2 = True

    failed_3 = False
    result_3 = check_3()
    full_log["Step 3 fetched IDs"] = len(result_3)
    if result_3:
        # save result into the json file for future use
        #with open('results/' + datetime.date.today().strftime("%F") +'_step_3.json', 'w') as f:
        #    json.dump(result_3, f)
        list_3 = ""
        i = 0
        for id in result_3:
            if i == 0:
                list_3 = str(id)
            else:
                list_3 += ", " + str(id)
            i += 1
        if len(list_3) > 0:
            result_file = open("results/" + datetime.date.today().strftime("%F") + "_step_3.txt", "w")
            result_file.write(list_3)
            result_file.close()
    else:
        failed_3 = True

    print("\nFull log:")
    print(json.dumps(full_log, indent=2))

    if not failed_1 and not failed_2 and not failed_3:
        # All steps are successfull
        postgres_update_failure_count(failure_count_to_zero)
        message_dict['status'] = "OK"
        message_dict['info'] = "GDPR monitoring: step has finished successfully"
        message_dict['full_log'] = full_log
    else:
        # some step has failed
        postgres_update_failure_count(failure_count_increase)
        message_dict['status'] = "FAILED"
        message_dict['info'] = "GDPR monitoring: some steps have failed"
        if failed_1:
            message_dict['step_1_details'] = "Found " + str(len(result_1)) + " failed step-functions"
        if failed_2:
            message_dict['step_2_details'] = "Count > 0"
        if failed_3:
            message_dict['step_3_details'] = "No data received from snowflake"
        message_dict['full_log'] = full_log

    return message_dict

# Close all connections on exit
def exit():
    global sf_conn
    global pg_conn

    if sf_conn:
        sf_conn.close()

    if pg_conn:
        pg_conn.close()


if __name__ == '__main__':
    # # check command line parameters and print or email result
    parser = create_arg_parser()
    arg_list = parser.parse_args()
    if len(sys.argv[1:]) > 0:
        result = run_steps()
        with open('/tmp/GDPR_weekly_friday_log.json', 'w') as json_file:
            json.dump(result, json_file, indent=4)

        if arg_list.recipients:
            recipients = arg_list.recipients[0]
            send_mail(recipients=recipients, message=json.dumps(result, indent=4), status=result['status'])

        if arg_list.print:
            print("\nResults:")
            print(json.dumps(result, indent=4))
    else:
        parser.print_help()

exit()