#! /usr/bin/env python3
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from contextlib import closing
from logging import info
from sql import sql_queries
import snowflake.connector
import psycopg2
import smtplib
import re
import argparse
import sys
import json
import datetime
import boto3

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"]


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)
    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] RTI Push Notifications (daily)"
    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:
    ctx = snowflake.connector.connect(
        user=sf_user,
        password=sf_password,
        account=sf_account,
        warehouse=sf_warehouse
    )
    cs = ctx.cursor(snowflake.connector.DictCursor)
    try:
        cs.execute(sql_query)
        result_dict = cs.fetchall()
    finally:
        cs.close()
    ctx.close()
    return result_dict

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 check_1_to_6() -> dict:
    result_dict = {}
    full_log = {}
## Check_1 Snowflake transform status - system history


## Check_4 Data consistency - data input check
    tmp_sf_check_4 = snowflake_fetch_data(sql_queries.sf_data_input_check_4)
    sf_check_4_result = tmp_sf_check_4[0]['count']
    print(f"CHECK_4: - {sf_check_4_result}")
    full_log["CHECK_4"] = {'CHECK_4_RESULT': sf_check_4_result}

    if sf_check_4_result == 0:
        result_dict["CHECK_4"] = {'CHECK_4_RESULT': 'Failed: Query returned ' + str(sf_check_4_result)}
        full_log["CHECK_4"] = result_dict["CHECK_4"]
    else:
        full_log["CHECK_4"] = {'CHECK_4_RESULT': 'OK: Query returned ' + str(sf_check_4_result) + ' count'}

## Check_5 Data consistency - mapping check
    tmp_sf_check_5 = snowflake_fetch_data(sql_queries.sf_mapping_check_5)
    if tmp_sf_check_5:
        print(f"CHECK_5: - {sf_check_5_result}")
        full_log["CHECK_5"] = {'CHECK_5_RESULT': 'Failed: Query returned not empty result'}
        result_dict["CHECK_5"] = full_log["CHECK_5"]
    else:
        print(f"CHECK_5: OK: Query returned empty result")
        full_log["CHECK_5"] = {'CHECK_5_RESULT': 'OK: Query returned empty result'}

## Check_6 Step-function status
    
    # 1 Step-function status
    # Step-function failed less than 1 times since 05:00 am UTC yesterday to 05:00 am UTC today
    # AWS credentials are stored in /root/.aws/config for the root user

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

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

    # find the date of yesterday
    today = datetime.datetime.utcnow().date()
    yesterday = today - datetime.timedelta(days=1)
    
    prev_day = yesterday.strftime("%Y-%m-%d") + " 05: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:
        result_dict["CHECK_6"] = {'CHECK_6_RESULT': 'Failed: found " + str(len(sfn_total_list)) + " failed step-functions'}
        full_log["CHECK_6"] = result_dict["CHECK_6"]
        print("Check 6: - Found " + str(len(sfn_total_list)) + " failed step-functions")
    else:
        full_log["CHECK_6"] = {'CHECK_6_RESULT': 'No failed step-functions found'}
        print("Check 6: - No failed step-functions found")

    return result_dict, full_log


################# run the all steps ######################
def run_steps() -> dict:
    message_dict = {'name':"[PROD] RTI Push Notifications (daily)",'status':"",'date':"", 'info':"", 'details':"", 'full_log':""}
    # set date according to sql searching date, currently it's `today`
    message_dict["date"] = datetime.datetime.now().strftime("%Y-%m-%d")

    result_dict, full_log = check_1_to_6()
        
    message_dict['full_log'] = {**full_log}
    message_dict['details'] = {**result_dict}

    if result_dict:
        postgres_update_failure_count(sql_queries.failure_count_increase_daily)
        message_dict['status'] = "FAILED"
        message_dict['info'] = "Some steps failed"

    else:
        postgres_update_failure_count(sql_queries.failure_count_to_zero_daily)
        message_dict['status'] = "OK"
        message_dict['info'] = "All steps executed without errors"

    return message_dict


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/rti_push_notifications_daily_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(json.dumps(result, indent=4))
    else:
        parser.print_help()