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

from remedapy.decorator import make_data_last

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

T = TypeVar('T')


@overload
def set(data: dict[Key, T], prop: Key, v: T, /) -> dict[Key, T]: ...


@overload
def set(prop: Key, v: T, /) -> Callable[[dict[Key, T]], dict[Key, T]]: ...


@make_data_last
def set(data: dict[Key, T], key: Key, value: T, /) -> dict[Key, T]:
    """
    Returns a dict with the given key-value pair set.

    Parameters
    ----------
    data: dict[K, V]
        The dict to set the key-value pair in.
    key: K
        The key to set the value for.
    value: V
        The value to set for the key.

    Returns
    -------
    dict[K, V]
        The new dict with the key-value pair set.

    Examples
    --------
    Data first:
    >>> R.set({'a': 1}, 'a', 2)
    {'a': 2}
    >>> R.set({'a': 1}, 'b', 2)
    {'a': 1, 'b': 2}

    Data last:
    >>> R.pipe({'a': 1}, R.set('a', 2))
    {'a': 2}

    """
    return data | {key: value}
