"""Transports are used for direct communication with the Riemann server. They
are usually used inside a :py:class:`.Client`, and are used to send and receive
protocol buffer objects."""

from __future__ import absolute_import

import abc
import socket
import ssl
import struct

import riemann_client.riemann_pb2


# Default arguments
HOST = 'localhost'
PORT = 5555
TIMEOUT = None


def socket_recvall(socket, length, bufsize=4096):
    """A helper method to read of bytes from a socket to a maximum length"""
    data = b""
    while len(data) < length:
        data += socket.recv(bufsize)
    return data


class RiemannError(Exception):
    """Raised when the Riemann server returns an error message"""
    pass


class Transport(object):
    """Abstract transport definition

    Subclasses must implement the :py:meth:`.connect`, :py:meth:`.disconnect`
    and :py:meth:`.send` methods.

    Can be used as a context manager, which will call :py:meth:`.connect` on
    entry and :py:meth:`.disconnect` on exit.
    """

    __metaclass__ = abc.ABCMeta

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

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

    @abc.abstractmethod
    def connect(self):
        pass

    @abc.abstractmethod
    def disconnect(self):
        pass

    @abc.abstractmethod
    def send(self):
        pass


class SocketTransport(Transport):
    """Provides common methods for Transports that use a sockets"""

    def __init__(self, host=HOST, port=PORT):
        self.host = host
        self.port = port

    @property
    def address(self):
        """
        :returns: A tuple describing the address to connect to
        :rtype: (host, port)
        """
        return self.host, self.port

    @property
    def socket(self):
        """Returns the socket after checking it has been created"""
        if not hasattr(self, '_socket'):
            raise RuntimeError("Transport has not been connected!")
        return self._socket

    @socket.setter
    def socket(self, value):
        self._socket = value


class UDPTransport(SocketTransport):
    def connect(self):
        """Creates a UDP socket"""
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

    def disconnect(self):
        """Closes the socket"""
        self.socket.close()

    def send(self, message):
        """Sends a message, but does not return a response

        :returns: None - can't receive a response over UDP
        """
        self.socket.sendto(message.SerializeToString(), self.address)
        return None


class TCPTransport(SocketTransport):
    def __init__(self, host=HOST, port=PORT, timeout=TIMEOUT):
        """Communicates with Riemann over TCP

        :param str host: The hostname to connect to
        :param int port: The port to connect to
        :param int timeout: The time in seconds to wait before raising an error
        """
        super(TCPTransport, self).__init__(host, port)
        self.timeout = timeout

    def connect(self):
        """Connects to the given host"""
        self.socket = socket.create_connection(self.address, self.timeout)

    def disconnect(self):
        """Closes the socket"""
        self.socket.close()

    def send(self, message):
        """Sends a message to a Riemann server and returns it's response

        :param message: The message to send to the Riemann server
        :returns: The response message from Riemann
        :raises RiemannError: if the server returns an error
        """
        message = message.SerializeToString()
        self.socket.sendall(struct.pack('!I', len(message)) + message)

        length = struct.unpack('!I', self.socket.recv(4))[0]
        response = riemann_client.riemann_pb2.Msg()
        response.ParseFromString(socket_recvall(self.socket, length))

        if not response.ok:
            raise RiemannError(response.error)

        return response


class TLSTransport(TCPTransport):
    def __init__(self, host=HOST, port=PORT, timeout=TIMEOUT, ca_certs=None):
        """Communicates with Riemann over TCP + TLS

        Options are the same as :py:class:`.TCPTransport` unless noted

        :param str ca_certs: Path to a CA Cert bundle used to create the socket
        """
        super(TLSTransport, self).__init__(host, port, timeout)
        self.ca_certs = ca_certs

    def connect(self):
        """Connects using :py:meth:`TLSTransport.connect` and wraps with TLS"""
        super(TLSTransport, self).connect()
        self.socket = ssl.wrap_socket(
            self.socket,
            ssl_version=ssl.PROTOCOL_TLSv1,
            cert_reqs=ssl.CERT_REQUIRED,
            ca_certs=self.ca_certs)


class BlankTransport(Transport):
    """A transport that collects events in a list, and has no connection

    Used by ``--transport none``, which is useful for testing commands without
    contacting a Riemann server. This is also used by the automated tests in
    ``riemann_client/tests/test_riemann_command.py``.
    """

    def __init__(self, *args, **kwargs):
        self.events = []

    def connect(self):
        """Creates a list to hold messages"""
        pass

    def send(self, message):
        """Adds a message to the list, returning a fake 'ok' response

        :returns: A response message with ``ok = True``
        """
        for event in message.events:
            self.events.append(event)
        reply = riemann_client.riemann_pb2.Msg()
        reply.ok = True
        return reply

    def disconnect(self):
        """Clears the list of messages"""
        pass

    def __len__(self):
        return len(self.events)
