"""
Usage: log-uploader [options] LOG_FILE LOG_GROUP LOG_STREAM
       log-uploader --version


This will upload the given log to stream, each line as a log event. Every
line in the log should start with an '[<ISO8601-TIMESTAMP>]' to allow this
script to correctly add that data to logged data.

This script assumes we can write data faster towards AWS than it is being
generated by the serial logger to the log file(s).

Options:
  --watch              Stay on foreground and upload new events.
  --ignore-timestamp   Ignore stored timestamp and upload all events.
  --verbose            Enable verbose logging.
"""

import boto3
import botocore
import docopt
import logging
import os
import time
import datetime
import re
import itertools
from typing import Optional, NamedTuple, Iterable
import multiprocessing as mp
import queue
from pathlib import Path

TIMESTAMP_FILE_SUFFIX = ".lus"
AWS_MAX_BATCH_SIZE = 1048576
AWS_MAX_EVENT_TIME_SPAN = 24 * 3600

logger = logging.getLogger(__name__)

here = Path(__file__).parent.resolve()


class ChunkUploader:
    def __init__(self,
                 client,
                 group: str,
                 stream: str,
                 upload_sequence_token: str,
                 min_time_between_requests: float,
                 max_lines_in_batch: int,
                 max_time_between_uploads: float,
                 max_batch_size: int = AWS_MAX_BATCH_SIZE,
                 max_event_time_span: int = AWS_MAX_EVENT_TIME_SPAN,
                 timestamp_file_name: Optional[str] = None):
        self.batch = []
        self.min_time_between_requests = min_time_between_requests
        self.max_time_between_uploads = max_time_between_uploads
        self.max_event_time_span = max_event_time_span
        self.max_batch_size = max_batch_size
        self.prev_ts = time.time()
        self.client = client
        self.group = group
        self.stream = stream

        self.timestamp_file_name = timestamp_file_name
        self.upload_sequence_token = upload_sequence_token
        self._last_event_ts = None
        self._current_batch_size = 0
        self.max_lines_in_batch = max_lines_in_batch

    def append(self, event):
        # Timestamps need to be increasing in a batch.
        if self._last_event_ts and self._last_event_ts > event.timestamp:
            self.upload_current_batch()

        if self.batch and (event.timestamp - self.batch[0].timestamp) > datetime.timedelta(seconds=self.max_event_time_span):
            self.upload_current_batch()

        EVENT_HEADER_SIZE = 26
        event_size = len(event.message.encode('utf-8')) + EVENT_HEADER_SIZE
        if event_size + self._current_batch_size > self.max_batch_size:
            self.upload_current_batch()

        if event.message:
            self.batch.append(event)
            self._current_batch_size += event_size
        self._last_event_ts = event.timestamp

        self.maybe_upload()

    def maybe_upload(self):
        if len(self.batch) >= self.max_lines_in_batch:
            self.upload_current_batch()
        else:
            elapsed = self._elapsed()
            if elapsed and elapsed > self.max_time_between_uploads:
                self.upload_current_batch()

    def _update_timestamp(self, ts: datetime.datetime):
        if not self.timestamp_file_name:
            return

        try:
            with open(self.timestamp_file_name, "w") as fp:
                fp.write(ts.isoformat())
        except PermissionError:
            logger.debug("Can't write timestamp file")

    def upload_current_batch(self):
        if not self.batch:
            return

        self._maybe_throttle()
        logger.debug(f"Uploading batch")

        # From https://docs.aws.amazon.com/AmazonCloudWatchLogs/latest/APIReference/API_PutLogEvents.html
        #  The batch of events must satisfy the following constraints:
        #
        # The maximum batch size is 1,048,576 bytes. This size is
        # calculated as the sum of all event messages in UTF-8, plus 26
        # bytes for each log event.
        # None of the log events in the batch can be more than 2 hours in
        # the future.
        # None of the log events in the batch can be older than 14 days or
        # older than the retention period of the log group.
        # The log events in the batch must be in chronological order by
        # their timestamp. The timestamp is the time the event occurred,
        # expressed as the number of milliseconds after Jan 1,
        # 1970 00:00:00 UTC. (In AWS Tools for PowerShell and the AWS SDK
        # for .NET, the timestamp is specified in .NET format:
        # yyyy-mm-ddThh:mm:ss. For example, 2017-09-15T13:45:30.)
        # A batch of log events in a single request cannot span more than
        # 24 hours. Otherwise, the operation fails.
        # The maximum number of log events in a batch is 10,000.
        # There is a quota of 5 requests per second per log stream.
        # Additional requests are throttled. This quota can't be changed.

        self.upload_sequence_token = upload_batch(self.client, self.group,
                                                  self.stream, self.batch,
                                                  self.upload_sequence_token)
        self._update_timestamp(self._last_event_ts)
        self.batch = []
        self.prev_ts = time.time()
        self._current_batch_size = 0

    def _maybe_throttle(self):
        if self.prev_ts is not None:
            elapsed = self._elapsed()
            if elapsed < self.min_time_between_requests:
                throttle_seconds = self.min_time_between_requests - elapsed
                logger.info(
                    f"Throttling for {throttle_seconds}s to avoid hitting "
                    f"request quota")
                time.sleep(throttle_seconds)

    def _elapsed(self):
        if not self.prev_ts:
            return None
        else:
            elapsed = time.time() - self.prev_ts
            return elapsed


class EOF:
    pass


def log_entry_producer(q, log_filename, watch: bool = False, read_timestamp: bool = True):
    try:
        for ev in get_log_entries(log_filename, watch=watch, read_timestamp=read_timestamp):
            q.put(ev)
        q.put(EOF())
    except Exception as e:
        q.put(e)
        raise


def upload_log(name: str, group: str, stream: str, watch: bool = False,
               max_time_between_uploads: float = 5,
               min_time_between_requests: float = 0.2,
               max_lines_in_batch: int = 5000,
               max_batch_size: int = AWS_MAX_BATCH_SIZE,
               timestamp_file_name: Optional[str] = None,
               read_timestamp: bool = True):
    """

    :param name: Name of the log file.
    :param group: Log Group in AWS CloudWatch.
    :param stream: Log Stream in AWS CloudWatch.
    :param watch: Instead of returning when log file is at the end, keep reading and waiting for more input.
    :param max_time_between_uploads: If there is any input, avoid waiting longer than this to upload.
    :param min_time_between_requests: Wait at least this long between calls to avoid hitting request quota.
    :param max_lines_in_batch: Maximum number of events in a single batch.
    :param max_batch_size: Maximum size, in bytes, for a single batch.
    :param timestamp_file_name: Timestamp file name, leave as None for default.
    :param read_timestamp: If True, skip events while events have earlier time than the timestamp found in the `timestamp_file_name`.
    :return:
    """
    client = boto3.client('logs')

    upload_sequence_token = get_next_sequence_token(client, group, stream)

    if timestamp_file_name is None:
        timestamp_file_name = get_stamp_filename(name)
    uploader = ChunkUploader(client, group, stream, upload_sequence_token,
                             min_time_between_requests,
                             max_lines_in_batch,
                             max_time_between_uploads,
                             max_batch_size,
                             timestamp_file_name=timestamp_file_name)

    q = mp.Queue()
    producer = mp.Process(target=log_entry_producer, args=(q, name),
                          kwargs={'watch': watch, 'read_timestamp': read_timestamp})
    producer.daemon = True
    producer.start()

    while True:
        try:
            ev = q.get(block=True, timeout=max(max_time_between_uploads/5, 0.1))
            if isinstance(ev, EOF):
                logger.info("Got EOF")
                break
            elif isinstance(ev, Exception):
                logger.error(f"Reader encountered error; {ev}", exc_info=ev)
                break
            else:
                logger.debug(f"got line {ev.message}")
                uploader.append(ev)
        except queue.Empty:
            uploader.maybe_upload()
            if not producer.is_alive():
                logger.error("Producer seems to have died.")
                break

    logger.info("stopping")
    uploader.upload_current_batch()
    q.close()
    producer.join()


def upload_batch(client, group, stream, batch, token):
    rest = {}
    if token:
        rest['sequenceToken'] = token
    response = client.put_log_events(
        logGroupName=group,
        logStreamName=stream,
        logEvents=[dict(timestamp=int(ev.timestamp.timestamp() * 1000),
                        message=ev.message) for ev in batch],
        **rest)
    logger.info(f"Wrote {len(batch)} log events, rejections: "
                f"{response.get('rejectedLogEventsInfo', '<none>')}")
    return response.get('nextSequenceToken', "")


def get_next_sequence_token(client, group, stream):
    response = client.describe_log_streams(logGroupName=group,
                                           logStreamNamePrefix=stream)
    matching = list(filter(lambda xx: xx['logStreamName'] == stream, response['logStreams']))
    if matching:
        stream_description = matching[0]
        logger.debug(f"Stream description for {stream}: {stream_description}")
        sequence_token = stream_description.get('uploadSequenceToken', "")
    else:
        try:
            client.create_log_stream(logGroupName=group, logStreamName=stream)
        except botocore.exceptions.ClientError as e:
            logger.warning(f"Failed to create stream {stream} to log group {group}; {e}")
        sequence_token = ""
    return sequence_token


class Event(NamedTuple):
    timestamp: datetime.datetime
    message: str


def parse_log_line(line: Optional[str]) -> Optional[Event]:
    if line is None:
        return None

    m = re.search(r"(?P<timestamp>\d+-\d+-\d+[T ]\d{2}:\d{2}:\d{2}(\.\d+)?([+-]\d{2}:\d{2})?)]?:? ?(?P<msg>.*$)", line)
    if m is None:
        logger.warning(f"line \"{line}\" did not match when finding timestamp")
        return None

    timestamp = m.group("timestamp")
    try:
        dt = datetime.datetime.fromisoformat(timestamp)
    except ValueError:
        logger.warning(f"line \"{line}\" had an invalid timestamp.")
        return None

    if dt.tzinfo is None:
        # Treat as old format localtime log timestamp
        dt = dt.astimezone()
    msg = m.group("msg")
    return Event(timestamp=dt, message=msg)


def get_log_entries(name: str, watch: bool = False, read_timestamp: bool = True) -> Iterable[Event]:
    """
    Return a generator yielding log events. If `watch` is True, tracks file
    changes to continue reading from a newly created file with the same name.

    :param name:
    :param watch:
    :param read_timestamp:
    :return:
    """
    if read_timestamp:
        timestamp = get_timestamp(name)
    else:
        timestamp = None

    def get_lines(fp):
        line = "start"
        while line or watch:
            line = fp.readline()
            if line:
                yield line
            else:
                if watch:
                    # in case file is just being rotated.
                    if os.path.exists(name) and os.stat(name).st_ino != current_inode:
                        return
                    else:
                        time.sleep(0.5)
                        continue
                else:
                    return None

    while True:
        try:
            current_inode = os.stat(name).st_ino
            with open(name, errors="backslashreplace") as logfile:
                for event in itertools.dropwhile(
                        lambda arg: timestamp is not None and arg and arg.timestamp <= timestamp,
                        filter(None, (parse_log_line(line) for line in get_lines(logfile)))):
                    yield event
        except OSError as e:
            if e.errno == 2:
                time.sleep(0.5)
            else:
                raise

        if not watch:
            return


def get_timestamp(name) -> Optional[datetime.datetime]:
    state_file_name = get_stamp_filename(name)
    if os.path.exists(state_file_name):
        line = parse_log_line(open(state_file_name).read().strip())
        if line:
            return line.timestamp
        else:
            return None
    else:
        return None


def get_stamp_filename(name):
    return name + TIMESTAMP_FILE_SUFFIX


def output_version():
    print("Metsuri log collector, version: " + open(here / "VERSION").read().strip())


def main():
    mp.set_start_method('spawn')

    opts = docopt.docopt(__doc__)

    if opts['--version']:
        output_version()
        return

    if opts['--verbose']:
        level = logging.DEBUG
    else:
        level = logging.INFO
    logging.basicConfig(level=level,
                        format='%(asctime)s [%(levelname)s] %('
                               'filename)s:%(lineno)s %(funcName)s %('
                               'message)s',
                        datefmt="%Y-%m-%dT%H:%M:%S%z"
                        )

    upload_log(opts["LOG_FILE"], opts["LOG_GROUP"], opts["LOG_STREAM"],
               watch=opts['--watch'], read_timestamp=not opts['--ignore-timestamp'])
    logger.info("exiting")


if __name__ == "__main__":
    main()
