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

from .decorator import make_data_last

T = TypeVar('T')
K = TypeVar('K')


@overload
def group_by(data: Iterable[T], fn: Callable[[T], K], /) -> dict[K, list[T]]: ...


@overload
def group_by(data: Iterable[T], fn: Callable[[T], K | None], /) -> dict[K, list[T]]: ...


@overload
def group_by(fn: Callable[[T], K], /) -> Callable[[Iterable[T]], dict[K, list[T]]]: ...


@overload
def group_by(fn: Callable[[T], K | None], /) -> Callable[[Iterable[T]], dict[K, list[T]]]: ...


@make_data_last
def group_by(iterable: Iterable[T], function: Callable[[T], K | None], /) -> dict[K, list[T]]:
    """
    Groups elements of an iterable by a key function.

    Returns a dict with keys being the results of applying the function to the elements of the iterable.

    Values of the returned dict are lists of elements that map to the same key. Order is preserved.

    Elements mapping to None are skipped.

    Parameters
    ----------
    iterable: Iterable[T]
        The iterable (positional-only).
    function: Callable[[T], K | None]
        Function to apply to each element (positional-only).

    Returns
    -------
    dict[K,T]

    See Also
    --------
    group_by_prop
    index_by

    Examples
    --------
    Data first:
    >>> R.group_by([{'a': 'cat'}, {'a': 'dog'}], R.prop('a'))
    {'cat': [{'a': 'cat'}], 'dog': [{'a': 'dog'}]}
    >>> R.group_by([0, 1], lambda x: 'even' if R.is_even(x) else None)
    {'even': [0]}

    Data last:
    >>> R.pipe([{'a': 'cat'}, {'a': 'dog'}], R.group_by(R.prop('a')))
    {'cat': [{'a': 'cat'}], 'dog': [{'a': 'dog'}]}
    >>> R.pipe([0, 1], R.group_by(lambda x: 'even' if R.is_even(x) else None))
    {'even': [0]}

    """
    # TODO: index and whole array
    result: defaultdict[K, list[T]] = defaultdict(list)
    for d in iterable:
        if (k := function(d)) is not None:
            result[k].append(d)
    return dict(result)
