# -*- coding: UTF-8 -*-
# pylint: disable=line-too-long
"""
This module provides functionality to support "async steps" (coroutines)
in a step-module with behave. This functionality simplifies to test
frameworks and protocols that make use of `asyncio.coroutines`_ or
provide `asyncio.coroutines`_.

EXAMPLE:

.. code-block:: python

    # -- FILE: features/steps/my_async_steps.py
    # EXAMPLE REQUIRES: Python >= 3.5
    from behave import step
    from behave.api.async_step import async_run_until_complete

    @step('an async coroutine step waits {duration:f} seconds')
    @async_run_until_complete
    async def step_async_step_waits_seconds(context, duration):
        await asyncio.sleep(duration)

.. code-block:: python

    # -- FILE: features/steps/my_async_steps2.py
    # EXAMPLE REQUIRES: Python >= 3.4
    from behave import step
    from behave.api.async_step import async_run_until_complete
    import asyncio

    @step('a tagged-coroutine async step waits {duration:f} seconds')
    @async_run_until_complete
    @asyncio.coroutine
    def step_async_step_waits_seconds2(context, duration):
        yield from asyncio.sleep(duration)


.. requires:: Python 3.5 (or 3.4) or :mod:`asyncio` backport (like :pypi:`trollius`)
.. seealso::
    https://docs.python.org/3/library/asyncio.html

.. _asyncio.coroutines: https://docs.python.org/3/library/asyncio-task.html#coroutines
"""
# pylint: enable=line-too-long

from __future__ import print_function
# -- REQUIRES: Python >= 3.4
# MAYBE BACKPORT: trollius
import functools
from six import string_types
try:
    import asyncio
    has_asyncio = True
except ImportError:
    has_asyncio = False

# -----------------------------------------------------------------------------
# ASYNC STEP DECORATORS:
# -----------------------------------------------------------------------------
def async_run_until_complete(astep_func=None, loop=None, timeout=None,
                             async_context=None, should_close=False):
    """Provides a function decorator for async-steps (coroutines).
    Provides an async event loop and runs the async-step until completion
    (or timeout, if specified).

    .. code-block:: python

        from behave import step
        from behave.api.async_step import async_run_until_complete
        import asyncio

        @step("an async step is executed")
        @async_run_until_complete
        async def astep_impl(context)
            await asycio.sleep(0.1)

        @step("an async step is executed")
        @async_run_until_complete(timeout=1.2)
        async def astep_impl2(context)
            # -- NOTE: Wrapped event loop waits with timeout=1.2 seconds.
            await asycio.sleep(0.3)

    Parameters:
        astep_func: Async step function (coroutine)
        loop (asyncio.EventLoop):   Event loop to use or None.
        timeout (int, float):       Timeout to wait for async-step completion.
        async_context (name):       Async_context name or object to use.
        should_close (bool):        Indicates if event loop should be closed.

    .. note::

        * If :param:`loop` is None, the default event loop will be used
          or a new event loop is created.
        * If :param:`timeout` is provided, the event loop waits only the
          specified time.
        * :param:`async_context` is only used, if :param:`loop` is None.
        * If :param:`async_context` is a name, it will be used to retrieve
          the real async_context object from the context.

    """
    @functools.wraps(astep_func)
    def step_decorator(astep_func, context, *args, **kwargs):
        loop = kwargs.pop("_loop", None)
        timeout = kwargs.pop("_timeout", None)
        async_context = kwargs.pop("_async_context", None)
        should_close = kwargs.pop("_should_close", None)

        if isinstance(loop, string_types):
            loop = getattr(context, loop, None)
        elif async_context:
            if isinstance(async_context, string_types):
                name = async_context
                async_context = use_or_create_async_context(context, name)
                loop = async_context.loop
            else:
                assert isinstance(async_context, AsyncContext)
                loop = async_context.loop
        if loop is None:
            loop = asyncio.get_event_loop() or asyncio.new_event_loop()

        # -- WORKHORSE:
        try:
            if timeout is None:
                loop.run_until_complete(astep_func(context, *args, **kwargs))
            else:
                # MAYBE: loop = asyncio.new_event_loop()
                # MAYBE: should_close = True
                task = loop.create_task(astep_func(context, *args, **kwargs))
                done, pending = loop.run_until_complete(
                    asyncio.wait([task], timeout=timeout))
                assert not pending, "TIMEOUT-OCCURED: timeout=%s" % timeout
                finished_task = done.pop()
                exception = finished_task.exception()
                if exception:
                    raise exception
        finally:
            if loop and should_close:
                # -- MAYBE-AVOID:
                loop.close()

    if astep_func is None:
        # -- CASE: @decorator(timeout=1.2, ...)
        # MAYBE: return functools.partial(step_decorator,
        def wrapped_decorator1(astep_func):
            @functools.wraps(astep_func)
            def wrapped_decorator2(context, *args, **kwargs):
                return step_decorator(astep_func, context, *args,
                                      _loop=loop,
                                      _timeout=timeout,
                                      _async_context=async_context,
                                      _should_close=should_close, **kwargs)
            assert callable(astep_func)
            return wrapped_decorator2
        return wrapped_decorator1
    else:
        # -- CASE: @decorator ... or astep = decorator(astep)
        # MAYBE: return functools.partial(step_decorator, astep_func=astep_func)
        assert callable(astep_func)
        @functools.wraps(astep_func)
        def wrapped_decorator(context, *args, **kwargs):
            return step_decorator(astep_func, context, *args, **kwargs)
        return wrapped_decorator

# -- ALIAS:
run_until_complete = async_run_until_complete

# -----------------------------------------------------------------------------
# ASYNC STEP UTILITY CLASSES:
# -----------------------------------------------------------------------------
class AsyncContext(object):
    # pylint: disable=line-too-long
    """Provides a context object for "async steps" to keep track:

    * which event loop is used
    * which (asyncio) tasks are used or of interest

    .. attribute:: loop
        Event loop object to use.
        If none is provided, the current event-loop is used
        (or a new one is created).

    .. attribute:: tasks
        List of started :mod:`asyncio` tasks (of interest).

    .. attribute:: name

        Optional name of this object (in the behave context).
        If none is provided, :attr:`AsyncContext.default_name` is used.

    .. attribute:: should_close
        Indicates if the :attr:`loop` (event-loop) should be closed or not.

    EXAMPLE:

    .. code-block:: python

        # -- FILE: features/steps/my_async_steps.py
        # REQUIRES: Python 3.5
        from behave import given, when, then, step
        from behave.api.async_step import AsyncContext

        @when('I dispatch an async-call with param "{param}"')
        def step_impl(context, param):
            async_context = getattr(context, "async_context", None)
            if async_context is None:
                async_context = context.async_context = AsyncContext()
            task = async_context.loop.create_task(my_async_func(param))
            async_context.tasks.append(task)

        @then('I wait at most {duration:f} seconds until all async-calls are completed')
        def step_impl(context, duration):
            async_context = context.async_context
            assert async_context.tasks
            done, pending = async_context.loop.run_until_complete(asyncio.wait(
                async_context.tasks, loop=async_context.loop, timeout=duration))
            assert len(pending) == 0

        # -- COROUTINE:
        async def my_async_func(param):
            await asyncio.sleep(0.5)
            return param.upper()
    """
    # pylint: enable=line-too-long
    default_name = "async_context"

    def __init__(self, loop=None, name=None, should_close=False, tasks=None):
        self.loop = loop or asyncio.get_event_loop() or asyncio.new_event_loop()
        self.tasks = tasks or []
        self.name = name or self.default_name
        self.should_close = should_close

    def __del__(self):
        if self.loop and self.should_close:
            self.close()

    def close(self):
        if self.loop and not self.loop.is_closed():
            self.loop.close()
            self.loop = None


# -----------------------------------------------------------------------------
# ASYNC STEP UTILITY FUNCTIONS:
# -----------------------------------------------------------------------------
def use_or_create_async_context(context, name=None, loop=None, **kwargs):
    """Utility function to be used in step implementations to ensure that an
    :class:`AsyncContext` object is stored in the :param:`context` object.

    If no such attribute exists (under the given name),
    a new :class:`AsyncContext` object is created with the provided args.
    Otherwise, the existing context attribute is used.

    EXAMPLE:

    .. code-block:: python

        # -- FILE: features/steps/my_async_steps.py
        # EXAMPLE REQUIRES: Python 3.5
        from behave import when
        from behave.api.async_step import use_or_create_async_context

        @when('I dispatch an async-call with param "{param}"')
        def step_impl(context, param):
            async_context = use_or_create_async_context(context, "async_context")
            task = async_context.loop.create_task(my_async_func(param))
            async_context.tasks.append(task)

        # -- COROUTINE:
        async def my_async_func(param):
            await asyncio.sleep(0.5)
            return param.upper()

    :param context:     Behave context object to use.
    :param name:        Optional name of async-context object (as string or None).
    :param loop:        Optional event_loop object to use for create call.
    :param kwargs:      Optional :class:`AsyncContext` params for create call.
    :return: :class:`AsyncContext` object from the param:`context`.
    """
    if name is None:
        name = AsyncContext.default_name
    async_context = getattr(context, name, None)
    if async_context is None:
        async_context = AsyncContext(loop=loop, name=name, **kwargs)
        setattr(context, async_context.name, async_context)
    assert isinstance(async_context, AsyncContext)
    assert getattr(context, async_context.name) is async_context
    return async_context
