import datetime
import os
import random
import string

import boto3

client = boto3.client('stepfunctions', region_name='us-east-1')

MAX_EXECUTIONS_TO_RETRY = int(os.environ.get(
    'MAX_EXECUTIONS_TO_RETRY', 1000))
STATE_MACHINE_ARN = os.environ.get(
    'STATE_MACHINE_ARN',
    'arn:aws:states:us-east-1:437795906767:stateMachine:prod_video_approval_workflow')

def main():
    "Retry all failed workflows from the past day"
    failed_executions = client.list_executions(
        stateMachineArn=STATE_MACHINE_ARN,
        statusFilter='FAILED',
        maxResults=MAX_EXECUTIONS_TO_RETRY,
    )

    five_days_ago = datetime.datetime.now(
        tz=datetime.timezone.utc) + datetime.timedelta(days=-5)
    for execution in failed_executions['executions']:
        print(f'-----')
        print(f'Execution is {execution}')
        # If less than five days old, retry
        if execution['stopDate'] > five_days_ago:
            root_name = execution['name'].split('_')[0]
            # Check to make sure a retry is not running
            retry_running = False
            active_executions = client.list_executions(
                stateMachineArn=STATE_MACHINE_ARN,
                statusFilter='RUNNING',
            )
            for active_execution in active_executions['executions']:
                if active_execution['name'].startswith(root_name):
                    print(f'Retry already running for {root_name} '
                          f'with name {active_execution["name"]}')
                    retry_running = True
                    break
            previously_succeeded = False
            succeeded_executions = client.list_executions(
                stateMachineArn=STATE_MACHINE_ARN,
                statusFilter='SUCCEEDED',
                maxResults=MAX_EXECUTIONS_TO_RETRY,
            )
            for succeeded_execution in succeeded_executions['executions']:
                if succeeded_execution['name'].startswith(root_name):
                    print(f'Run already succeeded for {root_name} '
                          f'with name {succeeded_execution["name"]}')
                    previously_succeeded = True
                    break
            if not retry_running and not previously_succeeded:
                execution_info = client.describe_execution(
                    executionArn=execution['executionArn']
                )
                name_suffix = ''.join(
                    random.choices(string.ascii_letters, k=10))
                if len(execution_info["name"].split('_')) > 1:
                    new_root_name = execution_info["name"].split('_')[0] + '_' + execution_info["name"].split('_')[1]
                    new_name = f'{new_root_name}_retry_{name_suffix}'
                    print(f'New name is {new_name}')
                    new_execution = client.start_execution(
                        stateMachineArn=STATE_MACHINE_ARN,
                        name=new_name,
                        input=execution_info['input'],
                    )
                    print(f'Execution started at {new_execution["startDate"]} '
                          f'with ARN {new_execution["executionArn"]}')
        else:
            print(f'Execution {execution["name"]} more than one day old')


if __name__ == "__main__":
    main()