import inspect
from collections.abc import Callable
from typing import TypeVar, cast, overload

from remedapy.decorator import make_data_last

Key = TypeVar('Key')
V = TypeVar('V')
NewV = TypeVar('NewV')

Fn = Callable[[V, Key], NewV] | Callable[[V], NewV] | Callable[[], NewV]


@overload
def map_values(data: dict[Key, V], fn: Fn[V, Key, NewV], /) -> dict[Key, NewV]: ...


@overload
def map_values(fn: Fn[V, Key, NewV], /) -> Callable[[dict[Key, V]], dict[Key, NewV]]: ...


@make_data_last
def map_values(data: dict[Key, V], function: Fn[V, Key, NewV], /) -> dict[Key, NewV]:
    """
    Given a dict and a mapping function returns a dict with values mapped.

    Parameters
    ----------
    data : dict[Key, V]
        Original dict (positional-only).
    function: Callable[[V, Key], NewV] | Callable[[V], NewV] | Callable[[], NewV]
        The function to apply to each element of the iterable (positional-only).

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

    Examples
    --------
    Data first:
    >>> R.map_values({'a': 1, 'b': 2}, lambda v, key: f'{v}{key}')
    {'a': '1a', 'b': '2b'}
    >>> R.map_values({'a': 1, 'b': 2}, R.add(1))
    {'a': 2, 'b': 3}
    >>> R.map_values({'a': 1, 'b': 2}, R.constant('1'))
    {'a': '1', 'b': '1'}

    Data last:
    >>> R.pipe({'a': 1, 'b': 2}, R.map_values(lambda v, key: f'{v}{key}'))
    {'a': '1a', 'b': '2b'}

    """
    sig = inspect.signature(function)
    param_count = len(sig.parameters.values())
    if param_count > 1:
        function = cast(Callable[[V, Key], NewV], function)
        return {key: function(value, key) for key, value in data.items()}
    elif param_count == 1:
        function = cast(Callable[[V], NewV], function)
        return {key: function(value) for key, value in data.items()}
    else:
        function = cast(Callable[[], NewV], function)
        return {key: function() for key, _ in data.items()}
