"""Module for inspecting chalice logs.

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

"""
import time
from datetime import datetime, timedelta
from collections import defaultdict

from typing import Any, Optional, Iterator, Dict, IO, Callable, Set  # noqa
from botocore.session import Session  # noqa
from attr import attrs, attrib

from chalice.awsclient import TypedAWSClient, CWLogEvent  # noqa
from chalice.utils import TimestampConverter


@attrs
class LogRetrieveOptions(object):
    max_entries = attrib(default=None)               # type: Optional[int]
    start_time = attrib(default=None)                # type: Optional[datetime]
    include_lambda_messages = attrib(default=False)  # type: bool

    @classmethod
    def create(cls, follow=False, since=None, **kwargs):
        # type: (bool, Optional[str], **Any) -> LogRetrieveOptions
        start_time = None
        if since is not None:
            start_time = TimestampConverter().timestamp_to_datetime(since)
        elif follow:
            # If they've specified --follow with no start time, we'll default
            # to 10 minutes prior to now.  This is a backwards compat
            # quirk where we can't change the default for `chalice logs`, but
            # for the --follow mode, we can default to only showing 10 minutes
            # worth of logs before we start polling for new logs.  In a
            # future major version we can just change the default for all log
            # retrieval to 10 minutes.
            start_time = datetime.utcnow() - timedelta(minutes=10)
        kwargs['start_time'] = start_time
        return cls(**kwargs)


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


class LogRetriever(object):
    def __init__(self, log_event_generator, log_group_name):
        # type: (BaseLogEventGenerator, str) -> None
        self._log_event_generator = log_event_generator
        self._log_group_name = log_group_name

    @classmethod
    def create_from_lambda_arn(cls, log_event_generator, lambda_arn):
        # type: (BaseLogEventGenerator, 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(log_event_generator, log_group_name)

    def _is_lambda_message(self, event):
        # type: (CWLogEvent) -> 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, retrieve_options):
        # type: (LogRetrieveOptions) -> Iterator[CWLogEvent]
        """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.

        """
        max_entries = retrieve_options.max_entries
        shown = 0
        events = self._log_event_generator.iter_log_events(
            self._log_group_name, retrieve_options)
        for event in events:
            if not retrieve_options.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


class BaseLogEventGenerator(object):
    def __init__(self, client):
        # type: (TypedAWSClient) -> None
        self._client = client

    def iter_log_events(self, log_group_name, options):
        # type: (str, LogRetrieveOptions) -> Iterator[CWLogEvent]
        pass


class LogEventGenerator(BaseLogEventGenerator):
    def iter_log_events(self, log_group_name, options):
        # type: (str, LogRetrieveOptions) -> Iterator[CWLogEvent]
        logs = self._client.iter_log_events(log_group_name=log_group_name,
                                            start_time=options.start_time)
        for log_message in logs:
            yield log_message


class FollowLogEventGenerator(BaseLogEventGenerator):

    _POLL_TIME = 5

    def __init__(self, client, sleep=time.sleep, poll_time=_POLL_TIME):
        # type: (TypedAWSClient, Callable[[int], None], int) -> None
        self._client = client
        self._sleep = sleep
        self._event_id_cache = defaultdict(
            set)  # type: Dict[datetime, Set[str]]
        self._poll_time = poll_time

    def iter_log_events(self, log_group_name, options):
        # type: (str, LogRetrieveOptions) -> Iterator[CWLogEvent]
        start_time = options.start_time
        try:
            for event in self._loop_on_filter_log_events(log_group_name,
                                                         start_time):
                yield event
        except KeyboardInterrupt:
            return

    def _loop_on_filter_log_events(self, log_group_name, start_time):
        # type: (str, Optional[datetime]) -> Iterator[CWLogEvent]
        self._event_id_cache.clear()
        kwargs = {
            'log_group_name': log_group_name,
            'start_time': start_time,
        }  # type: Dict[str, Any]
        while True:
            response = self._client.filter_log_events(**kwargs)
            for event in response['events']:
                if not self._in_cache(event):
                    self._add_to_cache(event)
                    yield event
            if 'nextToken' in response:
                # If there's more pages we go through the normal pagination
                # logic of following nextTokens.
                kwargs['next_token'] = response['nextToken']
            else:
                kwargs.pop('next_token', None)
                if self._event_id_cache:
                    # However, if there's no nextToken it means we've iterated
                    # through all the existing log events.  We now need to
                    # start polling for new events.  To do this, we need
                    # to use start time and not nextToken, because once
                    # we've exhausted the iteration we won't see new events
                    # from any log streams we've already iterated through.
                    # Querying only based on start time ensures we see new
                    # messages from all new streams.
                    # This isn't going to be perfect.  It's possible that
                    # we miss a gap between when we finished searching through
                    # a log stream and the new start time we're going to use
                    # to start polling, especially at high rates of log
                    # generation.
                    most_recent_start_time = max(self._event_id_cache)
                    kwargs['start_time'] = most_recent_start_time
                    self._prune_old_cache_entries(most_recent_start_time)
                    self._sleep(self._poll_time)

    def _in_cache(self, event):
        # type: (CWLogEvent) -> bool
        return event['eventId'] in self._event_id_cache[event['timestamp']]

    def _add_to_cache(self, event):
        # type: (CWLogEvent) -> None
        self._event_id_cache[event['timestamp']].add(event['eventId'])

    def _prune_old_cache_entries(self, timestamp):
        # type: (datetime) -> None
        surviving_entries = self._event_id_cache[timestamp]
        self._event_id_cache.clear()
        self._event_id_cache[timestamp] = surviving_entries
