"""Clients manage the main user facing API, and provide functions for sending
events and querying the Riemann server. UDP, TCP and TLS transports are
provided by the :py:mod:`riemann_client.transport` module, and the protocol
buffer objects are provided by the :py:mod:`riemann_client.riemann_pb2` module.
"""

from __future__ import absolute_import

import logging
try:
    from logging import NullHandler
except ImportError:
    # Create a NullHandler class in logging for python 2.6
    class NullHandler(logging.Handler):
        def emit(self, record):
            pass

import socket
try:
    from threading import RLock
    from threading import Timer
except ImportError:
    RLock = None
    Timer = None
import time

import riemann_client.riemann_pb2
import riemann_client.transport
import google.protobuf.descriptor


logger = logging.getLogger(__name__)
logger.addHandler(NullHandler())


class Client(object):
    """A client for sending events and querying a Riemann server.

    Two sets of methods are provided - an API dealing directly with protocol
    buffer objects and an extended API that takes and returns dictionaries
    representing events.

    Protocol buffer API:

        - :py:meth:`.send_event`
        - :py:meth:`.send_events`
        - :py:meth:`.send_query`

    Extended API:

        - :py:meth:`.event`
        - :py:meth:`.events`
        - :py:meth:`.query`

    Clients do not directly manage connections to a Riemann server - these are
    managed by :py:class:`riemann_client.transport.Transport` instances, which
    provide methods to read and write messages to the server. Client instances
    can be used as a context manager, and will connect and disconnect the
    transport when entering and exiting the context.

        >>> with Client(transport) as client:
        ...     # Calls transport.connect()
        ...     client.query('true')
        ...     # Calls transport.disconnect()
    """

    def __init__(self, transport=None):
        if transport is None:
            transport = riemann_client.transport.TCPTransport()
        self.transport = transport

    def __enter__(self):
        self.transport.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.transport.disconnect()

    @staticmethod
    def create_event(data):
        """Translates a dictionary of event attributes to an Event object

        :param dict data: The attributes to be set on the event
        :returns: A protocol buffer ``Event`` object
        """
        event = riemann_client.riemann_pb2.Event()
        event.host = socket.gethostname()
        event.tags.extend(data.pop('tags', []))

        for key, value in data.pop('attributes', {}).items():
            attribute = event.attributes.add()
            attribute.key, attribute.value = key, value

        for name, value in data.items():
            if value is not None:
                setattr(event, name, value)
        return event

    def send_events(self, events):
        """Sends multiple events to Riemann in a single message

        :param events: A list or iterable of ``Event`` objects
        :returns: The response message from Riemann
        """
        message = riemann_client.riemann_pb2.Msg()
        for event in events:
            message.events.add().MergeFrom(event)
        return self.transport.send(message)

    def send_event(self, event):
        """Sends a single event to Riemann

        :param event: An ``Event`` protocol buffer object
        :returns: The response message from Riemann
        """
        return self.send_events((event,))

    def events(self, *events):
        """Sends multiple events in a single message

        >>> client.events({'service': 'riemann-client', 'state': 'awesome'})

         :param \*events: event dictionaries for :py:func:`create_event`
         :returns: The response message from Riemann
        """
        return self.send_events(self.create_event(e) for e in events)

    def event(self, **data):
        """Sends an event, using keyword arguments to create an Event

        >>> client.event(service='riemann-client', state='awesome')

        :param \*\*data: keyword arguments used for :py:func:`create_event`
        :returns: The response message from Riemann
        """
        return self.send_event(self.create_event(data))

    @staticmethod
    def create_dict(event):
        """Translates an Event object to a dictionary of event attributes

        All attributes are included, so ``create_dict(create_event(input))``
        may return more attributes than were present in the input.

        :param event: A protocol buffer ``Event`` object
        :returns: A dictionary of event attributes
        """

        data = dict()

        for descriptor, value in event.ListFields():
            if descriptor.name == 'tags':
                value = list(value)
            elif descriptor.name == 'attributes':
                value = dict(((a.key, a.value) for a in value))
            elif descriptor.type == \
                    google.protobuf.descriptor.FieldDescriptor.TYPE_FLOAT:
                value = round(value, 1)
            data[descriptor.name] = value

        return data

    def send_query(self, query):
        """Sends a query to the Riemann server

        :returns: The response message from Riemann
        """
        message = riemann_client.riemann_pb2.Msg()
        message.query.string = query
        return self.transport.send(message)

    def query(self, query):
        """Sends a query to the Riemann server

        >>> client.query('true')

        :returns: A list of event dictionaries taken from the response
        :raises Exception: if used with a :py:class:`.UDPTransport`
        """
        if isinstance(self.transport, riemann_client.transport.UDPTransport):
            raise Exception('Cannot query the Riemann server over UDP')
        response = self.send_query(query)
        return [self.create_dict(e) for e in response.events]


class QueuedClient(Client):
    """A Riemann client using a queue that can be used to batch send events.

    A message object is used as a queue, with the :py:meth:`.send_event` and
    :py:meth:`.send_events` methods adding new events to the message and the
    :py:meth:`.flush` sending the message.
    """

    def __init__(self, transport=None):
        super(QueuedClient, self).__init__(transport)
        self.clear_queue()

    def flush(self):
        """Sends the waiting message to Riemann

        :returns: The response message from Riemann
        """
        response = self.transport.send(self.queue)
        self.clear_queue()
        return response

    def send_event(self, event):
        """Adds a single event to the queued message

        :returns: None - nothing has been sent to the Riemann server yet
        """
        self.send_events((event,))
        return None

    def send_events(self, events):
        """Adds multiple events to the queued message

        :returns: None - nothing has been sent to the Riemann server yet
        """
        for event in events:
            self.queue.events.add().MergeFrom(event)
        return None

    def clear_queue(self):
        """Resets the message/queue to a blank :py:class:`.Msg` object"""
        self.queue = riemann_client.riemann_pb2.Msg()


if RLock and Timer:  # noqa
    class AutoFlushingQueuedClient(QueuedClient):
        """A Riemann client using a queue and a timer that will automatically
        flush its contents if either:
            - the queue size exceeds :param max_batch_size: or
            - more than :param max_delay: has elapsed since the last flush and
             the queue is non-empty.

        if :param stay_connected: is False, then the transport will be
        disconnected after each flush and reconnected at the beginning of
        the next flush.
        if :param clear_on_fail: is True, then the client will discard its
        buffer after the second retry in the event of a socket error.

        A message object is used as a queue, and the following methods are
        given:
            - :py:meth:`.send_event` - add a new event to the queue
            - :py:meth:`.send_events` add a tuple of new events to the queue
            - :py:meth:`.event` - add a new event to the queue from
               keyword arguments
            - :py:meth:`.events` - add new events to the queue from
               dictionaries
            - :py:meth:`.flush` - manually force flush the queue to the
               transport
        """

        def __init__(self, transport, max_delay=0.5, max_batch_size=100,
                     stay_connected=False, clear_on_fail=False):
            super(AutoFlushingQueuedClient, self).__init__(transport)
            self.stay_connected = stay_connected
            self.clear_on_fail = clear_on_fail
            self.max_delay = max_delay
            self.max_batch_size = max_batch_size
            self.lock = RLock()
            self.event_counter = 0
            self.last_flush = time.time()
            self.timer = None

            # start the timer
            self.start_timer()

        def connect(self):
            """Connect the transport if it is not already connected."""
            if not self.is_connected():
                self.transport.connect()

        def is_connected(self):
            """Check whether the transport is connected."""
            try:
                # this will throw an exception whenever socket isn't connected
                self.transport.socket.type
                return True
            except (AttributeError, RuntimeError, socket.error):
                return False

        def event(self, **data):
            """Enqueues an event, using keyword arguments to create an Event

            >>> client.event(service='riemann-client', state='awesome')

            :param \*\*data: keyword arguments used for :py:func:`create_event`
            """
            self.send_events((self.create_event(data),))

        def events(self, *events):
            """Enqueues multiple events in a single message

            >>> client.events({'service': 'riemann-client',
            >>>                'state': 'awesome'})

             :param \*events: event dictionaries for :py:func:`create_event`
             :returns: The response message from Riemann
            """
            self.send_events(self.create_event(evd) for evd in events)

        def send_events(self, events):
            """Enqueues multiple events

            :param events: A list or iterable of ``Event`` objects
            :returns: The response message from Riemann
            """
            with self.lock:
                for event in events:
                    self.queue.events.add().MergeFrom(event)
                    self.event_counter += 1
                    self.check_for_flush()

        def flush(self):
            """Sends the events in the queue to Riemann in a single protobuf
            message

            :returns: The response message from Riemann
            """
            response = None
            with self.lock:
                if not self.is_connected():
                    self.connect()
                try:
                    response = super(AutoFlushingQueuedClient, self).flush()
                except socket.error:
                    # log and retry
                    logger.warn("Socket error on flushing. "
                                "Attempting reconnect and retry...")
                    try:
                        self.transport.disconnect()
                        self.connect()
                        response = (
                            super(AutoFlushingQueuedClient, self).flush())
                    except:
                        logger.warn("Socket error on flushing "
                                    "second attempt. Batch discarded.")
                        self.transport.disconnect()
                        if self.clear_on_fail:
                            self.clear_queue()
                self.event_counter = 0
                if not self.stay_connected:
                    self.transport.disconnect()
                self.last_flush = time.time()
            self.start_timer()
            return response

        def check_for_flush(self):
            """Checks the conditions for flushing the queue"""
            if (self.event_counter >= self.max_batch_size or
                    (time.time() - self.last_flush) >= self.max_delay):
                self.flush()

        def start_timer(self):
            """Cycle the timer responsible for periodically flushing the queue
            """
            if self.timer:
                self.timer.cancel()
            self.timer = Timer(self.max_delay, self.check_for_flush)
            self.timer.daemon = True
            self.timer.start()

        def stop_timer(self):
            """Stops the current timer

            a :py:meth:`.flush` event will reactviate the timer
            """
            self.timer.cancel()
