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

import remedapy as R
from remedapy.decorator import make_data_last

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


@overload
def set_path(data: dict[Key, T], path: list[int | str], v: Any, /) -> dict[Key, T]: ...


@overload
def set_path(path: list[int | str], v: Any, /) -> Callable[[dict[Key, T]], dict[Key, T]]: ...


def _set_path_helper(data: dict[Key, T], path: list[int | str], v: Any, /) -> Any:
    if not path:
        return v
    return data | {path[0]: _set_path_helper(R.get(data, path[0], {}), path[1:], v)}  # pyright: ignore[reportArgumentType, reportUnknownArgumentType]


@make_data_last
def set_path(data: dict[Key, T], path: list[int | str], value: Any, /) -> dict[Key, T]:
    """
    Returns a dict with the value set at the given path.

    Parameters
    ----------
    data: dict[K, V]
        The dict to set the value in.
    path: list[int | str]
        The path to set the value at.
    value: V
        The value to set.

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

    Examples
    --------
    Data first:
    >>> R.set_path({'a': {'b': 1}}, ['a', 'b'], 2)
    {'a': {'b': 2}}
    >>> R.set_path({'a': {'b': 1}}, ['a', 'c'], 2)
    {'a': {'b': 1, 'c': 2}}

    Data last:
    >>> R.pipe({'a': {'b': 1}}, R.set_path(['a', 'b'], 2))
    {'a': {'b': 2}}
    >>> R.pipe({'a': {'b': 1}}, R.set_path(['a', 'c'], 2))
    {'a': {'b': 1, 'c': 2}}

    """
    return _set_path_helper(data, path, value)
