"""
This module provides a remote procedure call (RPC) mechanism over sockets
between conventional computers (PCs) running Python. It strives to be
transparent and uses ``artiq.protocols.pyon`` internally so that e.g. Numpy
arrays can be easily used.

Note that the server operates on copies of objects provided by the client,
and modifications to mutable types are not written back. For example, if the
client passes a list as a parameter of an RPC method, and that method
``append()s`` an element to the list, the element is not appended to the
client's list.
"""

import socket
import asyncio
import traceback
import threading
import time
import logging
import inspect
from operator import itemgetter

from artiq.protocols import pyon
from artiq.protocols.asyncio_server import AsyncioServer as _AsyncioServer


logger = logging.getLogger(__name__)


class AutoTarget:
    """Use this as target value in clients for them to automatically connect
    to the target exposed by the server. Servers must have only one target."""
    pass


class RemoteError(Exception):
    """Raised when a RPC failed or raised an exception on the remote (server)
    side."""
    pass


class IncompatibleServer(Exception):
    """Raised by the client when attempting to connect to a server that does
    not have the expected target."""
    pass


_init_string = b"ARTIQ pc_rpc\n"


def _validate_target_name(target_name, target_names):
    if target_name is AutoTarget:
        if len(target_names) > 1:
            raise ValueError("Server has multiple targets: " +
                             " ".join(sorted(target_names)))
        else:
            target_name = target_names[0]
    elif target_name not in target_names:
        raise IncompatibleServer(
            "valid target name(s): " +
                " ".join(sorted(target_names)))
    return target_name


class Client:
    """This class proxies the methods available on the server so that they
    can be used as if they were local methods.

    For example, if the server provides method ``foo``, and ``c`` is a local
    ``Client`` object, then the method can be called as: ::

        result = c.foo(param1, param2)

    The parameters and the result are automatically transferred with the
    server.

    Only methods are supported. Attributes must be accessed by providing and
    using "get" and/or "set" methods on the server side.

    At object initialization, the connection to the remote server is
    automatically attempted. The user must call ``close_rpc`` to
    free resources properly after initialization completes successfully.

    :param host: Identifier of the server. The string can represent a
        hostname or a IPv4 or IPv6 address (see
        ``socket.create_connection`` in the Python standard library).
    :param port: TCP port to use.
    :param target_name: Target name to select. ``IncompatibleServer`` is
        raised if the target does not exist.
        Use ``AutoTarget`` for automatic selection if the server has only one
        target.
        Use ``None`` to skip selecting a target. The list of targets can then
        be retrieved using ``get_rpc_id`` and then one can be selected later
        using ``select_rpc_target``.
    """
    def __init__(self, host, port, target_name=AutoTarget):
        self.__socket = socket.create_connection((host, port))

        try:
            self.__socket.sendall(_init_string)

            server_identification = self.__recv()
            self.__target_names = server_identification["targets"]
            self.__description = server_identification["description"]
            if target_name is not None:
                self.select_rpc_target(target_name)
        except:
            self.__socket.close()
            raise

    def select_rpc_target(self, target_name):
        """Selects a RPC target by name. This function should be called
        exactly once if the object was created with ``target_name=None``."""
        target_name = _validate_target_name(target_name, self.__target_names)
        self.__socket.sendall((target_name + "\n").encode())

    def get_rpc_id(self):
        """Returns a tuple (target_names, description) containing the
        identification information of the server."""
        return (self.__target_names, self.__description)

    def close_rpc(self):
        """Closes the connection to the RPC server.

        No further method calls should be done after this method is called.
        """
        self.__socket.close()

    def __send(self, obj):
        line = pyon.encode(obj) + "\n"
        self.__socket.sendall(line.encode())

    def __recv(self):
        buf = self.__socket.recv(4096).decode()
        while "\n" not in buf:
            more = self.__socket.recv(4096)
            if not more:
                break
            buf += more.decode()
        return pyon.decode(buf)

    def __do_action(self, action):
        self.__send(action)

        obj = self.__recv()
        if obj["status"] == "ok":
            return obj["ret"]
        elif obj["status"] == "failed":
            raise RemoteError(obj["message"])
        else:
            raise ValueError

    def __do_rpc(self, name, args, kwargs):
        obj = {"action": "call", "name": name, "args": args, "kwargs": kwargs}
        return self.__do_action(obj)

    def get_rpc_method_list(self):
        obj = {"action": "get_rpc_method_list"}
        return self.__do_action(obj)

    def __getattr__(self, name):
        def proxy(*args, **kwargs):
            return self.__do_rpc(name, args, kwargs)
        return proxy


class AsyncioClient:
    """This class is similar to :class:`artiq.protocols.pc_rpc.Client`, but
    uses ``asyncio`` instead of blocking calls.

    All RPC methods are coroutines.

    Concurrent access from different asyncio tasks is supported; all calls
    use a single lock.
    """
    def __init__(self):
        self.__lock = asyncio.Lock()
        self.__reader = None
        self.__writer = None
        self.__target_names = None
        self.__description = None

    async def connect_rpc(self, host, port, target_name):
        """Connects to the server. This cannot be done in __init__ because
        this method is a coroutine. See ``Client`` for a description of the
        parameters."""
        self.__reader, self.__writer = \
            await asyncio.open_connection(host, port)
        try:
            self.__writer.write(_init_string)
            server_identification = await self.__recv()
            self.__target_names = server_identification["targets"]
            self.__description = server_identification["description"]
            if target_name is not None:
                self.select_rpc_target(target_name)
        except:
            self.close_rpc()
            raise

    def select_rpc_target(self, target_name):
        """Selects a RPC target by name. This function should be called
        exactly once if the connection was created with ``target_name=None``.
        """
        target_name = _validate_target_name(target_name, self.__target_names)
        self.__writer.write((target_name + "\n").encode())

    def get_rpc_id(self):
        """Returns a tuple (target_names, description) containing the
        identification information of the server."""
        return (self.__target_names, self.__description)

    def close_rpc(self):
        """Closes the connection to the RPC server.

        No further method calls should be done after this method is called.
        """
        self.__writer.close()
        self.__reader = None
        self.__writer = None
        self.__target_names = None
        self.__description = None

    def __send(self, obj):
        line = pyon.encode(obj) + "\n"
        self.__writer.write(line.encode())

    async def __recv(self):
        line = await self.__reader.readline()
        return pyon.decode(line.decode())

    async def __do_rpc(self, name, args, kwargs):
        await self.__lock.acquire()
        try:
            obj = {"action": "call", "name": name,
                   "args": args, "kwargs": kwargs}
            self.__send(obj)

            obj = await self.__recv()
            if obj["status"] == "ok":
                return obj["ret"]
            elif obj["status"] == "failed":
                raise RemoteError(obj["message"])
            else:
                raise ValueError
        finally:
            self.__lock.release()

    def __getattr__(self, name):
        async def proxy(*args, **kwargs):
            res = await self.__do_rpc(name, args, kwargs)
            return res
        return proxy


class BestEffortClient:
    """This class is similar to :class:`artiq.protocols.pc_rpc.Client`, but
    network errors are suppressed and connections are retried in the
    background.

    RPC calls that failed because of network errors return ``None``. Other RPC
    calls are blocking and return the correct value.

    :param firstcon_timeout: Timeout to use during the first (blocking)
        connection attempt at object initialization.
    :param retry: Amount of time to wait between retries when reconnecting
        in the background.
    """
    def __init__(self, host, port, target_name,
                 firstcon_timeout=0.5, retry=5.0):
        self.__host = host
        self.__port = port
        self.__target_name = target_name
        self.__retry = retry

        self.__conretry_terminate = False
        self.__socket = None
        try:
            self.__coninit(firstcon_timeout)
        except:
            logger.warning("first connection attempt to %s:%d[%s] failed, "
                           "retrying in the background",
                           self.__host, self.__port, self.__target_name,
                           exc_info=True)
            self.__start_conretry()
        else:
            self.__conretry_thread = None

    def __coninit(self, timeout):
        if timeout is None:
            self.__socket = socket.create_connection(
                (self.__host, self.__port))
        else:
            self.__socket = socket.create_connection(
                (self.__host, self.__port), timeout)
            self.__socket.settimeout(None)
        self.__socket.sendall(_init_string)
        server_identification = self.__recv()
        target_name = _validate_target_name(self.__target_name,
                                            server_identification["targets"])
        self.__socket.sendall((target_name + "\n").encode())

    def __start_conretry(self):
        self.__conretry_thread = threading.Thread(target=self.__conretry)
        self.__conretry_thread.start()

    def __conretry(self):
        while True:
            try:
                self.__coninit(None)
            except:
                if self.__conretry_terminate:
                    break
                time.sleep(self.__retry)
            else:
                break
        if not self.__conretry_terminate:
            logger.warning("connection to %s:%d[%s] established in "
                           "the background",
                           self.__host, self.__port, self.__target_name)
        if self.__conretry_terminate and self.__socket is not None:
            self.__socket.close()
        # must be after __socket.close() to avoid race condition
        self.__conretry_thread = None

    def close_rpc(self):
        """Closes the connection to the RPC server.

        No further method calls should be done after this method is called.
        """
        if self.__conretry_thread is None:
            if self.__socket is not None:
                self.__socket.close()
        else:
            # Let the thread complete I/O and then do the socket closing.
            # Python fails to provide a way to cancel threads...
            self.__conretry_terminate = True

    def __send(self, obj):
        line = pyon.encode(obj) + "\n"
        self.__socket.sendall(line.encode())

    def __recv(self):
        buf = self.__socket.recv(4096).decode()
        while "\n" not in buf:
            more = self.__socket.recv(4096)
            if not more:
                break
            buf += more.decode()
        return pyon.decode(buf)

    def __do_rpc(self, name, args, kwargs):
        if self.__conretry_thread is not None:
            return None

        obj = {"action": "call", "name": name, "args": args, "kwargs": kwargs}
        try:
            self.__send(obj)
            obj = self.__recv()
        except:
            logger.warning("connection failed while attempting "
                           "RPC to %s:%d[%s], re-establishing connection "
                           "in the background",
                           self.__host, self.__port, self.__target_name)
            self.__start_conretry()
            return None
        else:
            if obj["status"] == "ok":
                return obj["ret"]
            elif obj["status"] == "failed":
                raise RemoteError(obj["message"])
            else:
                raise ValueError

    def __getattr__(self, name):
        def proxy(*args, **kwargs):
            return self.__do_rpc(name, args, kwargs)
        return proxy


def _format_arguments(arguments):
    fmtargs = []
    for k, v in sorted(arguments.items(), key=itemgetter(0)):
        fmtargs.append(k + "=" + repr(v))
    if fmtargs:
        return ", ".join(fmtargs)
    else:
        return ""


class _PrettyPrintCall:
    def __init__(self, obj):
        self.obj = obj

    def __str__(self):
        r = self.obj["name"] + "("
        args = ", ".join([repr(a) for a in self.obj["args"]])
        r += args
        kwargs = _format_arguments(self.obj["kwargs"])
        if args and kwargs:
            r += ", "
        r += kwargs
        r += ")"
        return r


class Server(_AsyncioServer):
    """This class creates a TCP server that handles requests coming from
    ``Client`` objects.

    The server is designed using ``asyncio`` so that it can easily support
    multiple connections without the locking issues that arise in
    multi-threaded applications. Multiple connection support is useful even in
    simple cases: it allows new connections to be be accepted even when the
    previous client failed to properly shut down its connection.

    :param targets: A dictionary of objects providing the RPC methods to be
        exposed to the client. Keys are names identifying each object.
        Clients select one of these objects using its name upon connection.
    :param description: An optional human-readable string giving more
        information about the server.
    :param builtin_terminate: If set, the server provides a built-in
        ``terminate`` method that unblocks any tasks waiting on
        ``wait_terminate``. This is useful to handle server termination
        requests from clients.
    """
    def __init__(self, targets, description=None, builtin_terminate=False):
        _AsyncioServer.__init__(self)
        self.targets = targets
        self.description = description
        self.builtin_terminate = builtin_terminate
        if builtin_terminate:
            self._terminate_request = asyncio.Event()

    async def _handle_connection_cr(self, reader, writer):
        try:
            line = await reader.readline()
            if line != _init_string:
                return

            obj = {
                "targets": sorted(self.targets.keys()),
                "description": self.description
            }
            line = pyon.encode(obj) + "\n"
            writer.write(line.encode())
            line = await reader.readline()
            if not line:
                return
            target_name = line.decode()[:-1]
            try:
                target = self.targets[target_name]
            except KeyError:
                return

            while True:
                line = await reader.readline()
                if not line:
                    break
                obj = pyon.decode(line.decode())
                try:
                    if obj["action"] == "get_rpc_method_list":
                        members = inspect.getmembers(target, inspect.ismethod)
                        doc = {
                            "docstring": inspect.getdoc(target),
                            "methods": {}
                        }
                        for name, method in members:
                            if name.startswith("_"):
                                continue
                            method = getattr(target, name)
                            argspec = inspect.getfullargspec(method)
                            doc["methods"][name] = (dict(argspec._asdict()),
                                                    inspect.getdoc(method))
                        if self.builtin_terminate:
                            doc["methods"]["terminate"] = (
                                {
                                    "args": ["self"],
                                    "defaults": None,
                                    "varargs": None,
                                    "varkw": None,
                                    "kwonlyargs": [],
                                    "kwonlydefaults": [],
                                },
                                "Terminate the server.")
                        obj = {"status": "ok", "ret": doc}
                    elif obj["action"] == "call":
                        logger.debug("calling %s", _PrettyPrintCall(obj))
                        if self.builtin_terminate and obj["name"] == "terminate":
                            self._terminate_request.set()
                            obj = {"status": "ok", "ret": None}
                        else:
                            method = getattr(target, obj["name"])
                            ret = method(*obj["args"], **obj["kwargs"])
                            if inspect.iscoroutine(ret):
                                ret = await ret
                            obj = {"status": "ok", "ret": ret}
                    else:
                        raise ValueError("Unknown action: {}"
                                         .format(obj["action"]))
                except Exception:
                    obj = {"status": "failed",
                           "message": traceback.format_exc()}
                line = pyon.encode(obj) + "\n"
                writer.write(line.encode())
        finally:
            writer.close()

    async def wait_terminate(self):
        await self._terminate_request.wait()


def simple_server_loop(targets, host, port, description=None):
    """Runs a server until an exception is raised (e.g. the user hits Ctrl-C)
    or termination is requested by a client.

    See ``Server`` for a description of the parameters.
    """
    loop = asyncio.get_event_loop()
    try:
        server = Server(targets, description, True)
        loop.run_until_complete(server.start(host, port))
        try:
            loop.run_until_complete(server.wait_terminate())
        finally:
            loop.run_until_complete(server.stop())
    finally:
        loop.close()
