#!/usr/bin/env python3

"""
** Makes the communication medium and its use independent. **
-------------------------------------------------------------

Whether it is tcp sockets, files, mailboxes or a drive like dropbox,
we want to be able to use all these media indifferently...
That's the role of this module.
It allows to add an abstraction layer to unify the communication API.
"""

import abc
import socket
import threading
import time
import uuid

from raisin.serialization import deserialize, serialize
from raisin.serialization.iter_tools import anticipate
from raisin.serialization.core import size2tag, tag2size
from raisin.encapsulation.formatting import format_package, unpack_message


__pdoc__ = {
    'AbstractConn.__enter__' : True,
    'AbstractConn.__exit__' : True,
}


class AbstractConn(abc.ABC):
    """
    ** Allows for easier communication. **

    It is the basic abstract class that allows to make information
    carriers of very different nature communicate in the same way.
    """
    @abc.abstractmethod
    def get_conn_names(self):
        """
        ** Retrieves the names of all accessible connections. **

        Must be overwritten.

        Returns
        -------
        conn_names : set
            All the identifiers of the recipients that
            can be contacted via the keyword *dest*.
        """
        raise NotImplementedError('this abstract method must be overwritten')

    @abc.abstractmethod
    def send(self, raw, dest=None):
        """
        ** Sends a data stream to the recipient. **

        Must be overwritten.

        Parameters
        ----------
        raw : iterator
            A generator that yields byte strings. The division is arbitrary.
            On reception, the packets received are not necessarily split in the same way.
            This flow can typically be generated by the ``raisin.serialization.serialize`` function.
        dest : str, default=None
            The address of the recipient(s) of the message.
            By default (None), the message is sent to the main recipient.

        Raises
        ------
        KeyError
            If the address does not exist.
        """
        raise NotImplementedError('this abstract method must be overwritten')

    def send_obj(self, obj, dest=None, **kwargs):
        """
        ** Sends an object to the recipient. **

        Can be overwritten.

        Parameters
        ----------
        obj
            It is the serialisable object that is sent to the recipient(s).
        dest : str, default=None
            Same as ``AbstractConn.send``.
        **kwargs : dict
            These are the parameters passed to the ``raisin.serialization.serialize`` function.
        """
        self.send(serialize(obj, **kwargs), dest=dest)

    @abc.abstractmethod
    def recv(self, dest=None):
        """
        ** Recovers the raw data sent by the sender. **

        Must be overwritten.

        Parameters
        ----------
        dest : str, default=None
            The address of the sender of the message.
            By default (None), the message is read to the main connection.

        Yields
        ------
        bytes
            The packets of bytes emitted by the ``AbstractConn.send``
            method of the sender. Packet decoding may differ, but the order is preserved.

        Raises
        ------
        ConnectionError
            If the connection is broken.
        KeyError
            If the address does not exist.
        """
        raise NotImplementedError('this abstract method must be overwritten')

    def recv_obj(self, **kwargs):
        """
        ** Recover and deserialize the sent object. **

        Can be overwritten.

        Parameters
        ----------
        **kwargs : dict
            These are the parameters passed to the ``raisin.serialization.deserialize`` function.

        Returns
        -------
        obj
            The emitted object.
        """
        return deserialize(self.recv(), **kwargs)

    @abc.abstractmethod
    def close(self):
        """
        ** Definitely close the connection. **

        Must be overwritten.
        """
        raise NotImplementedError('this abstract method must be overwritten')

    def __del__(self):
        """
        ** Help for the garage-collector. **
        """
        self.close()

    def __enter__(self):
        """
        ** Allows to manage the closure through a context manager. **
        """
        return self

    def __exit__(self, exc_type, value, traceback):
        """
        ** Guaranteed that the connection closes in all cases. **
        """
        self.close()


class SelectiveConn(AbstractConn):
    """
    ** Allows you to receive a specific message. **

    In particular, this class allows to retrieve a message without losing the others.
    For example if a thread is waiting for a response, the message that is intended
    for it must not be intercepted by another thread.
    This class allows to manage these conflicts.
    """
    def __init__(self):
        self.lock_send = threading.Lock()
        self.lock_recv = threading.Lock()
        self.lock_unpack = threading.Lock()
        self.res_buff = {} # to each signature, associates the received object

    def send_formatted(self, content, kind, signature=None, **kwargs):
        """
        ** Allows you to send information by indicating its nature. **

        For example, it is necessary to be able to separate questions from answers.

        Parameters
        ----------
        content
            Forwarded to ``raisin.encapsulation.formatting.format_package``.
        kind : str
            Forwarded to ``raisin.encapsulation.formatting.format_package``.
        signature : bytes
            Forwarded to ``raisin.encapsulation.formatting.format_package``.
        **kwargs : dict
            The parameters directly transmitted to the ``AbstractConn.send_obj`` method.

        Notes
        -----
        This method is thread-safe because it uses a lock.
        """
        message = format_package(content, kind, signature)
        with self.lock_send:
            self.send_obj(message, **kwargs)

    def recv_formatted(self, *, kind=None, signature=None, get_context=False, **kwargs):
        """
        ** Receives an object and makes verifications on its content. **

        Parameters
        ----------
        kind : str or set, optional
            Selects only the packages that match this type of packages.
        signature : bytes, optional
            Select only the package with this signature.
        get_context : boolean
            If True, also returns the context, not just the content.
        **kwargs : dict
            The parameters directly transmitted to the ``AbstractConn.recv_obj`` method.

        Returns
        -------
        content
            The content of the package that meets the selection criteria.
        context : namedtuple
            Returns this value only if *get_context* is True.

        Examples
        --------
        >>> import socket
        >>> from raisin.communication.abstraction import SocketConn
        >>> soc1, soc2 = socket.socketpair()
        >>> with SocketConn(soc1) as conn1, SocketConn(soc2) as conn2:
        ...     conn1.send_formatted(0, kind='answer', signature=b'0')
        ...     conn1.send_formatted(1, kind='answer', signature=b'1')
        ...     conn1.send_formatted(2, kind='answer', signature=b'2')
        ...     conn1.send_formatted('a', kind='request')
        ...     conn1.send_formatted('b', kind='request')
        ...     conn2.recv_formatted(signature=b'1')
        ...     conn2.recv_formatted(kind='request')
        ...     conn2.recv_formatted(kind='request', get_context=True)
        ...     conn2.recv_formatted(signature=b'0')
        ...     conn2.recv_formatted(signature=b'2', get_context=True)
        ...
        1
        'a'
        ('b', context(kind='request', signature=None))
        0
        (2, context(kind='answer', signature=b'2'))
        >>>
        """
        if isinstance(kind, str):
            kind = {kind}
        while True:
            with self.lock_unpack:
                for context in self.res_buff:
                    if kind is None or context.kind in kind:
                        if (signature is None
                            or context._asdict().get('signature', None) == signature
                        ):
                            if get_context:
                                return self.res_buff.pop(context), context
                            return self.res_buff.pop(context)
            if self.lock_recv.locked():
                time.sleep(0.5)
                continue
            with self.lock_recv:
                message = self.recv_obj(**kwargs)
                with self.lock_unpack:
                    content, context = unpack_message(message)
                    self.res_buff[context] = content

    def dialog(self, request, **kwargs):
        """
        ** Send a request and wait for the answer. **

        This method can be called from a thread without any risk
        of interference, because everything is set up to make it work.
        This method is blocking: it waits until the result is available.

        Parameters
        ----------
        request
            The content of the request.
        **kwargs : dict
            The parameters directly transmitted to the ``SelectiveConn.send_formatted``
            and ``SelectiveConn.recv_formatted`` method.

        Returns
        -------
        answer
            The content of the result is linked to the issued request.

        Examples
        --------
        >>> import uuid
        >>> uuid.uuid4() # doctest: +SKIP
        4545454
        >>>
        >>> import socket
        >>> from raisin.communication.abstraction import SocketConn
        >>> from raisin.communication.handler import Handler
        >>> soc1, soc2 = socket.socketpair()
        >>> with SocketConn(soc1) as conn, Handler(SocketConn(soc2)) as handler:
        ...     handler.start()
        ...     conn.dialog('hello')
        ...
        'hello'
        >>>
        """
        signature = uuid.uuid4().bytes
        self.send_formatted(request, kind='request', signature=signature, **kwargs)
        return self.recv_formatted(kind='answer', signature=signature, **kwargs)


class SocketConn(SelectiveConn):
    """
    ** Adds an abstraction layer to TCP sockets. **

    Attributes
    ----------
    conn : socket.socket
        A socket capable of reading and sending data.
    """

    def __init__(self, conn):
        SelectiveConn.__init__(self)
        self.conn = conn
        self._recv_pack = b'' # buffer memory

    def get_conn_names(self):
        """
        ** Implementation of ``AbstractConn.get_conn_names``. **
        """
        return {None}

    def send(self, raw, dest=None):
        """
        ** Implementation of ``AbstactConn.send``. **
        """
        if dest is not None:
            raise KeyError('this connection is linked to a single recipient')
        for is_end, pack in anticipate(raw):
            signed_pack = bytes([is_end]) + size2tag(len(pack)) + pack
            self.conn.sendall(signed_pack)

    def recv(self, dest=None):
        """
        ** Implementation of ``AbstraConn.recv``. **
        """
        def gen():
            while True:
                try:
                    yield self.conn.recv(4096)
                except OSError as err:
                    raise ConnectionError('communication seems to be interrupted') from err
        if dest is not None:
            raise KeyError('this connection is linked to a single recipient')
        iter_pack = iter(gen())
        while True:
            # check for len and last pack
            while not self._recv_pack:
                self._recv_pack += next(iter_pack)
            is_end = bool(self._recv_pack[0])
            self._recv_pack = self._recv_pack[1:]
            size, self._recv_pack, _ = tag2size(pack=self._recv_pack, gen=iter_pack)
            while len(self._recv_pack) < size:
                self._recv_pack += next(iter_pack)
            yield self._recv_pack[:size]
            self._recv_pack = self._recv_pack[size:]
            if is_end:
                break

    def close(self):
        """
        ** Implementation of ``AbstractConn.close``. **
        """
        try:
            self.conn.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        self.conn.close()


class BridgeConn(SelectiveConn):
    """
    ** Allows one connection to reach others. **

    Attributes
    ----------
    main_conn : AbstractConn
        The main connection, the one that is intended
        when the target is not specified.
    other_conn : dict
        To each connection name, associate the connection.
    """

    def __init__(self, main_conn, other_conn):
        SelectiveConn.__init__(self)
        self.main_conn = main_conn
        self.other_conn = other_conn

    def get_conn_names(self):
        """
        ** Implementation of ``AbstractConn.get_conn_names``. **
        """
        return {None} | set(self.other_conn)

    def send(self, raw, dest=None):
        """
        ** Implementation of ``AbstactConn.send``. **
        """
        if dest is None:
            return self.main_conn.send(raw)
        try:
            return self.other_conn[dest].send(raw)
        except KeyError as err:
            raise KeyError(
                f"only the recipients {list(self.other_conn.keys())} exist, not '{dest}'"
            ) from err

    def recv(self, dest=None):
        """
        ** Implementation of ``AbstraConn.recv``. **
        """
        if dest is None:
            yield from self.main_conn.recv()
        elif dest in self.other_conn:
            yield from self.other_conn[dest].recv()
        else:
            raise KeyError(
                f"only the recipients {list(self.other_conn)} exist, not '{dest}'"
            )

    def close(self):
        """
        ** Implementation of ``AbstractConn.close``. **
        """
        self.main_conn.close()
        for conn in self.other_conn.values():
            conn.close()


class ServerConn(SelectiveConn):
    """
    ** Allows you to bridge the gap between clients. **

    The 'server' class inherits from this class.
    Thus, it is possible to make the different clients connected
    to the same server communicate with each other.

    Attributes
    ----------
    clients_conn : dict
        To each client identifier, associates the connection of this client.
    """

    def __init__(self):
        SelectiveConn.__init__(self)
        self.clients_conn = {}

    def _conn_expansion(self, abstract_conn):
        """
        ** Memorizes and expands a connection. **

        The server calls this method when the handler is created.
        This is to extend the *abstract_conn* connection so that
        the handler can have access to other clients.

        Parameters
        ----------
        abstract_conn : AbstractConn
            A simple and minimalist connection that
            allows you to get in touch with a single client.

        Returns
        -------
        extended_abstract_conn : AbstractConn
            A connection that allows of course to continue to communicate with the customers.
            That is, it has a similar behavior to the *abstract_conn* argument,
            but it also allows to address the connections of other clients.
        """
        self.clients_conn['id client'] = abstract_conn
        return BridgeConn(main_conn=abstract_conn, other_conn=self.clients_conn)

    def get_conn_names(self):
        """
        ** Implementation of ``AbstractConn.get_conn_names``. **
        """
        return {None} | set(self.clients_conn)

    def send(self, raw, dest=None):
        """
        ** Implementation of ``AbstactConn.send``. **
        """
        if dest is None:
            raise KeyError(
                'there is no default receiver, you have to specify: '
                f'{list(self.clients_conn)}'
            )
        try:
            return self.clients_conn[dest].send(raw)
        except KeyError as err:
            raise KeyError(
                f"only the recipients {list(self.clients_conn)} exist, not '{dest}'"
            ) from err

    def recv(self, dest=None):
        """
        ** Implementation of ``AbstraConn.recv``. **
        """
        if dest is None:
            raise KeyError(
                'there is no default receiver, you have to specify: '
                f'{list(self.clients_conn.keys())}'
            )
        if dest not in self.clients_conn:
            raise KeyError(
                f"only the recipients {list(self.clients_conn)} exist, not '{dest}'"
            )
        yield from self.clients_conn[dest].recv()

    def close(self):
        """
        ** Implementation of ``AbstractConn.close``. **
        """
        for conn in self.clients_conn.values():
            conn.close()
