"""Thread handling functions.

Sample usage:

    >>> @threaded
    >>> def such_long():
    >>>     time.sleep(2)
    >>>     return 'much wiki'

    >>> @threaded
    >>> def much_slow():
    >>>     time.sleep(1)
    >>>     return 'such spin'

    >>> @threaded
    >>> def nope():
    >>>     raise NotImplementedError

    >>> join(
    >>>     such_long(),
    >>>     much_slow()
    >>> )
    ('much wiki', 'such spin')

    >>> ops = (such_long, much_slow)
    >>> join(op() for op in ops)
    ('much wiki', 'such spin')

    >>> ops = (such_long, much_slow)
    >>> join(op() for op in ops)
    ('much wiki', 'such spin')

    >>> join(
    >>>     such_long(),
    >>>     nope(),
    >>>     much_slow()
    >>> )
    Traceback (most recent call last):
      ...
    NotImplementedError
"""

import queue
import threading


def join(*threads):
    """Wait for all passed `@threaded` functions return values.

    If any exception was raised, it is automatically re-raised. Passed
    functions are always joined.
    """
    exc = None

    if len(threads) == 0:  # pragma: nocover
        return []
    elif len(threads) == 1:
        if hasattr(threads[0], "__iter__"):
            threads = list(threads[0])

    for thread in threads:
        thread.join()

        if not thread.e.empty():
            exc = thread.e.get()

    if exc:
        raise exc

    return [thread.q.get() for thread in threads]


def threaded(f, daemon=False):
    """Declare a function as a thread.

    This high-level implementation relies on `threading.Thread` but adds 2
    queues for safe operations:

    - `q`: Thread's return value (after it finishes)
    - `e`: Thread's raised exception

    These queues are mutually exclusive and since some operations on queues are
    blocking, some time should be spent on the documentation. Alternatively, a
    high-level `join()` implementation is provided.
    """

    def wrapped_f(q, e, *args, **kwargs):
        try:
            q.put(f(*args, **kwargs))
        except Exception as exc:
            e.put(exc)

    def wrap(*args, **kwargs):
        q = queue.Queue()
        e = queue.Queue()

        t = threading.Thread(
            target=wrapped_f, args=(q, e) + args, kwargs=kwargs
        )

        t.daemon = daemon
        t.start()

        t.q = q
        t.e = e

        return t

    return wrap