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 index_by(data: Iterable[T], fn: Callable[[T], K], /) -> dict[K, T]: ...


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


@make_data_last
def index_by(iterable: Iterable[T], function: Callable[[T], K], /) -> dict[K, T]:
    """
    Given an iterable and a function returns a dict.

    The dicts keys are the results of applying the function to the elements of the iterable.
    The values are the values from the iterable.

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

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

    See Also
    --------
    group_by

    Examples
    --------
    Data first:
    >>> R.index_by(['one', 'two', 'three'], R.length)
    {3: 'two', 5: 'three'}

    Data last:
    >>> R.pipe(['one', 'two', 'three'], R.index_by(R.length))
    {3: 'two', 5: 'three'}

    """
    result: dict[K, T] = {}
    for d in iterable:
        result[function(d)] = d
    return result
