"""
forbiddenfp - Some builtin/itertools functions that can be chained after objects, in favor of more functional programming.

https://github.com/yx-z/forbiddenfp
"""
import builtins
import collections
import functools
import itertools
import operator
from numbers import Number
from typing import Callable, List, NoReturn, Optional, Dict, TypeVar, Iterable, Sequence, Tuple, Union, Set, Type, \
    ContextManager

from forbiddenfruit import curse
from typing_extensions import ParamSpec

_P = ParamSpec("_P")
_R = TypeVar("_R")
_F = Callable[[_P], _R]
_Pred = Callable[[_P], bool]
_P2 = ParamSpec("_P2")
_R2 = TypeVar("_R2")


def chain_as(name: str, cls: Type = object) -> Callable[[_F], _F]:
    def decorator(func: _F) -> _F:
        curse(cls, name, func)
        return func

    return decorator


def chain_for(cls: Type) -> Callable[[_F], _F]:
    def decorator(func: _F) -> _F:
        return chain_as(func.__name__, cls)(func)

    return decorator


def chainable(func: _F) -> _F:
    return chain_as(func.__name__)(func)


def always(val: _P) -> Callable[..., _P]:
    return lambda *args, **kwargs: val


def greater_than(val: _P) -> _Pred:
    return lambda a: a > val


chain_as("greater_than")(lambda self, val: greater_than(val)(self))

positive = greater_than(0)
chain_as("positive")(lambda self: positive(self))


def less_than(val: _P) -> _Pred:
    return lambda a: a < val


chain_as("less_than")(lambda self, val: less_than(val)(self))

negative = less_than(0)
chain_as("negative")(lambda self: negative(self))


def greater_than_or_equals(val: _P) -> _Pred:
    return lambda a: a >= val


chain_as("greater_than_or_equals")(lambda self, val: greater_than_or_equals(val)(self))

non_negative = greater_than_or_equals(0)
chain_as("non_negative")(lambda self: non_negative(self))


def less_than_or_equals(val: _P) -> _Pred:
    return lambda a: a >= val


chain_as("less_than_or_equals")(lambda self, val: less_than_or_equals(val)(self))

non_positive = less_than_or_equals(0)
chain_as("non_positive")(lambda self: non_positive(self))


def identity(val: _P) -> _P:
    return val


@chain_as("not")
@chainable
def negate(val: _P) -> _P:
    return not val


@chain_as("not_func")
@chainable
def negate_func(func: _Pred) -> _Pred:
    return lambda *args, **kwargs: not func(*args, **kwargs)


def equals(val: _P) -> _Pred:
    return lambda a: a == val


chain_as("equals")(lambda self, val: equals(val)(self))
equals_zero = chain_as("equals_zero")(equals(0))


def not_equals(val: _P) -> Callable[[Iterable[_P]], bool]:
    return lambda a: val != a


chain_as("not_equals")(lambda self, val: not_equals(val)(self))
not_equals_zero = chain_as("not_equals_zero")(not_equals(0))


def contains(val: _P) -> Callable[[Iterable[_P]], bool]:
    return lambda a: val in a


chain_as("contains")(lambda self, val: contains(val)(self))


def not_contains(val: _P) -> Callable[[Iterable[_P]], bool]:
    return lambda a: val not in a


chain_as("not_contains")(lambda self, val: not_contains(val)(self))

truthful = chain_as("truthful")(bool)
falseful = chain_as("falseful")(negate_func(bool))


def in_iter(it: Iterable[_P]) -> Callable[[_P], bool]:
    return lambda a: a in it


chain_as("in_iter")(lambda self, val: in_iter(val)(self))


def not_in(it: Iterable[_P]) -> Callable[[_P], bool]:
    return lambda a: a not in it


chain_as("not_in")(lambda self, val: not_in(val)(self))


@chainable
def asserting(self: _P, pred: _Pred) -> _P:
    assert pred(self)
    return self


@chainable
def is_none(val: _P) -> bool:
    return val is None


@chainable
def is_not_none(val: _P) -> bool:
    return val is not None


def add(val: _P) -> _F:
    return lambda a: a + val


chain_as("add")(lambda self, val: add(val)(self))


def subtract(val: _P) -> _F:
    return lambda a: a - val


chain_as("subtract")(lambda self, val: subtract(val)(self))


def multiply(val: _P) -> _F:
    return lambda a: a * val


chain_as("multiply")(lambda self, val: multiply(val)(self))


def divide(val: _P) -> _F:
    return lambda a: a / val


chain_as("divide")(lambda self, val: divide(val)(self))


def and_val(self: _P, other: _P2) -> _R:
    return self and other


@chainable
def and_func(self: _Pred, other: _Pred) -> _Pred:
    return lambda *args, **kwargs: self(*args, **kwargs) and other(*args, **kwargs)


def or_val(self: _P, other: _P2) -> _R:
    return self or other


@chainable
def or_func(self: _Pred, other: _Pred) -> _Pred:
    return lambda *args, **kwargs: self(*args, **kwargs) or other(*args, **kwargs)


@chainable
def map_val(self: _P, vals: Optional[Dict[_P, _R]] = None, default: Optional[_R] = None, **kwargs: _R) -> Optional[_R]:
    return {**(vals or {}), **kwargs}.get(self, default)


@chainable
def map_pred(self: _P, pred_and_val: Dict[_Pred, _R], default: Optional[_R] = None) -> Optional[_R]:
    for pred, val in pred_and_val.items():
        if pred(self):
            return val
    return default


@chainable
def match_val(self: _P,
              val_to_action: Optional[Dict[_P, _F]] = None,
              default: _F = always(None), **kwargs: _F) -> Optional[_R]:
    return {**(val_to_action or {}), **kwargs}.get(self, default)(self)


@chainable
def match_pred(self: _P, pred_and_action: Dict[_Pred, _R], default: _F = always(None)) -> Optional[_R]:
    for pred, action in pred_and_action.items():
        if pred(self):
            return action(self)
    return default(self)


def isinstance_val(t: _P) -> _Pred:
    return lambda a: isinstance(a, t)


chain_as("isinstance")(lambda self, t: isinstance_val(t)(self))


@chain_as("invoke")
@chainable
def call(self: _F, *args: _P.args, **kwargs: _P.kwargs) -> _R:
    return self(*args, **kwargs)


@chainable
def compose(*funcs: Iterable[Callable[[_P], _R]]) -> Callable[[_P2], _R2]:
    def _compose2(f: Callable[[_R], _R2], g: _F) -> Callable[[_P], _R2]:
        return lambda *args, **kwargs: f(g(*args, **kwargs))

    return functools.reduce(_compose2, funcs)


@chainable
def compose_r(*funcs: Iterable[Callable[[_P], _R]]) -> Callable[[_P2], _R2]:
    return compose(*reversed(funcs))


@chainable
def flatten(self: List, recurse: bool = False) -> List:
    res = []
    for x in self:
        if isinstance(x, list):
            res += flatten(x, recurse=True) if recurse else x
        else:
            res.append(x)
    return res


@chainable
def partial(self: _F, *args: _P.args, **kwargs: _P.kwargs) -> _F:
    return functools.partial(self, *args, **kwargs)


@chainable
def apply(self: _P, func: _F) -> _P:
    func(self)
    return self


# don't clash function name with builtin print. But we can still chain as `object().print()`.
@chain_as("print")
def apply_print(self: _P, print_func: _F = builtins.print) -> _P:
    print_func(self)
    return self


@chainable
def print_list(self: Iterable[_P], print_func: _F = builtins.print) -> List[_P]:
    ls = list(self)
    print_func(ls)
    return ls


@chainable
def pair(self: _P, func: _F) -> Tuple[_P, _R]:
    return self, func(self)


@chainable
def pairwise(self: Iterable[_P]) -> Iterable[Tuple[_P, _R]]:
    return itertools.pairwise(self)


@chainable
def also(self: _P, _: _R) -> _P:
    # side effect is evaluated before passing in to this function
    return self


@chainable
def then(self: _P, func: _F) -> _R:
    return func(self)


@chainable
def then_use(self: _P, val: _R) -> _R:
    return val


@chain_as("setattr")
def set_attr(self: _P, **kwargs: _P.kwargs) -> _P:
    for k, v in kwargs.items():
        builtins.setattr(self, k, v)
    return self


@chainable
def setitem(self: _P, key: _P2, val: _R) -> _P:
    self[key] = val
    return self


@chainable
def apply_unpack(self: _P, func: _F) -> _R:
    func(*self)
    return self


@chainable
def then_unpack(self: _P, func: _F) -> _R:
    return func(*self)


@chainable
def empty(self: Iterable[_P]) -> bool:
    return all(False for _ in self)


def tee(self: Iterable[_P], n: int = 2) -> Tuple[_P, ...]:
    return itertools.tee(self, n)


@chain_as("min")
def min_iter(self: Iterable[_P], key: _F = identity) -> _P:
    return builtins.min(self, key=key)


@chain_as("max")
def max_iter(self: Iterable[_P], key: _F = identity) -> _P:
    return builtins.max(self, key=key)


@chain_as("range")
def range_up_to(self: int, start: Optional[int] = None, step: Optional[int] = None) -> Iterable[int]:
    if start is None and step is None:
        return builtins.range(self)
    if step is None:
        return builtins.range(start, self)
    return builtins.range(start, self, step)


@chain_as("all")
def all_iter(self: Iterable[_P], predicate: _Pred = truthful) -> bool:
    return builtins.all(predicate(i) for i in self)


@chain_as("any")
def any_iter(self: Iterable[_P], predicate: _Pred = truthful) -> bool:
    return builtins.any(predicate(i) for i in self)


@chain_as("map")
def map_iter(self: Iterable[_P], func: _F, *other: Iterable[_P2]) -> Iterable[_R]:
    return builtins.map(func, self, *other)


@chainable
def map_dict(self: Dict[_P, _P2], key_func: _F = identity, val_func: Callable[[_P2], _R2] = identity) -> Dict[
    _R, _R2]:
    return {key_func(k): val_func(v) for k, v in self.items()}


@chain_as("filter")
def filter_iter(self: Iterable[_P], predicate: _Pred = truthful) -> Iterable[_P]:
    return builtins.filter(predicate, self)


@chainable
def filter_key(self: Dict[_P, _P2], predicate: Callable[[_P], bool]) -> Iterable[Tuple[_P, _P2]]:
    return ((k, v) for k, v in self.items() if predicate(k))


@chainable
def filter_val(self: Dict[_P, _P2], predicate: Callable[[_P2], bool]) -> Iterable[Tuple[_P, _P2]]:
    return ((k, v) for k, v in self.items() if predicate(v))


@chainable
def last(self: Sequence[_P], predicate: _Pred = always(True)) -> Optional[_P]:
    return builtins.next((i for i in builtins.reversed(self) if predicate(i)), None)


@chain_as("next")
@chainable
def first(self: Iterable[_P], predicate: _Pred = always(True)) -> Optional[_P]:
    return builtins.next((i for i in self if predicate(i)), None)


@chain_as("sum")
def sum_iter(self: Iterable[_P], of: Callable[[_P], Number] = identity, predicate: _Pred = truthful) -> int:
    return builtins.sum(of(i) for i in self if predicate(i))


@chain_as("len")
def len_iter(self: Iterable[_P], predicate: _Pred = truthful) -> int:
    return builtins.sum(1 for i in self if predicate(i))


@chain_as("reversed")
def reversed_iter(self: Sequence[_P]) -> Sequence[_P]:
    return builtins.reversed(self)


@chain_as("sorted")
def sorted_iter(self: Iterable[_P], key: _F = identity, reverse: bool = False) -> List[_P]:
    return builtins.sorted(self, key=key, reverse=reverse)


@chain_as("foldl")
@chain_as("fold")
@chainable
def reduce(self: Iterable[_P], func: Callable[[_P, _P], _R], initial: Optional[_R] = None) -> _R:
    return functools.reduce(func, self, initial) if initial is not None else functools.reduce(func, self)


@chainable
def foldr(self: Iterable[_P], func: Callable[[_P, _P], _R], initial: Optional[_R] = None) -> _R:
    rev = reversed(self)
    return functools.reduce(func, rev, initial) if initial is not None else functools.reduce(func, rev)


@chainable
def separate(self: Iterable[_P], predicate: _Pred = truthful) -> Tuple[List[_P], List[_P]]:
    true_part = []
    false_part = []
    for i in self:
        (true_part if predicate(i) else false_part).append(i)
    return true_part, false_part


@chainable
def counter(self: Iterable[_P]) -> Dict[_P, int]:
    return collections.Counter(self)


@chainable
def groupby(self: Iterable[_P], key: _F) -> Dict[_R, List[_P]]:
    return itertools.groupby(self, key=key)


@chainable
def chain(*self: Iterable[Iterable[_P]]) -> Iterable[_P]:
    return itertools.chain(*self)


@chain_as("zip")
def zip_iter(*self: Iterable[_P]) -> Iterable[Tuple[_P, ...]]:
    return builtins.zip(*self)


@chain_as("enumerate")
def enumerate_iter(self: Iterable[_P]) -> Iterable[Tuple[int, _P]]:
    return builtins.enumerate(self)


@chain_as("tuple")
def tuple_iter(self: Iterable[_P]) -> Tuple[_P, ...]:
    return builtins.tuple(self)


@chain_as("list")
def list_iter(self: Iterable[_P]) -> List[_P]:
    return builtins.list(self)


@chain_as("set")
def set_iter(self: Iterable[_P]) -> Set[_P]:
    return builtins.set(self)


@chain_as("dict")
def dict_iter(self: Iterable[Tuple[_P, _P2]]) -> Dict[_P, _P2]:
    return builtins.dict(self)


@chainable
def join(self: Iterable[_P], sep: builtins.str = "", to_str: Callable[[_P], str] = builtins.str) -> builtins.str:
    return sep.join(builtins.map(to_str, self))


@chainable
def starmap(self: Iterable[Iterable[_P]], func: Callable[[Iterable[_P]], _R]) -> Iterable[_R]:
    return itertools.starmap(func, self)


@chainable
def accumulate(self: Iterable[Iterable[_P]], func: Callable[[_R, _P], _R] = operator.add,
               initial: Optional[_R] = None) -> Iterable[_R]:
    return (itertools.accumulate(self, func, initial=initial) if initial is not None
            else itertools.accumulate(self, func))


@chainable
def pairwise(self: Iterable[_P]) -> Iterable[Tuple[_P, _P]]:
    return itertools.pairwise(self)


@chainable
def product(self: Iterable[_P], repeat: int = 1) -> Iterable[Tuple[_P, ...]]:
    return itertools.product(*self, repeat=repeat)


@chainable
def repeat(self: _P, times: Optional[int] = None) -> Iterable[_P]:
    return itertools.repeat(self, times)


@chainable
def inifinite(self: _P) -> Iterable[_P]:
    return itertools.repeat(self, None)


@chainable
def cycle(self: _P) -> Iterable[_P]:
    return itertools.cycle(self)


@chainable
def takewhile(self: Iterable[_P], predicate: _Pred = truthful) -> Iterable[_P]:
    return itertools.takewhile(predicate, self)


@chain_as("next_n")
@chain_as("first_n")
@chainable
def islice_up_to(self: Iterable[_P], stop: int, predicate: _Pred = always(True)) -> Iterable[_P]:
    return itertools.islice(builtins.filter(predicate, self), stop)


@chainable
def islice(
        self: Iterable[_P], start: int, stop: Optional[int], step: int = 1, predicate: _Pred = always(True)
) -> Iterable[_P]:
    return itertools.islice(builtins.filter(predicate, self), start, stop, step)


@chainable
def dropwhile(self: Iterable[_P], predicate: _Pred = truthful) -> Iterable[_P]:
    return itertools.dropwhile(predicate, self)


@chain_as("str")
def str_obj(self: _P) -> builtins.str:
    return builtins.str(self)


@chain_as("repr")
def repr_obj(self: _P) -> builtins.str:
    return builtins.repr(self)


@chain_as("format")
def format_obj(self: _P, format_str: builtins.str) -> builtins.str:
    return builtins.format(self, format_str)


@chainable
def if_branches(self: _P, func_true: _F, func_false: _F, predicate: _Pred = truthful) -> _R:
    return func_true(self) if predicate(self) else func_false(self)


@chainable
def if_true(self: _P, func: _F, predicate: _Pred = truthful) -> Optional[_R]:
    return func(self) if predicate(self) else None


@chainable
def if_false(self: _P, func: _F, predicate: _Pred = truthful) -> Optional[_R]:
    return func(self) if not predicate(self) else None


@chainable
def or_else(self: _P, val: _R, predicate: _Pred = truthful) -> _R:
    return self if predicate(self) else val


@chainable
def or_eval(self: _P, func: _F, predicate: _Pred = truthful) -> _R:
    return self if predicate(self) else func(self)


ExceptionFuncOrException = Union[Callable[[_P], Exception], Exception]


@chainable
def or_raise(self: _P, func_or_val: ExceptionFuncOrException, predicate: _Pred = truthful) -> _P:
    if predicate(self):
        return self
    if callable(func_or_val):
        raise func_or_val(self)
    raise func_or_val


@chainable
def raise_as(self: _P, func_or_val: ExceptionFuncOrException) -> NoReturn:
    if callable(func_or_val):
        raise func_or_val(self)
    raise func_or_val


ExceptionType = Type[Exception]


@chainable
def then_catch(self: _P, exception_type: ExceptionType, func: _F,
               exception_handler: Callable[[_P, ExceptionType], _R2]) -> Union[_R, _R2]:
    try:
        return func(self)
    except exception_type as e:
        return exception_handler(self, e)


@chainable
def apply_while(self: _P, func: Callable[[_P], _P], predicate: _Pred = truthful) -> _P:
    while predicate(self):
        self = func(self)
    return self


@chainable
def generate_while(self: _P, func: Callable[[_P], _P], predicate: _Pred = truthful) -> Iterable[_P]:
    while predicate(self):
        yield self
        self = func(self)


@chainable
def with_context(self: _P, context_func: Callable[[_P], ContextManager],
                 then: Union[Callable[[_P, _R], _R2], Callable[[_P], _R2]]) -> _R2:
    with context_func(self) as r:
        if r is None:
            return then(self)
        return then(self, r)


@chainable
def take_unpacked(self: Iterable[_P],
                  func_iterable: Callable[[Iterable[_P], Callable[[Iterable[_P]], _R], _P.kwargs], _R2],
                  func_vararg: Callable[[_P.args], _R]) -> _R2:
    """
    Given an iterable `self`, a `func_iterable` that takes a (function that acts on the single parameter iterable),
    and a `func_vararg` that takes a (function that acts on variable length args, as opposed to the single parameter),
    Call `func_iterable on the iterable but with `func_vararg`.

    E.g. Given `filter = lambda iterable, predicate: [i for i in iterable if predicate(i)]`,
    `[(1,2), (3,4)].filter(lambda tup: tup[1] == 2)  # [(1,2)]`, where `tup` is an iterable.
    It can be rewritten as `[(1,2), (3,4)].take_unpacked(filter, lambda _, snd: snd == 2)`
    """
    return func_iterable(self, lambda iterable: func_vararg(*iterable))
