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

from .decorator import make_data_last

K = TypeVar('K')
V = TypeVar('V')


@overload
def group_by_prop(iterable: Iterable[dict[K, V]], prop: K, /) -> dict[V, list[dict[K, V]]]: ...


@overload
def group_by_prop(prop: K, /) -> Callable[[Iterable[dict[K, V]]], dict[V, list[dict[K, V]]]]: ...


@make_data_last
def group_by_prop(iterable: Iterable[dict[K, V]], prop: K, /) -> dict[V, list[dict[K, V]]]:
    """
    Groups dicts by the value under a key.

    Returns a dict with keys being the values of the given key in the dicts in the iterable.

    Values of the returned dict are lists of dicts with the same value for the key. Order is preserved.

    If the value is not in the dict or if it is None the dict is skipped.

    Parameters
    ----------
    iterable: Iterable[dict[K, V]
        The iterable (positional-only).
    prop: K
        Function to apply to each element (positional-only).

    Returns
    -------
    dictVK,T]

    See Also
    --------
    group_by

    Examples
    --------
    Data first:
    >>> R.group_by_prop([{'a': 'cat'}, {'a': 'dog'}, {}, {'b': 'cat'}], 'a')
    {'cat': [{'a': 'cat'}], 'dog': [{'a': 'dog'}]}

    Data last:
    >>> R.pipe([{'a': 'cat'}, {'a': 'dog'}], R.group_by_prop('a'))
    {'cat': [{'a': 'cat'}], 'dog': [{'a': 'dog'}]}

    """
    result: defaultdict[V, list[dict[K, V]]] = defaultdict(list)
    for d in iterable:
        if (k := d.get(prop)) is not None:
            result[k].append(d)
    return dict(result)
