"""Module for inspecting chalice logs.

This module provides APIs for searching, interacting
with the logs generated by AWS Lambda.

"""
from typing import Any, Optional, Iterator, Dict, IO  # noqa
from botocore.session import Session  # noqa

from chalice.awsclient import TypedAWSClient  # noqa


def display_logs(retriever, max_entries,
                 include_lambda_messages, stream):
    # type: (LogRetriever, int, bool, IO[str]) -> None
    events = retriever.retrieve_logs(
        include_lambda_messages=include_lambda_messages,
        max_entries=max_entries)
    for event in events:
        stream.write('%s %s %s\n' % (
            event['timestamp'],
            event['logShortId'],
            event['message'].strip()))


class LogRetriever(object):
    def __init__(self, client, log_group_name):
        # type: (TypedAWSClient, str) -> None
        self._client = client
        self._log_group_name = log_group_name

    @classmethod
    def create_from_lambda_arn(cls, client, lambda_arn):
        # type: (TypedAWSClient, str) -> LogRetriever
        """Create a LogRetriever from a client and lambda arn.

        :type client: botocore.client.Logs
        :param client: A ``logs`` client.

        :type lambda_arn: str
        :param lambda_arn: The ARN of the lambda function.

        :return: An instance of ``LogRetriever``.

        """
        lambda_name = lambda_arn.split(':')[6]
        log_group_name = '/aws/lambda/%s' % lambda_name
        return cls(client, log_group_name)

    def _is_lambda_message(self, event):
        # type: (Dict[str, Any]) -> bool
        # Lambda will also inject log messages into your log streams.
        # They look like:
        # START RequestId: guid Version: $LATEST
        # END RequestId: guid
        # REPORT RequestId: guid    Duration: 0.35 ms   Billed Duration: ...

        # By default, these message are included in retrieve_logs().
        # But you can also request that retrieve_logs() filter out
        # these message so that we only include log messages generated
        # by your chalice app.
        msg = event['message'].strip()
        return msg.startswith(('START RequestId',
                               'END RequestId',
                               'REPORT RequestId'))

    def retrieve_logs(self, include_lambda_messages=True, max_entries=None):
        # type: (bool, Optional[int]) -> Iterator[Dict[str, Any]]
        """Retrieve logs from a log group.

        :type include_lambda_messages: boolean
        :param include_lambda_messages: Include logs generated by the AWS
            Lambda service.  If this value is False, only chalice logs will be
            included.

        :type max_entries: int
        :param max_entries: Maximum number of log messages to include.

        :rtype: iterator
        :return: An iterator that yields event dicts.  Each event
            dict has these keys:

            * logStreamName -> (string) The name of the log stream.
            * timestamp -> (datetime.datetime) - The timestamp for the msg.
            * message -> (string) The data contained in the log event.
            * ingestionTime -> (datetime.datetime) Ingestion time of event.
            * eventId -> (string) A unique identifier for this event.
            * logShortId -> (string) Short identifier for logStreamName.

        """
        # TODO: Add support for startTime/endTime.
        shown = 0
        for event in self._client.iter_log_events(self._log_group_name,
                                                  interleaved=True):
            if not include_lambda_messages and \
                    self._is_lambda_message(event):
                continue
            # logStreamName is: '2016/07/05/[id]hash'
            # We want to extract the hash portion and
            # provide a short identifier.
            identifier = event['logStreamName']
            if ']' in identifier:
                index = identifier.find(']')
                identifier = identifier[index + 1:index + 7]
            event['logShortId'] = identifier
            yield event
            shown += 1
            if max_entries is not None and shown >= max_entries:
                return
