"""Use pika with the Gevent IOLoop."""

import functools
import logging
import os
import threading
import weakref

try:
    import queue
except ImportError:  # Python <= v2.7
    import Queue as queue

import gevent
import gevent.hub
import gevent.socket

import pika.compat
from pika.adapters.base_connection import BaseConnection
from pika.adapters.utils.io_services_utils import check_callback_arg
from pika.adapters.utils.nbio_interface import (
    AbstractIOReference,
    AbstractIOServices,
)
from pika.adapters.utils.selector_ioloop_adapter import (
    AbstractSelectorIOLoop,
    SelectorIOServicesAdapter,
)

LOGGER = logging.getLogger(__name__)


class GeventConnection(BaseConnection):
    """Implementation of pika's ``BaseConnection``.

    An async selector-based connection which integrates with Gevent.
    """

    def __init__(self,
                 parameters=None,
                 on_open_callback=None,
                 on_open_error_callback=None,
                 on_close_callback=None,
                 custom_ioloop=None,
                 internal_connection_workflow=True):
        """Create a new GeventConnection instance and connect to RabbitMQ on
        Gevent's event-loop.

        :param pika.connection.Parameters|None parameters: The connection
            parameters
        :param callable|None on_open_callback: The method to call when the
            connection is open
        :param callable|None on_open_error_callback: Called if the connection
            can't be established or connection establishment is interrupted by
            `Connection.close()`:
            on_open_error_callback(Connection, exception)
        :param callable|None on_close_callback: Called when a previously fully
            open connection is closed:
            `on_close_callback(Connection, exception)`, where `exception` is
            either an instance of `exceptions.ConnectionClosed` if closed by
            user or broker or exception of another type that describes the
            cause of connection failure
        :param gevent._interfaces.ILoop|nbio_interface.AbstractIOServices|None
            custom_ioloop: Use a custom Gevent ILoop.
        :param bool internal_connection_workflow: True for autonomous connection
            establishment which is default; False for externally-managed
            connection workflow via the `create_connection()` factory
        """
        if pika.compat.ON_WINDOWS:
            raise RuntimeError('GeventConnection is not supported on Windows.')

        custom_ioloop = (custom_ioloop or
                         _GeventSelectorIOLoop(gevent.get_hub()))

        if isinstance(custom_ioloop, AbstractIOServices):
            nbio = custom_ioloop
        else:
            nbio = _GeventSelectorIOServicesAdapter(custom_ioloop)

        super(GeventConnection, self).__init__(
            parameters,
            on_open_callback,
            on_open_error_callback,
            on_close_callback,
            nbio,
            internal_connection_workflow=internal_connection_workflow)

    @classmethod
    def create_connection(cls,
                          connection_configs,
                          on_done,
                          custom_ioloop=None,
                          workflow=None):
        """Implement
        :py:classmethod::`pika.adapters.BaseConnection.create_connection()`.
        """
        custom_ioloop = (custom_ioloop or
                         _GeventSelectorIOLoop(gevent.get_hub()))

        nbio = _GeventSelectorIOServicesAdapter(custom_ioloop)

        def connection_factory(params):
            """Connection factory."""
            if params is None:
                raise ValueError('Expected pika.connection.Parameters '
                                 'instance, but got None in params arg.')
            return cls(parameters=params,
                       custom_ioloop=nbio,
                       internal_connection_workflow=False)

        return cls._start_connection_workflow(
            connection_configs=connection_configs,
            connection_factory=connection_factory,
            nbio=nbio,
            workflow=workflow,
            on_done=on_done)


class _TSafeCallbackQueue(object):
    """Dispatch callbacks from any thread to be executed in the main thread
    efficiently with IO events.
    """

    def __init__(self):
        """
        :param _GeventSelectorIOLoop loop: IO loop to add callbacks to.
        """
        # Thread-safe, blocking queue.
        self._queue = queue.Queue()
        # PIPE to trigger an event when the queue is ready.
        self._read_fd, self._write_fd = os.pipe()
        # Lock around writes to the PIPE in case some platform/implementation
        # requires this.
        self._write_lock = threading.RLock()

    @property
    def fd(self):
        """The file-descriptor to register for READ events in the IO loop."""
        return self._read_fd

    def add_callback_threadsafe(self, callback):
        """Add an item to the queue from any thread. The configured handler
        will be invoked with the item in the main thread.

        :param item: Object to add to the queue.
        """
        self._queue.put(callback)
        with self._write_lock:
            # The value written is not important.
            os.write(self._write_fd, b'\xFF')

    def run_next_callback(self):
        """Invoke the next callback from the queue.

        MUST run in the main thread. If no callback was added to the queue,
        this will block the IO loop.

        Performs a blocking READ on the pipe so must only be called when the
        pipe is ready for reading.
        """
        try:
            callback = self._queue.get_nowait()
        except queue.Empty:
            # Should never happen.
            LOGGER.warning("Callback queue was empty.")
        else:
            # Read the byte from the pipe so the event doesn't re-fire.
            os.read(self._read_fd, 1)
            callback()


class _GeventSelectorIOLoop(AbstractSelectorIOLoop):
    """Implementation of `AbstractSelectorIOLoop` using the Gevent event loop.

    Required by implementations of `SelectorIOServicesAdapter`.
    """
    # Gevent's READ and WRITE masks are defined as 1 and 2 respectively. No
    # ERROR mask is defined.
    # See http://www.gevent.org/api/gevent.hub.html#gevent._interfaces.ILoop.io
    READ = 1
    WRITE = 2
    ERROR = 0

    def __init__(self, gevent_hub=None):
        """
        :param gevent._interfaces.ILoop gevent_loop:
        """
        self._hub = gevent_hub or gevent.get_hub()
        self._io_watchers_by_fd = {}
        # Used to start/stop the loop.
        self._waiter = gevent.hub.Waiter()

        # For adding callbacks from other threads. See `add_callback(..)`.
        self._callback_queue = _TSafeCallbackQueue()

        def run_callback_in_main_thread(fd, events):
            """Swallow the fd and events arguments."""
            del fd
            del events
            self._callback_queue.run_next_callback()

        self.add_handler(self._callback_queue.fd, run_callback_in_main_thread,
                         self.READ)

    def close(self):
        """Release the loop's resources."""
        self._hub.loop.destroy()
        self._hub = None

    def start(self):
        """Run the I/O loop. It will loop until requested to exit. See `stop()`.
        """
        LOGGER.debug("Passing control to Gevent's IOLoop")
        self._waiter.get()  # Block until 'stop()' is called.

        LOGGER.debug("Control was passed back from Gevent's IOLoop")
        self._waiter.clear()

    def stop(self):
        """Request exit from the ioloop. The loop is NOT guaranteed to
        stop before this method returns.

        To invoke `stop()` safely from a thread other than this IOLoop's thread,
        call it via `add_callback_threadsafe`; e.g.,

            `ioloop.add_callback(ioloop.stop)`
        """
        self._waiter.switch(None)

    def add_callback(self, callback):
        """Requests a call to the given function as soon as possible in the
        context of this IOLoop's thread.

        NOTE: This is the only thread-safe method in IOLoop. All other
        manipulations of IOLoop must be performed from the IOLoop's thread.

        For example, a thread may request a call to the `stop` method of an
        ioloop that is running in a different thread via
        `ioloop.add_callback_threadsafe(ioloop.stop)`

        :param callable callback: The callback method
        """
        if gevent.get_hub() == self._hub:
            # We're in the main thread; just add the callback.
            LOGGER.debug("Adding callback from main thread")
            self._hub.loop.run_callback(callback)
        else:
            # This isn't the main thread and Gevent's hub/loop don't provide
            # any thread-safety so enqueue the callback for it to be registered
            # in the main thread.
            LOGGER.debug("Adding callback from another thread")
            callback = functools.partial(self._hub.loop.run_callback, callback)
            self._callback_queue.add_callback_threadsafe(callback)

    def call_later(self, delay, callback):
        """Add the callback to the IOLoop timer to be called after delay seconds
        from the time of call on best-effort basis. Returns a handle to the
        timeout.

        :param float delay: The number of seconds to wait to call callback
        :param callable callback: The callback method
        :returns: handle to the created timeout that may be passed to
            `remove_timeout()`
        :rtype: object
        """
        timer = self._hub.loop.timer(delay)
        timer.start(callback)
        return timer

    def remove_timeout(self, timeout_handle):
        """Remove a timeout

        :param timeout_handle: Handle of timeout to remove
        """
        timeout_handle.close()

    def add_handler(self, fd, handler, events):
        """Start watching the given file descriptor for events

        :param int fd: The file descriptor
        :param callable handler: When requested event(s) occur,
            `handler(fd, events)` will be called.
        :param int events: The event mask (READ|WRITE)
        """
        io_watcher = self._hub.loop.io(fd, events)
        self._io_watchers_by_fd[fd] = io_watcher
        io_watcher.start(handler, fd, events)

    def update_handler(self, fd, events):
        """Change the events being watched for.

        :param int fd: The file descriptor
        :param int events: The new event mask (READ|WRITE)
        """
        io_watcher = self._io_watchers_by_fd[fd]
        # Save callback from the original watcher. The close the old watcher
        # and create a new one using the saved callback and the new events.
        callback = io_watcher.callback
        io_watcher.close()
        del self._io_watchers_by_fd[fd]
        self.add_handler(fd, callback, events)

    def remove_handler(self, fd):
        """Stop watching the given file descriptor for events

        :param int fd: The file descriptor
        """
        io_watcher = self._io_watchers_by_fd[fd]
        io_watcher.close()
        del self._io_watchers_by_fd[fd]


class _GeventSelectorIOServicesAdapter(SelectorIOServicesAdapter):
    """SelectorIOServicesAdapter implementation using Gevent's DNS resolver."""

    def getaddrinfo(self,
                    host,
                    port,
                    on_done,
                    family=0,
                    socktype=0,
                    proto=0,
                    flags=0):
        """Implement :py:meth:`.nbio_interface.AbstractIOServices.getaddrinfo()`.
        """
        resolver = _GeventAddressResolver(native_loop=self._loop,
                                          host=host,
                                          port=port,
                                          family=family,
                                          socktype=socktype,
                                          proto=proto,
                                          flags=flags,
                                          on_done=on_done)
        resolver.start()
        # Return needs an implementation of `AbstractIOReference`.
        return _GeventIOLoopIOHandle(resolver)


class _GeventIOLoopIOHandle(AbstractIOReference):
    """Implement `AbstractIOReference`.

    Only used to wrap the _GeventAddressResolver.
    """

    def __init__(self, subject):
        """
        :param subject: subject of the reference containing a `cancel()` method
        """
        self._cancel = subject.cancel

    def cancel(self):
        """Cancel pending operation

        :returns: False if was already done or cancelled; True otherwise
        :rtype: bool
        """
        return self._cancel()


class _GeventAddressResolver(object):
    """Performs getaddrinfo asynchronously Gevent's configured resolver in a
    separate greenlet and invoking the provided callback with the result.

    See: http://www.gevent.org/dns.html
    """
    __slots__ = (
        '_loop',
        '_on_done',
        '_greenlet',
        # getaddrinfo(..) args:
        '_ga_host',
        '_ga_port',
        '_ga_family',
        '_ga_socktype',
        '_ga_proto',
        '_ga_flags')

    def __init__(self, native_loop, host, port, family, socktype, proto, flags,
                 on_done):
        """Initialize the `_GeventAddressResolver`.

        :param AbstractSelectorIOLoop native_loop:
        :param host: `see socket.getaddrinfo()`
        :param port: `see socket.getaddrinfo()`
        :param family: `see socket.getaddrinfo()`
        :param socktype: `see socket.getaddrinfo()`
        :param proto: `see socket.getaddrinfo()`
        :param flags: `see socket.getaddrinfo()`
        :param on_done: on_done(records|BaseException) callback for reporting
            result from the given I/O loop. The single arg will be either an
            exception object (check for `BaseException`) in case of failure or
            the result returned by `socket.getaddrinfo()`.
        """
        check_callback_arg(on_done, 'on_done')

        self._loop = native_loop
        self._on_done = on_done
        # Reference to the greenlet performing `getaddrinfo`.
        self._greenlet = None
        # getaddrinfo(..) args.
        self._ga_host = host
        self._ga_port = port
        self._ga_family = family
        self._ga_socktype = socktype
        self._ga_proto = proto
        self._ga_flags = flags

    def start(self):
        """Start an asynchronous getaddrinfo invocation."""
        if self._greenlet is None:
            self._greenlet = gevent.spawn_raw(self._resolve)
        else:
            LOGGER.warning("_GeventAddressResolver already started")

    def cancel(self):
        """Cancel the pending resolver."""
        changed = False

        if self._greenlet is not None:
            changed = True
            self._stop_greenlet()

        self._cleanup()
        return changed

    def _cleanup(self):
        """Stop the resolver and release any resources."""
        self._stop_greenlet()
        self._loop = None
        self._on_done = None

    def _stop_greenlet(self):
        """Stop the greenlet performing getaddrinfo if running.

        Otherwise, this is a no-op.
        """
        if self._greenlet is not None:
            gevent.kill(self._greenlet)
            self._greenlet = None

    def _resolve(self):
        """Call `getaddrinfo()` and return result via user's callback
        function on the configured IO loop.
        """
        try:
            # NOTE(JG): Can't use kwargs with getaddrinfo on Python <= v2.7.
            result = gevent.socket.getaddrinfo(self._ga_host, self._ga_port,
                                               self._ga_family,
                                               self._ga_socktype,
                                               self._ga_proto, self._ga_flags)
        except Exception as exc:  # pylint: disable=broad-except
            LOGGER.error('Address resolution failed: %r', exc)
            result = exc

        callback = functools.partial(self._dispatch_callback, result)
        self._loop.add_callback(callback)

    def _dispatch_callback(self, result):
        """Invoke the configured completion callback and any subsequent cleanup.

        :param result: result from getaddrinfo, or the exception if raised.
        """
        try:
            LOGGER.debug(
                'Invoking async getaddrinfo() completion callback; host=%r',
                self._ga_host)
            self._on_done(result)
        finally:
            self._cleanup()