"""Dependency utilities"""

from types import ModuleType

from ipyparallel.client.asyncresult import AsyncResult
from ipyparallel.error import UnmetDependency
from ipyparallel.serialize import can
from ipyparallel.util import interactive


class depend:
    """Dependency decorator, for use with tasks.

    `@depend` lets you define a function for engine dependencies
    just like you use `apply` for tasks.


    Examples
    --------
    ::

        @depend(df, a,b, c=5)
        def f(m,n,p)

        view.apply(f, 1,2,3)

    will call df(a,b,c=5) on the engine, and if it returns False or
    raises an UnmetDependency error, then the task will not be run
    and another engine will be tried.
    """

    def __init__(self, _wrapped_f, *args, **kwargs):
        self.f = _wrapped_f
        self.args = args
        self.kwargs = kwargs

    def __call__(self, f):
        return dependent(f, self.f, *self.args, **self.kwargs)


class dependent:
    """A function that depends on another function.
    This is an object to prevent the closure used
    in traditional decorators, which are not picklable.
    """

    def __init__(self, _wrapped_f, _wrapped_df, *dargs, **dkwargs):
        self.f = _wrapped_f
        name = getattr(_wrapped_f, '__name__', 'f')
        self.__name__ = name
        self.df = _wrapped_df
        self.dargs = dargs
        self.dkwargs = dkwargs

    def check_dependency(self):
        if self.df(*self.dargs, **self.dkwargs) is False:
            raise UnmetDependency()

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


@interactive
def _require(*modules, **mapping):
    """Helper for @require decorator."""
    from ipyparallel.error import UnmetDependency
    from ipyparallel.serialize import uncan

    user_ns = globals()
    for name in modules:
        try:
            exec(f'import {name}', user_ns)
        except ImportError:
            raise UnmetDependency(name)

    for name, cobj in mapping.items():
        user_ns[name] = uncan(cobj, user_ns)
    return True


def require(*objects, **mapping):
    """Simple decorator for requiring local objects and modules to be available
    when the decorated function is called on the engine.

    Modules specified by name or passed directly will be imported
    prior to calling the decorated function.

    Objects other than modules will be pushed as a part of the task.
    Functions can be passed positionally,
    and will be pushed to the engine with their __name__.
    Other objects can be passed by keyword arg.

    Examples
    --------
    ::

        In [1]: @ipp.require('numpy')
           ...: def norm(a):
           ...:     return numpy.linalg.norm(a,2)

    ::

        In [2]: foo = lambda x: x*x
        In [3]: @ipp.require(foo)
           ...: def bar(a):
           ...:     return foo(1-a)
    """
    names = []
    for obj in objects:
        if isinstance(obj, ModuleType):
            obj = obj.__name__

        if isinstance(obj, str):
            names.append(obj)
        elif hasattr(obj, '__name__'):
            mapping[obj.__name__] = obj
        else:
            raise TypeError(
                "Objects other than modules and functions "
                f"must be passed by kwarg, but got: {type(obj)}"
            )

    for name, obj in mapping.items():
        mapping[name] = can(obj)
    return depend(_require, *names, **mapping)


class Dependency(set):
    """An object for representing a set of msg_id dependencies.

    Subclassed from set().

    Parameters
    ----------
    dependencies: list/set of msg_ids or AsyncResult objects or output of Dependency.as_dict()
        The msg_ids to depend on
    all : bool [default True]
        Whether the dependency should be considered met when *all* depending tasks have completed
        or only when *any* have been completed.
    success : bool [default True]
        Whether to consider successes as fulfilling dependencies.
    failure : bool [default False]
        Whether to consider failures as fulfilling dependencies.

    If `all=success=True` and `failure=False`, then the task will fail with an ImpossibleDependency
        as soon as the first depended-upon task fails.
    """

    all = True
    success = True
    failure = True

    def __init__(self, dependencies=[], all=True, success=True, failure=False):
        if isinstance(dependencies, dict):
            # load from dict
            all = dependencies.get('all', True)
            success = dependencies.get('success', success)
            failure = dependencies.get('failure', failure)
            dependencies = dependencies.get('dependencies', [])
        ids = []

        # extract ids from various sources:
        if isinstance(dependencies, (str, AsyncResult)):
            dependencies = [dependencies]
        for d in dependencies:
            if isinstance(d, str):
                ids.append(d)
            elif isinstance(d, AsyncResult):
                ids.extend(d.msg_ids)
            else:
                raise TypeError(f"invalid dependency type: {type(d)!r}")

        set.__init__(self, ids)
        self.all = all
        if not (success or failure):
            raise ValueError("Must depend on at least one of successes or failures!")
        self.success = success
        self.failure = failure

    def check(self, completed, failed=None):
        """check whether our dependencies have been met."""
        if len(self) == 0:
            return True
        against = set()
        if self.success:
            against = completed
        if failed is not None and self.failure:
            against = against.union(failed)
        if self.all:
            return self.issubset(against)
        else:
            return not self.isdisjoint(against)

    def unreachable(self, completed, failed=None):
        """return whether this dependency has become impossible."""
        if len(self) == 0:
            return False
        against = set()
        if not self.success:
            against = completed
        if failed is not None and not self.failure:
            against = against.union(failed)
        if self.all:
            return not self.isdisjoint(against)
        else:
            return self.issubset(against)

    def as_dict(self):
        """Represent this dependency as a dict. For json compatibility."""
        return dict(
            dependencies=list(self),
            all=self.all,
            success=self.success,
            failure=self.failure,
        )


__all__ = ['depend', 'require', 'dependent', 'Dependency']
