from collections.abc import Callable, Iterable
from typing import TypeVar, overload

from .decorator import make_data_last
from .find import find

T = TypeVar('T')
U = TypeVar('U')


@overload
def intersection_with(
    data: Iterable[T],
    other: Iterable[U],
    is_equal: Callable[[T, U], bool],
    /,
) -> Iterable[T]: ...


@overload
def intersection_with(
    other: Iterable[U],
    is_equal: Callable[[T, U], bool],
    /,
) -> Callable[[Iterable[T]], Iterable[T]]: ...


@make_data_last
def intersection_with(
    iterable: Iterable[T],
    other: Iterable[U],
    equality_function: Callable[[T, U], bool],
    /,
) -> Iterable[T] | Callable[[Iterable[T]], Iterable[T]]:
    """
    Yields elements of the first iterable that equal to an element in the other iterable.

    Given two iterables and an equality function.

    The inputs are treated as multi-sets/bags (multiple copies of items are treated as unique items).
    The order of elements is maintained.

    Parameters
    ----------
    iterable : Iterable[T]
        First iterable (positional-only).
    other: Iterable[U]
        Second iterable (positional-only).
    equality_function: Callable[[T, U], bool]
        Equality function, predicate of arity 2 (positional-only).

    Returns
    -------
    Iterable[T]
        Iterable of elements of the first iterable that appear in the other iterable.

    Examples
    --------
    Data first:
    >>> list(
    ...  R.intersection_with(
    ...   [{'id': 1, 'name': 'Ryan'}, {'id': 3, 'name': 'Emma'}],
    ...   [3, 5],
    ...   lambda d, x: d['id'] == x,
    ...  )
    ... )
    [{'id': 3, 'name': 'Emma'}]

    Data last:
    >>> R.pipe(
    ...  [{'id': 1, 'name': 'Ryan'}, {'id': 3, 'name': 'Emma'}],
    ...  R.intersection_with([3, 5], lambda d, x: d['id'] == x),
    ...  list,
    ... )
    [{'id': 3, 'name': 'Emma'}]

    """
    other_ = list(other)
    for x in iterable:
        if (i := find(other_, lambda a, xx=x: equality_function(xx, a))) is not None:
            other_.remove(i)
            yield x
