# -*- coding: utf-8 -*-
import logging
import socket
from time import time
from itertools import repeat
from contextlib import contextmanager

import amqp
import gevent
from gevent.queue import Queue
from gevent.event import Event

from base import safe_repr, AmqpRpcAbstract, safe_spawn
from consumer import FailsafeAmqpConsumer
from exception import WrongRequest, IgnoreRequest
from publisher import FailsafeAmqpPublisher
from serializer import RpcJsonSerializer


class _ConsumeLimit(object):
    def __init__(self, tag, max_calls=None, until=None):
        self.tag = tag
        self.max_calls = max_calls
        self.until = until
        self.dying = False
        self.__message_ticks = Queue()

    def __repr__(self):
        if self.until is None:
            remained = None
        else:
            remained = '{0:.4} sec'.format(self.until - time())
        dying = 'DYING ' if self.dying else ''
        return ('<ConsumeLimit id={0:x} {4}tag={1} max_calls: {2}, '
                'remained: {3}>').format(
                    id(self), self.tag, self.max_calls, remained, dying)

    @property
    def overlimited(self):
        if self.dying:
            return True
        if self.max_calls is not None and self.max_calls <= 0:
            return True
        if self.until is not None and self.until <= time():
            return True
        return False

    def tick(self):
        if self.max_calls is not None:
            self.max_calls -= 1
        self.__message_ticks.put(None)
        return self.overlimited

    def wait_until_overlimit(self):
        timer = None
        count = 0
        try:
            if self.until is not None:
                timer = gevent.Timeout(self.until - time())
                timer.start()
            if self.max_calls is not None:
                iterator = xrange(self.max_calls)
            else:
                iterator = repeat(None)
            for _ in iterator:
                self.__message_ticks.get(block=True)
                count += 1
        except gevent.Timeout as e:
            if e is not timer:
                raise
        finally:
            self.dying = True
            if timer is not None:
                timer.cancel()
        return count


class ConsumeObserver(object):
    def __init__(self):
        # {tag::str: limits::ConsumeLimit}
        self._active_by_tag = {}
        # {queue::str: event::Event}
        self._mutex_by_queue = {}
        # {queue::str: count::int}
        self._mutex_waiters_by_queue = {}

    class _ConsumeLimitBuilder(object):
        def __init__(self, observer):
            self._observer = observer

        def __call__(self, tag, max_calls=None, until=None):
            limit = _ConsumeLimit(tag, max_calls, until)
            self._observer._active_by_tag[tag] = limit
            return limit

    @contextmanager
    def creating_for_queue(self, queue):
        if queue not in self._mutex_by_queue:
            self._mutex_by_queue[queue] = Event()
        event = self._mutex_by_queue[queue]
        self._mutex_waiters_by_queue.setdefault(queue, 1)
        try:
            yield self._ConsumeLimitBuilder(self)
        finally:
            event.set()
            event.clear()
            self._mutex_waiters_by_queue[queue] -= 1
            assert self._mutex_waiters_by_queue[queue] >= 0
            if self._mutex_waiters_by_queue[queue] == 0:
                del self._mutex_by_queue[queue]
                del self._mutex_waiters_by_queue[queue]

    def by_tag(self, tag, queue_hint=None, timeout=2):
        try:
            return self._active_by_tag[tag]
        except KeyError:
            if queue_hint not in self._mutex_by_queue:
                raise
            # We're exactly in the race condition
            event = self._mutex_by_queue[queue_hint]
            while timeout > 0:
                iteration_start = time()
                if event.wait(timeout):
                    try:
                        return self._active_by_tag[tag]
                    except KeyError:
                        pass
                timeout -= time() - iteration_start
            raise KeyError(tag)

    def destroyed(self, consume):
        if self._active_by_tag.get(consume.tag) is consume:
            del self._active_by_tag[consume.tag]


class AmqpRpcServer(AmqpRpcAbstract):
    """Server class for AMQP RPC.

    :param service_name: Namespace of your RPC service.
    :type service_name: str
    :param serializer: Class for generating AMQP messages.
        JSON messages generated by default.
    :type serializer: ``AbstractRpcSerializer``
    :param amqp_host: IP address or hostname of AMQP server.
        By default ``127.0.0.1``.
    :type amqp_host: str
    :param amqp_port: TCP port of AMQP server. By default 5672.
    :type amqp_port: int
    :param amqp_user: Username for authentication in AMQP.
        By default ``guest``.
    :type amqp_user: str
    :param amqp_password: Password for authentication in AMQP.
        By default ``guest``.
    :type amqp_password: str
    :param amqp_vhost: Virtual host for authentication in AMQP.
        By default ``/``.
    :type amqp_host: str
    """
    logger = logging.getLogger('amphora.AmqpRpcServer')
    request_exchange_type = 'direct'

    def __init__(self, service_name, serializer=RpcJsonSerializer,
                 **kwargs):
        from gevent import socket as gevent_socket
        assert socket.socket == gevent_socket.socket, \
            "Program must be monkey-patched by gevent"

        super(AmqpRpcServer, self).__init__()
        self.service_name = service_name
        self.serializer = serializer

        self.functions = {}
        self.request_exchange = 'rpc_{0}_request'.format(
            self.service_name)
        self.request_queue = 'rpc_{0}_request'.format(
            self.service_name)
        self.response_exchange = 'rpc_{0}_response'.format(
            self.service_name)
        self.additional_queue_template = 'rpc_' + service_name + '_queue_{0}'

        self.publisher = FailsafeAmqpPublisher(**kwargs)
        self.publisher.before_start_cycle.add_handler(self._prepare_publisher)

        self.consumer = FailsafeAmqpConsumer(**kwargs)
        self.consumer.before_start_cycle.add_handler(self._prepare_consumer)
        self.consumer.on_message.add_handler(self._on_message)

        self._consume_observer = ConsumeObserver()

    @staticmethod
    def _add_function_decorator(func):
        """Change this function by class inheritance if you
        want to decorate every added function."""
        return func

    def add_function(self, func, name=None):
        """Add function for handling remote call. Can be used as
        common function or as decorator.

        :param func: Remote call handler.
        :type func: callable
        :param name: Remote call handler name.
            If omited then call names the same as handler name.
        :type name: str

        Example:

        .. code-block:: python

           # math_server.py
           from gevent import monkey; monkey.patch_socket()
           from amphora import AmqpRpcServer

           server = AmqpRpcServer('simplemath')

           def add(x, y):
               return x + y
           server.add_function(add)

           def substract(x, y):
               return x - y
           server.add_function(substract, name="sub")

           @server.add_function
           def multiply(x, y):
               return x * y

           @server.add_function('div')
           def divide(x, y):
               return x / y

           server.serve(nowait=False)


           # math_client.py
           from gevent import monkey; monkey.patch_socket()
           from amphora import AmqpRpcClient

           client = AmqpRpcClient('simplemath')
           # TODO: must work without __import__('gevent').sleep(1)
           print client.call.add(2, 3)  # 5
           print client.call.sub(10, 3)  # 7
           print client.call.multiply(5, 5)  # 25
           print client.call.div(64, 16)  # 4

        Added function can raise :py:class:`amphora.IgnoreRequest`
        for re-sending this request to another server.
        """
        if name is None:
            if not callable(func):
                name, func = func, None
            else:
                name = func.__name__
        else:
            name = str(name)

        def add_function_wrapper(func):
            self.functions[name] = self._add_function_decorator(func)
            self.logger.debug('Added function %s as %r', func, name)
            return func

        if func is None:
            return add_function_wrapper
        else:
            add_function_wrapper(func)
            return func

    def prepare_stop(self):
        """Reject any new requests. Use this method
        in couple with :py:meth:`~amphora.AmqpRpcServer.stop`
        when you want to do "warm shutdown"."""
        self.consumer.prepare_stop()

    def receive_from_queue(self, queue, max_calls=None,
                           timeout=None, block=True, block_until_closed=True):
        """Receive new calls from specified queue.

        :py:class:`amphora.AmqpRpcServer` authomatically
        starts receiving from queue ``rpc_{service_name}_request``,
        but you may enable receiving calls from several queues.

        :param queue: AMQP queue name with requests.
        :type queue: str
        :param max_calls: Maximum call count. By default unlimited.
        :type max_calls: int
        :param timeout: Stop receiving from queue after
            ``timeout`` seconds.
        :type timeout: float
        :param block: Block current greenlet until
            maximum calls reached or timed out.
        :type block: bool
        :param block_until_closed: Block current greenlet until
            all requests will be handled and channel will be closed.
            Ignored if ``block`` == ``False``
        :type block_until_closed: bool
        :raises gevent.Timeout: When timed out and maximum calls
            was not reached.
        :raises ValueError: If queue not found.

        If you set ``max_calls`` then `AmqpRpcServer` handles only
        specified calls. After reaching maximum `AmqpRpcServer` cancel
        consuming from queue but does not close channel until all handling
        messages acked or rejected. Same rules applies for ``timeout``.
        You can specify both parameters at once.

        When requests delivered too frequently, AMQP server can send
        some request messages just after `basic_cancel` call. These
        messages will be rejected with requeuing. In order to avoid
        infinite resending messages that can cause denial of service,
        AmqpRpcServer can reject messages with one second delay.

        When ``max_calls`` specified, `basic.qos` with similar
        prefetch count will be sent before start consuming.
        """
        queue = self.additional_queue_template.format(queue)
        if block:
            return self._receive_from_queue(
                queue, max_calls, timeout, block_until_closed)
        else:
            safe_spawn(self._receive_from_queue, queue, max_calls,
                       timeout, False)

    def _receive_from_queue(self, queue, max_calls, timeout,
                            block_until_closed):
        if timeout is None:
            until = None
        else:
            until = time() + timeout

        with self._consume_observer.creating_for_queue(queue) as create:
            tag = self.consumer.start_consume(queue, nowait=False, qos=max_calls)
            limits = create(tag, max_calls=max_calls, until=until)
        self.logger.debug(
            "Start handling requests from queue %s with %r", queue, limits)

        result = limits.wait_until_overlimit()
        self.logger.debug(
            "Stop handling requests from queue %s with %r", queue, limits)

        try:
            self.consumer.stop_consume(tag=tag, nowait=block_until_closed)
        finally:
            self._consume_observer.destroyed(limits)
        return result

    def _prepare_publisher(self, publisher):
        publisher.new_exchange(self.response_exchange, 'direct', nowait=True)

    def _prepare_consumer(self, consumer):
        consumer.new_exchange(
            self.request_exchange, self.request_exchange_type, nowait=True)
        consumer.new_queue(
            self.request_queue, self.request_exchange, nowait=True)
        safe_spawn(
            self._receive_from_queue, self.request_queue, None, None, False)

    def _on_message(self, message):
        safe_spawn(self._on_message_spawned, message)

    def _on_message_spawned(self, message):
        if ('reply_to' not in message.properties or
                'message_id' not in message.properties):
            self.logger.warning(
                "Strange message without proper headers: %s. (queue %s)"
                "Silently acking.", safe_repr(message.body), message.queue)
            message.ack()
            return

        try:
            limits = self._consume_observer.by_tag(
                message.consumer_tag, queue_hint=message.queue)
        except KeyError:
            self.logger.warning(
                "Got request %s from not listened queue %s. Rejecting.",
                safe_repr(message.body), message.queue)
            message.reject()
            return

        if limits.overlimited:
            self.logger.debug(
                "Got request %s from overlimited consume %r "
                "(queue %s). Rejecting.",
                safe_repr(message.body), limits, message.queue)
            message.reject()
            return

        self.logger.debug(
            "Got request message %s", safe_repr(message.body))

        serializer = self.serializer
        try:
            func_name, args, kwargs = serializer.deserialize_call(
                message.body)
            func = self.functions.get(func_name)
            if func is None:
                self.logger.warning("No function %r defined", func_name)
                raise WrongRequest("No function {0} defined".format(func_name))
        except Exception as e:
            # Wrong formatted messages aren't considered as executed calls
            self._send_result(message, serializer.serialize_result(e))
            message.ack()
        else:
            limits.tick()
            safe_spawn(self._run_function, message, func, args, kwargs)

    def _run_function(self, message, func, args, kwargs):
        try:
            result = func(*args, **kwargs)
            body = self.serializer.serialize_result(result)
        except IgnoreRequest:
            self._on_ignore_request(message)
            return
        except Exception as e:
            self.logger.warning(
                "When executed %s(*%r, **%r)", func, args, kwargs,
                exc_info=1)
            body = self.serializer.serialize_result(e)
        self._send_result(message, body)
        message.ack()

    def _on_ignore_request(self, message):
        message.reject()

    def _send_result(self, message, body):
        message_id = message.properties['message_id']
        reply_to = message.properties['reply_to']

        response_msg = amqp.Message(
            body=body,
            message_id=message_id,
            delivery_mode=2)  # PERSISTENT
        self.publisher.publish_message(
            response_msg, self.response_exchange, routing_key=reply_to)


class AmqpRpcFanoutServer(AmqpRpcServer):
    """Fanout server class for AMQP RPC.

    Fanout exchanges in AMQP used for broadcasting. If you
    execute remote function via :py:class:`~amphora.AmqpRpcFanoutClient`
    then request will be handled by every connected
    :py:class:`~amphora.AmqpRpcFanoutServer`.

    One difference from :py:class:`~amphora.AmqpRpcServer` it that
    added functions can raise :py:class:`amphora.IgnoreRequest`.
    When function raises it, server sends "ack" to request message
    but does not send any result or exception back to client.
    """
    request_exchange_type = 'fanout'

    def _on_ignore_request(self, message):
        pass  # simply drop, do not send reject
