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

from .decorator import make_data_last

Key = TypeVar('Key', int, str)
T = TypeVar('T')


@overload
def for_each_dict(count: dict[Key, T], fn: Callable[[], Any], /) -> dict[Key, T]: ...
@overload
def for_each_dict(count: dict[Key, T], fn: Callable[[T], Any], /) -> dict[Key, T]: ...
@overload
def for_each_dict(count: dict[Key, T], fn: Callable[[T, Key], Any], /) -> dict[Key, T]: ...
@overload
def for_each_dict(count: dict[Key, T], fn: Callable[[T, Key, dict[Key, T]], Any], /) -> dict[Key, T]: ...


@overload
def for_each_dict(fn: Callable[[], Any], /) -> Callable[[dict[Key, T]], dict[Key, T]]: ...
@overload
def for_each_dict(fn: Callable[[T], Any], /) -> Callable[[dict[Key, T]], dict[Key, T]]: ...
@overload
def for_each_dict(fn: Callable[[T, Key], Any], /) -> Callable[[dict[Key, T]], dict[Key, T]]: ...
@overload
def for_each_dict(fn: Callable[[T, Key, dict[Key, T]], Any], /) -> Callable[[dict[Key, T]], dict[Key, T]]: ...


@make_data_last
def for_each_dict(
    data: dict[Key, T],
    fn: Callable[[T, Key, dict[Key, T]], Any] | Callable[[T, Key], Any] | Callable[[T], Any] | Callable[[], Any],
    /,
) -> dict[Key, T]:
    """
    Applies the given function to each key-value pair of the dict and then yields the element.

    Parameters
    ----------
    data: dict[Key, T]
        The dict to iterate over (positional-only).
    fn: Callable[[], Any] |
       Callable[[T], Any] |
       Callable[[T, Key], Any] |
       Callable[[T, Key, dict[Key, T]], Any]
        The function to apply to each key-value pair of the dict (positional-only).

    Returns
    -------
    dict[Key, T]
        The given dict.

    Examples
    --------
    Data first:
    >>> x = []
    >>> result = R.for_each_dict({'a': 1, 'b': 2, 'c': 3}, lambda v, k: x.append(k + str(v)))
    >>> x
    ['a1', 'b2', 'c3']
    >>> result
    {'a': 1, 'b': 2, 'c': 3}

    Data last:
    >>> x = []
    >>> result = R.for_each_dict(lambda v, k: x.append(k + str(v)))({'a': 1, 'b': 2, 'c': 3})
    >>> x
    ['a1', 'b2', 'c3']
    >>> result
    {'a': 1, 'b': 2, 'c': 3}

    """
    sig = inspect.signature(fn)
    num_params = len(sig.parameters)
    match num_params:
        case 0:
            fn = cast(Callable[[], Any], fn)
            for _, _ in data.items():
                fn()
        case 1:
            fn = cast(Callable[[T], Any], fn)
            for _, value in data.items():
                fn(value)
        case 2:
            fn = cast(Callable[[T, Key], Any], fn)
            for key, value in data.items():
                fn(value, key)
        case 3:
            fn = cast(Callable[[T, Key, dict[Key, T]], Any], fn)
            for key, value in data.items():
                fn(value, key, data)
        case _:  # pragma: no cover
            raise ValueError(f'Unsupported number of parameters: {num_params}')
    return data
