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')
NewKey = TypeVar('NewKey')

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


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


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


@make_data_last
def map_keys(data: dict[Key, V], function: Fn[Key, V, NewKey], /) -> dict[NewKey, V]:
    """
    Given a dict and a function, returns a new dict with the same values but with keys mapped by the function.

    Parameters
    ----------
    data: dict[Key, V]
        Dict to map keys.
    function: Callable[[Key, V], NewKey] | Callable[[Key], NewKey] | Callable[[], NewKey]
        Function to map keys.

    Returns
    -------
    dict[NewKey, V]
        Dict with mapped keys.

    Examples
    --------
    Data first:
    >>> glue = lambda k, v: f'{k}{v}'
    >>> R.map_keys({'a': 1, 'b': 2}, glue)
    {'a1': 1, 'b2': 2}
    >>> R.map_keys({'a': 1, 'bbb': 2}, R.length)
    {1: 1, 3: 2}
    >>> R.map_keys({'a': 1, 'bbb': 2}, R.constant('a'))
    {'a': 2}

    Data last:
    >>> R.pipe({'a': 1, 'b': 2}, R.map_keys(glue))
    {'a1': 1, 'b2': 2}

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