#!/usr/bin/env python
"""
This script deletes empty log streams from cloudwatch log groups.

Your shell should be configured for connecting to the AWS API. This can be done
with the CLI command `aws configure`.

This script requires Python 3.6 or newer, and you must have boto3 installed.
"""

import argparse
import os
import boto3

LOG_ONLY = False
LOOP_UNTIL_ZERO = False

log_client = boto3.client('logs')

def get_log_groups():
    """Return a list containing the names of all log groups."""

    found_all = False
    next_token = None
    group_names = []

    while not found_all:
        if next_token is not None:
            response = log_client.describe_log_groups(nextToken=next_token)
        else:
            response = log_client.describe_log_groups()
        if 'logGroups' in response:
            for group in response['logGroups']:
                if 'logGroupName' in group:
                    group_names.append(group['logGroupName'])
        next_token = response.get('nextToken', None)
        if next_token is None:
            found_all = True

    return group_names


def get_log_streams(group_name, max=500):
    """Return a list of log stream names in group `group_name`."""
    found_all = False
    next_token = None
    stream_names = {}

    while not found_all and len(stream_names) <= max:
        if next_token is not None:
            response = log_client.describe_log_streams(
                logGroupName=group_name,
                orderBy='LastEventTime',
                nextToken=next_token
            )
        else:
            response = log_client.describe_log_streams(
                logGroupName=group_name,
                orderBy = 'LastEventTime',
            )
        if 'logStreams' in response:
            for stream in response['logStreams']:
                if 'logStreamName' in stream:
                    stream_names[stream['logStreamName']] = stream
        next_token = response.get('nextToken', None)
        if next_token is None:
            found_all = True

    return stream_names


def stream_has_events(group_name, stream_name):
    """Return whether or not `stream_name` in `group_name` has any events."""
    # Return the first event.
    response = log_client.get_log_events(logGroupName=group_name,
                                     logStreamName=stream_name,
                                     limit=1)

    if 'events' in response and len(response['events']) > 0:
        return True
    return False


def get_and_delete_group_empty_streams(group_name):
    streams = get_log_streams(group_name, max=100)
    print(f"Returned {len(streams)} groups for {group_name}")
    num_deleted = 0
    for stream_name, obj in streams.items():

        if not stream_has_events(group_name, stream_name):
            if LOG_ONLY:
                print(
                    f'Stream `{stream_name}` in group `{group_name}` is empty. '
                    'Would be deleted...'
                )
            else:
                print(
                    f'Stream `{stream_name}` in group `{group_name}` is empty. '
                    'Deleting...'
                )
                del_stream(group_name, stream_name)
                num_deleted += 1
    return num_deleted


def del_stream(group_name, stream_name):
    """Delete `stream_name` from `group_name`."""
    log_client = boto3.client('logs')
    log_client.delete_log_stream(logGroupName=group_name,
                             logStreamName=stream_name)

def del_empty_streams(group_name):
    """Finds all empty streams and deletes them.

    If the global variable `LOG_ONLY` is True, streams are only logged to the
    console and not deleted.
    """
    for group in [group_name]: # get_log_groups():
        print(f'Finding and deleting empty streams in group `{group}`')
        num_deleted = get_and_delete_group_empty_streams(group)
        if LOOP_UNTIL_ZERO and num_deleted > 0:
            while num_deleted > 0:
                print(f"Sleeping for 10 seconds")
                import time; time.sleep(10)
                print(f"Deleted {num_deleted} from `{group}` going for more.")
                num_deleted = get_and_delete_group_empty_streams(group)


def main():
    """Main handler for this script."""
    parser = argparse.ArgumentParser(description='Delete empty log streams.')
    parser.add_argument(
        'group_name',
        action="store",
        help="Name of log group"
    )
    parser.add_argument(
        '--log-only',
        action='store_true',
        help="If specified, only log to console, don't actually delete."
    )
    parser.add_argument(
        '--loop-until-zero',
        action='store_true',
        help="If specified, go until the group is empty, otherwise a max will be deleted"
    )
    args = parser.parse_args()
    global LOG_ONLY
    LOG_ONLY = args.log_only
    global LOOP_UNTIL_ZERO
    LOOP_UNTIL_ZERO = args.loop_until_zero

    del_empty_streams(args.group_name)


if __name__ == "__main__":
    main()