"""
.. module:: factory
   :synopsis: Functions and decorators to construct nuts.
"""
from __future__ import absolute_import

import functools

from nutsflow.base import Nut, NutSink, NutSource, NutFunction


def _arg_insert(args, arg, pos=0):
    """
    Insert arg in args at given position.

    :param tuple args: Some function arguments
    :param any arg: Some function argument
    :param int pos: Insert position. If None argument is appended.
    :return: List with arguments where arg is inserted
    :rtype: list
    """
    args = list(args)
    if pos is None:
        args.append(arg)
    else:
        args.insert(pos, arg)
    return args


def _wrap(wrappercls, func):
    """
    Return wrapped function.

    Used to ensure that decorated nut function has the correct docstring.

    :param class wrappercls: Nut wrapper class
    :param function func: Function to wrap
    :return: Wrapped function
    :rtype: function
    """
    @functools.wraps(func)
    def wrapper(*args, **kwds):
        return wrappercls(*args, **kwds)

    return wrapper


def _create_nut_wrapper(base_class, func, iterpos):
    """
    Return Nut for given function.

    :param class base_class: Base class, e.g. Nut, NutSink, NutFunction, ...
    :param function func: Function to wrap
    :param int iterpos: Argument position for iterable in function.
    :return: Nut that wraps the given function.
    :rtype: Nut
    """

    class Wrapper(base_class):

        def __rrshift__(self, iterable):
            args = _arg_insert(self.args, iterable, iterpos)
            return func(*args, **self.kwargs)

    return _wrap(Wrapper, func)


def _create_filter_wrapper(func, invert=False):
    """
    Return filter Nut for given function.

    :param func: Filter function to wrap
    :param invert: Filter is inverted.
    :return: Nut operates as a filter.
    :rtype: Nut
    """

    class Wrapper(Nut):

        def __rrshift__(self, iterable):
            for e in iterable:
                args = _arg_insert(self.args, e)
                if bool(func(*args, **self.kwargs)) != invert:
                    yield e

    return _wrap(Wrapper, func)


def nut_function(func):
    """
    Decorator for Nut functions.

    Example on how to define a custom function nut:

    .. code::

      @nut_function
      def TimesN(x, n):
          return x * n

      [1, 2, 3] >> TimesN(2) >> Collect()  -->  [2, 4, 6]

    :param function func: Function to decorate
    :return: Nut function for given function
    :rtype: NutFunction
    """

    class Wrapper(NutFunction):

        def __call__(self, element):
            return func(element, *self.args, **self.kwargs)

    return _wrap(Wrapper, func)


def nut_source(func):
    """
    Decorator for Nut sources.

    Example on how to define a custom source nut. Note that a source
    must return an iterable/generator and does not read any input.

    .. code::

      @nut_source
      def MyRange(start, end):
          return range(start, end)

      MyRange(0, 5) >> Collect()  --> [0, 1, 2, 3, 4]


    .. code::

      @nut_source
      def MyRange2(start, end):
          for i in range(start, end):
              yield i * 2

      MyRange2(0, 5) >> Collect()  --> [0, 2, 4, 6, 8]

    :param function func: Function to decorate
    :return: Nut source for given function
    :rtype: NutSource
    """

    class Wrapper(NutSource):

        def __iter__(self):
            return func(*self.args, **self.kwargs)

    return _wrap(Wrapper, func)


def nut_processor(func, iterpos=0):
    """
    Decorator for Nut processors.

    Examples on how to define a custom processor nut.
    Note that a processor reads an iterable and must return
    an iterable/generator

    .. code::

      @nut_processor
      def Twice(iterable):
          for e in iterable:
              yield e
              yield e

      [1, 2, 3] >> Twice() >> Collect()  --> [1, 1, 2, 2, 3, 3]


    .. code::

      @nut_processor
      def Odd(iterable):
          return (e for e in iterable if e % 2)

      [1, 2, 3, 4, 5] >> Odd() >> Collect()  --> [1, 3, 5]


    .. code::

      @nut_processor
      def Clone(iterable, n):
          for e in iterable:
              for _ in range(p):
                  yield e

      [1, 2, 3] >> Clone(2) >> Collect()  --> [1, 1, 2, 2, 3, 3]

    :param function func: Function to decorate
    :param iterpos: Position of iterable in function arguments
    :return: Nut processor for given function
    :rtype: Nut
    """
    return _create_nut_wrapper(Nut, func, iterpos)


def nut_sink(func, iterpos=0):
    """
    Decorator for Nut sinks.

    Example on how to define a custom sink nut:

    .. code::

      @nut_sink
      def ToList(iterable):
          return list(iterable)

      range(5) >> ToList()  -->   [0, 1, 2, 3, 4]


    .. code::

      @nut_sink
      def MyCollect(iterable, container):
          return container(iterable)

      range(5) >> MyCollect(tuple)  -->   (0, 1, 2, 3, 4)


    .. code::

      @nut_sink
      def MyProd(iterable):
          p = 1
          for e in iterable:
              p *= e
          return p

      [1, 2, 3] >> MyProd()  --> 12


    :param function func: Function to decorate
    :param iterpos: Position of iterable in function arguments
    :return: Nut sink for given function
    :rtype: NutSink
    """
    return _create_nut_wrapper(NutSink, func, iterpos)


def nut_filter(func):
    """
    Decorator for Nut filters.

    Also see nut_filerfalse().
    Example on how to define a custom filter nut:

    .. code::

      @nut_filter
      def Positive(x):
          return x > 0

      [-1, 1, -2, 2] >> Positive() >> Collect()  --> [1, 2]


    .. code::

      @nut_filter
      def GreaterThan(x, threshold):
          return x > threshold

      [1, 2, 3, 4] >> GreaterThan(2) >> Collect()  --> [3, 4] 

    :param function func: Function to decorate. Must return boolean value.
    :return: Nut filter for given function
    :rtype: Nut
    """
    return _create_filter_wrapper(func, invert=False)


def nut_filterfalse(func):
    """
    Decorator for Nut filters that are inverted.

    Also see nut_filter().
    Example on how to define a custom filter-false nut:

    .. code::

      @nut_filterfalse
      def NotGreaterThan(x, threshold):
          return x > threshold

      [1, 2, 3, 4] >> NotGreaterThan(2) >> Collect()  --> [1, 2]

    :param function func: Function to decorate
    :return: Nut filter for given function. . Must return boolean value.
    :rtype: Nut
    """

    return _create_filter_wrapper(func, invert=True)
