""" Internal utils related to Tornado

"""
from __future__ import absolute_import, print_function

import logging
log = logging.getLogger(__name__)

from tornado import gen

@gen.coroutine
def yield_for_all_futures(result):
    """ Converts result into a Future by collapsing any futures inside result.

    If result is a Future we yield until it's done, then if the value inside
    the Future is another Future we yield until it's done as well, and so on.
    """
    while True:

        # This is needed for Tornado >= 4.5 where convert_yielded will no
        # longer raise BadYieldError on None
        if result is None:
            break

        try:
            future = gen.convert_yielded(result)
        except gen.BadYieldError:
            # result is not a yieldable thing, we are done
            break
        else:
            result = yield future

    raise gen.Return(result)

class _AsyncPeriodic(object):
    """Like ioloop.PeriodicCallback except the 'func' can be async and
        return a Future, and we wait for func to finish each time
        before we call it again.  Plain ioloop.PeriodicCallback
        can "pile up" invocations if they are taking too long.

    """
    def __init__(self, func, period, io_loop):
        self._func = func
        self._loop = io_loop
        self._period = period
        self._started = False
        self._stopped = False

    # this is like gen.sleep but uses our IOLoop instead of the
    # current IOLoop
    def sleep(self):
        f = gen.Future()
        self._loop.call_later(self._period / 1000.0, lambda: f.set_result(None))
        return f

    def start(self):
        if self._started:
            raise RuntimeError("called start() twice on _AsyncPeriodic")
        self._started = True
        def invoke():
            # important to start the sleep before starting callback
            # so any initial time spent in callback "counts against"
            # the period.
            sleep_future = self.sleep()
            result = self._func()

            # This is needed for Tornado >= 4.5 where convert_yielded will no
            # longer raise BadYieldError on None
            if result is None:
                return sleep_future

            try:
                callback_future = gen.convert_yielded(result)
            except gen.BadYieldError:
                # result is not a yieldable thing
                return sleep_future
            else:
                return gen.multi([sleep_future, callback_future])
        def on_done(future):
            if not self._stopped:
                self._loop.add_future(invoke(), on_done)
            if future.exception() is not None:
                log.error("Error thrown from periodic callback: %r", future.exception())
        self._loop.add_future(self.sleep(), on_done)

    def stop(self):
        self._stopped = True

class _CallbackGroup(object):
    """ A collection of callbacks added to a Tornado IOLoop that we may
    want to remove as a group. """

    def __init__(self, io_loop=None):
        if io_loop is None:
            raise ValueError("must provide an io loop")
        self._loop = io_loop
        # dicts from callback to remove callable. These are
        # separate only because it's allowed to add the same
        # callback as multiple kinds of callback at once.
        self._next_tick_callbacks = {}
        self._timeout_callbacks = {}
        self._periodic_callbacks = {}

    def remove_all_callbacks(self):
        """ Removes all registered callbacks."""
        for cb in list(self._next_tick_callbacks.keys()):
            self.remove_next_tick_callback(cb)
        for cb in list(self._timeout_callbacks.keys()):
            self.remove_timeout_callback(cb)
        for cb in list(self._periodic_callbacks.keys()):
            self.remove_periodic_callback(cb)

    def _error_on_double_remove(self, callback, callbacks):
        if callback not in callbacks:
            raise ValueError("Removing a callback twice (or after it's already been run)")

    def _remover(self, callback, callbacks, cleanup):
        self._error_on_double_remove(callback, callbacks)
        del callbacks[callback]
        if cleanup is not None:
            cleanup(callback)

    def _wrap_next_tick(self, callback, cleanup):
        # this 'removed' flag is a hack because Tornado has no way
        # to remove a "next tick" callback added with
        # IOLoop.add_callback. So instead we make our wrapper skip
        # invoking the callback.
        handle = { 'removed' : False }
        def wrapper(*args, **kwargs):
            was_removed = handle['removed']
            if not was_removed:
                self.remove_next_tick_callback(callback)
                return callback(*args, **kwargs)
            else:
                return None
        wrapper.handle = handle
        return wrapper

    def add_next_tick_callback(self, callback, cleanup=None):
        """ Adds a callback to be run on the next tick.
        Returns a callable that removes the callback if called."""
        if callback in self._next_tick_callbacks:
            raise ValueError("Next-tick callback added twice")
        wrapper = self._wrap_next_tick(callback, cleanup)
        self._loop.add_callback(wrapper)
        def remover():
            wrapper.handle['removed'] = True
            self._remover(callback, self._next_tick_callbacks, cleanup)
        self._next_tick_callbacks[callback] = remover
        return remover

    def _remove(self, callback, callbacks):
        self._error_on_double_remove(callback, callbacks)
        callbacks[callback]()

    def remove_next_tick_callback(self, callback):
        """ Removes a callback added with add_next_tick_callback."""
        self._remove(callback, self._next_tick_callbacks)

    def _wrap_timeout(self, callback):
        def wrapper(*args, **kwargs):
            self.remove_timeout_callback(callback)
            return callback(*args, **kwargs)
        return wrapper

    def add_timeout_callback(self, callback, timeout_milliseconds, cleanup=None):
        """ Adds a callback to be run once after timeout_milliseconds.
        Returns a callable that removes the callback if called."""
        if callback in self._timeout_callbacks:
            raise ValueError("Callback added as a timeout twice")
        handle = self._loop.call_later(timeout_milliseconds / 1000.0,
                                       self._wrap_timeout(callback))
        def remover():
            self._loop.remove_timeout(handle)
            self._remover(callback, self._timeout_callbacks, cleanup)
        self._timeout_callbacks[callback] = remover
        return remover

    def remove_timeout_callback(self, callback):
        """ Removes a callback added with add_timeout_callback, before it runs."""
        self._remove(callback, self._timeout_callbacks)

    def add_periodic_callback(self, callback, period_milliseconds, cleanup=None):
        """ Adds a callback to be run every period_milliseconds until it is removed."""
        if callback in self._periodic_callbacks:
            raise ValueError("Callback added as a periodic callback twice")
        cb = _AsyncPeriodic(
            callback, period_milliseconds, io_loop=self._loop
        )
        def remover():
            cb.stop()
            self._remover(callback, self._periodic_callbacks, cleanup)
        self._periodic_callbacks[callback] = remover
        cb.start()
        return remover

    def remove_periodic_callback(self, callback):
        """ Removes a callback added with add_periodic_callback."""
        self._remove(callback, self._periodic_callbacks)
