#! /usr/bin/env python3

from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
import smtplib
import boto3
import json
from datetime import datetime, timedelta
import re
import argparse
import sys

### define dates
today = datetime.now().strftime("%Y-%m-%d")
yesterday = (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d")
# today = "2021-03-26"
# yesterday = "2021-03-25"


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}] SONY Daily Exploration Area (Step Functions) for {today}"
    msg['From'] = "production-support@report-vm.ipa.dataart.net"
    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()

### AWS credentials are stored in /root/.aws/config for the root user
def fetch_aws_sfn_data():
    client = boto3.client('stepfunctions')
    response = client.list_executions(
        stateMachineArn='arn:aws:states:us-east-1:323555055331:stateMachine:prod-delphi-explorationFlow',
        statusFilter='FAILED'
    )
    return response


### find FAILED functions only for today and yesterday
def check_sfn_failed(data: dict, acceptable_fails_count: int):
    sfn_total_list = []
    sfn_prefix_names_list = []
    re_pattern_sfn_prefix = r'\w+-\d+-\w+-\w+-\w\d'
    # since this script runs at 05:00 UTC, the below re_pattern grabs all step functions between:
    # {yesterday 05:00 - 23:59} and {today's full time range up to the running moment}.
    # In such case we are able to observe any FAILED step function that might appears between CRON running interval
    re_pattern_sfn_time_range = rf'(?:{yesterday} (?:0[5-9]|[1-2][0-9])|{today} \d+):\d+:\d+'

    # find FAILED {sfn} for {re_pattern_sfn_time_range} period
    for sfn in data["executions"]:
        sfn.pop('executionArn')
        sfn.pop('stateMachineArn')
        sfn['startDate'] = str(sfn['startDate'])
        sfn['stopDate'] = str(sfn['stopDate'])
        # if re is True then add {sfn} to the list
        if re.findall(re_pattern_sfn_time_range, sfn['startDate']):
            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]
    return message_data



######################################################################################################################
################### ENTRY POINT ######################################################################################
######################################################################################################################
######################################################################################################################
if __name__ == '__main__':
    parser = create_arg_parser()
    arg_list = parser.parse_args()

    if len(sys.argv[1:]) > 0:
        acceptable_fails_count = arg_list.acceptable_fails_count[0]
        aws_sfn_data = fetch_aws_sfn_data()
        result = check_sfn_failed(data=aws_sfn_data, acceptable_fails_count=acceptable_fails_count)
        if result:
            message = json.dumps(result, indent=4)
            status = "ALERT"
        else:
            message = "No FAILED Step Functions were found"
            status = "OK"

        if arg_list.recipients:
            recipients = arg_list.recipients[0]
            send_mail(recipients=recipients, message=message, status=status)
        if arg_list.print:
            print(message)
    else:
        parser.print_help()