import argparse
import random
import string
import sys

import boto3

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


def main():
    parser = argparse.ArgumentParser()
    
    parser.add_argument(
        '-u',
        '--upc',
        required=True,
        help='UPC to retry',
    )
    
    args = parser.parse_args()
    matched = False
    next_token = None

    failed_params = {
        'stateMachineArn': 'arn:aws:states:us-east-1:437795906767:stateMachine:prod_workflow_transfer_assets',
        'statusFilter': 'FAILED',
        'maxResults': 1000,
    }

    while not matched:
        failed_executions = client.list_executions(**failed_params)
        for execution in failed_executions['executions']:
            if args.upc in execution['executionArn']:
                print(f'{args.upc} matched in {execution["executionArn"]}')
                execution_info = client.describe_execution(
                    executionArn=execution['executionArn']
                )
                if len(execution_info["name"].split('_')) > 1:
                    name_suffix = ''.join(random.choices(string.ascii_letters, k=10))
                    new_root_name = execution_info["name"].split(
                        '_')[0] + '_' + execution_info["name"].split('_')[1]
                    new_name = f'{new_root_name}_retry_{name_suffix}'
                    new_execution = client.start_execution(
                        stateMachineArn='arn:aws:states:us-east-1:437795906767:stateMachine:prod_workflow_transfer_assets',
                        name=new_name,
                        input=execution_info['input'],
                    )
                    print(f'Execution started at {new_execution["startDate"]} with ARN {new_execution["executionArn"]}')
                    matched = True
                    break
        next_token = failed_executions.get('nextToken')
        if next_token:
            failed_params['nextToken'] = next_token
        else:
            sys.exit('upc not found')


if __name__ == '__main__':
    main()