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

from .decorator import make_data_last

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


@overload
def unique_by(it: Iterable[T], fn: Callable[[T], U], /) -> Iterable[T]: ...


@overload
def unique_by(fn: Callable[[T], U]) -> Callable[[Iterable[T]], Iterable[T]]: ...


@make_data_last
def unique_by(iterable: Iterable[T], function: Callable[[T], U], /) -> Iterable[T]:
    """
    Yields unique elements from the iterable, given the iterable and a mapping function.

    Uniqueness is determined by resulting in the same value when passed to the mapping function.

    Parameters
    ----------
    iterable: Iterable[T]
        Iterable (positional-only).
    function: Callable[[T], U]
        Function to apply to each element of the iterable (positional-only).

    Returns
    -------
    Iterable[T]
        Unique elements from the iterable.

    Examples
    --------
    Data first:
    >>> list(
    ...     R.unique_by(
    ...             [{'n': 1}, {'n': 2}, {'n': 2}, {'n': 5}, {'n': 1}, {'n': 6}, {'n': 7}],
    ...             R.prop('n'),
    ...         )
    ... )
    [{'n': 1}, {'n': 2}, {'n': 5}, {'n': 6}, {'n': 7}]

    Data last:
    >>> R.pipe(
    ...     [{'n': 1}, {'n': 2}, {'n': 2}, {'n': 5}, {'n': 1}, {'n': 6}, {'n': 7}],
    ...     R.unique_by(R.prop('n')),
    ...     R.take(3),
    ...     list,
    ... )
    [{'n': 1}, {'n': 2}, {'n': 5}]

    """
    seen = set[U]()
    for x in iterable:
        key = function(x)
        if key not in seen:
            seen.add(key)
            yield x
