"""Client for putting events on the events Kinesis stream."""

from __future__ import absolute_import

import base64
import json
import logging
import time
import uuid as uuid_mod

import boto3
from google.protobuf import message

from . import errors
from .proto import event_pb2


logger = logging.getLogger(__name__)

MAX_RECORD_BYTES = 1024 * 1024
MAX_PUT_RECORDS = 500
MAX_PUT_BYTES = 5 * 1024 * 1024


class EventBus(object):
    """AWS Kinesis stream that carries event protobufs.

    Args:
      stream_name: string name of the Kinesis stream name
      large_event_store: LargeEventStore object used to store large events.
      region: string name of the AWS region the stream is in

    Properties:
      kinesis: boto3 kinesis resource
      stream_name: Kinesis stream name
    """

    def __init__(self, stream_name, large_event_store, region='us-west-2'):
        self.kinesis = boto3.client('kinesis', region_name=region)
        self.stream_name = stream_name
        self.large_event_store = large_event_store

    def PutEvent(self, event):
        """Put Event protobuf on the Kinesis stream.

        Raises:
          EventInvalid if the event fields aren't set correct.
        """
        ValidateEvent(event)
        encoded_uuid = base64.b64encode(event.uuid)
        data = event.SerializeToString()
        if len(data) > MAX_RECORD_BYTES:
            logger.info('Item too large: %s', encoded_uuid)
            self.large_event_store.PutEvent(event)
            # TODO: publish event with field indicating where event body is.
        else:
            self.kinesis.put_record(
                StreamName=self.stream_name,
                Data=data,
                PartitionKey=encoded_uuid)
            logger.info('Published event UUID: %s', encoded_uuid)


def Event(event_type, body, uuid=None, timestamp=None, attributes=None):
    """Constructs Event protobuf.

    Args:
      event_type: event type string
      body: string of the body of the event
      uuid: optional byte string, if None, will generate one
      timestamp: time, in seconds, as a floating point, if None, will use "now"
      attributes: dict of queryable attributes to attach to this event.

    Returns:
      Event protobuf object
    """
    e = event_pb2.Event()
    if uuid:
        e.uuid = uuid
    else:
        e.uuid = uuid_mod.uuid4().bytes
    if timestamp is None:
        e.timestamp = time.time()
    else:
        e.timestamp = timestamp
    e.type = event_type
    e.body = body
    if attributes:
        for k, v in attributes.iteritems():
            attr = e.attributes.add()
            attr.key = k
            attr.value = v
    return e


def DecodeKinesisRecordData(data):
    """Decode a kinesis record payload back into an Event protobuf object.

    Args:
      data: kinesis record data string

    Returns:
      Event protobuf object

    Raises:
      EventDecodeError if it couldn't decode.
      EventInvalid if the event's fields weren't set correctly.
    """
    # Kinesis data is base64 encoded.
    try:
        payload = base64.b64decode(data)
    except TypeError as e:
        raise errors.EventDecodeError(
          'Failed to base64 decode kinesis record: {}'.format(e))

    # Deserialize payload into protobuf.
    try:
        event = event_pb2.Event.FromString(payload)
    except message.DecodeError as e:
        raise errors.EventDecodeError(
          'Failed to decode event data an Event protobuf: {}'.format(e))

    # Validate event has required fields.
    ValidateEvent(event)

    return event


def ValidateEvent(event):
    """Validate that an event's fields are set correctly.

    Args:
      event: Event protobuf object.

    Raises:
      EventInvalid on failure.
    """
    try:
        assert event.HasField('uuid'), 'uuid field not set'
        assert event.HasField('timestamp'), 'timestamp field not set'
        assert event.HasField('type'), 'type field not set'
        assert event.HasField('body'), 'body field not set'
    except AssertionError as e:
        raise errors.EventInvalid(str(e))
