"""Demo code for Demos and Donuts - AWS key injection in Jenkins. 2021-10 """
import boto3
import config


def get_full_execution_count_by_arn(sfn_arn, status='RUNNING'):
    """Get an aws SFN by sfn name."""

    # Connect to boto3 client
    sfn_client = boto3.client('stepfunctions')

    # Get first tranche of executions
    current_executions = sfn_client.list_executions(
        stateMachineArn=sfn_arn,
        statusFilter=status,
        maxResults=50
    )
    next_token = current_executions.get('nextToken')

    # Init counter
    current_sfn_count = len(current_executions['executions'])

    # Loop through tranches until no more are found
    while next_token:
        current_executions = sfn_client.list_executions(
            stateMachineArn=sfn_arn,
            statusFilter=status,
            maxResults=50,
            nextToken=next_token
        )
        current_sfn_count = current_sfn_count + len(
            current_executions['executions'])

        next_token = current_executions.get('nextToken')

    # Return full count
    return current_sfn_count


def main():
    """Main code for aws_jenkins_agent demo."""
    sfn_arn = config.BULK_ASSETS_INGESTER_SFN_ARN
    all_statuses = ['SUCCEEDED', 'TIMED_OUT', 'ABORTED', 'FAILED', 'RUNNING']

    msg = f'\n\n-------------------------' \
          f'\nExecution tally for SFN at {sfn_arn}:'
    print(msg)

    for status in all_statuses:
        execution_count = get_full_execution_count_by_arn(sfn_arn, status)

        msg = f'\'{status}\': {execution_count}'
        print(msg)


if __name__ == '__main__':
    main()