"""Module for inspecting chalice logs.

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

"""
from __future__ import annotations
import time
from datetime import datetime, timedelta
from dataclasses import dataclass
from collections import defaultdict

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

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


@dataclass
class LogRetrieveOptions(object):
    max_entries: Optional[int] = None
    start_time: Optional[datetime] = None
    include_lambda_messages: bool = False

    @classmethod
    def create(
        cls, follow: bool = False, since: Optional[str] = None, **kwargs: 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: LogRetriever,
    stream: IO[str],
    retrieve_options: 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: BaseLogEventGenerator, log_group_name: 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: BaseLogEventGenerator, lambda_arn: 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: 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: 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: TypedAWSClient) -> None:
        self._client = client

    def iter_log_events(
        self, log_group_name: str, options: LogRetrieveOptions
    ) -> Iterator[CWLogEvent]:
        raise NotImplementedError("iter_log_events")


class LogEventGenerator(BaseLogEventGenerator):
    def iter_log_events(
        self, log_group_name: str, options: 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: TypedAWSClient,
        sleep: Callable[[int], None] = time.sleep,
        poll_time: int = _POLL_TIME,
    ) -> None:
        self._client = client
        self._sleep = sleep
        self._event_id_cache: Dict[datetime, Set[str]] = defaultdict(set)
        self._poll_time = poll_time

    def iter_log_events(
        self, log_group_name: str, options: 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: str, start_time: Optional[datetime]
    ) -> Iterator[CWLogEvent]:
        self._event_id_cache.clear()
        kwargs: Dict[str, Any] = {
            'log_group_name': log_group_name,
            'start_time': start_time,
        }
        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: CWLogEvent) -> bool:
        return event['eventId'] in self._event_id_cache[event['timestamp']]

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

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