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

from remedapy.decorator import make_data_last

T = TypeVar('T')
R = TypeVar('R')


@overload
def map(data: Iterable[T], callbackfn: Callable[[], R], /) -> Iterable[R]: ...


@overload
def map(data: Iterable[T], callbackfn: Callable[[T], R], /) -> Iterable[R]: ...


@overload
def map(data: Iterable[T], callbackfn: Callable[[T, int], R], /) -> Iterable[R]: ...


@overload
def map(data: Iterable[T], callbackfn: Callable[[T, int, Sequence[T]], R], /) -> Iterable[R]: ...


@overload
def map(
    callbackfn: Callable[[T], R] | Callable[[T, int], R] | Callable[[T, int, Sequence[T]], R],
    /,
) -> Callable[[Iterable[T]], Iterable[R]]: ...


@make_data_last
def map(
    iterable: Iterable[T],
    function: Callable[[], R] | Callable[[T], R] | Callable[[T, int], R] | Callable[[T, int, Sequence[T]], R],
    /,
) -> Iterable[R]:
    """
    Maps iterable with a function.

    Parameters
    ----------
    iterable: Iterable[T]
        The iterable to map.
    function: Callable[[T], R] | Callable[[T, int], R] | Callable[[T, int, Sequence[T]], R]
        The function to apply to each element.

    Yields
    ------
    R
        Elements satisfying the predicate.

    Examples
    --------
    Data first:
    >>> list(R.map([1, 2, 3], R.multiply(2)))
    [2, 4, 6]
    >>> list(R.map([0, 0], R.add(1)))
    [1, 1]
    >>> list(R.map([0, 0], lambda value, index: value + index))
    [0, 1]

    Data last:
    >>> list(R.map(R.multiply(2))([1, 2, 3]))
    [2, 4, 6]
    >>> list(R.map(R.add(1))([0, 0]))
    [1, 1]
    >>> R.pipe([1, 2, 3], R.map(R.multiply(2)), list)
    [2, 4, 6]
    >>> R.pipe([0, 0], R.map(R.add(1)), list)
    [1, 1]

    """
    # TODO: create a decorator to handle caching this case
    sig = inspect.signature(function)
    num_params = len(sig.parameters)
    match num_params:
        case 0:
            function = cast(Callable[[], R], function)
            return (function() for _ in iterable)
        case 1:
            function = cast(Callable[[T], R], function)
            return (function(item) for item in iterable)
        case 2:
            function = cast(Callable[[T, int], R], function)
            return (function(item, index) for index, item in enumerate(iterable))
        case 3:
            function = cast(Callable[[T, int, Iterable[T]], R], function)
            data_list = list(iterable)
            return (function(item, index, data_list) for index, item in enumerate(data_list))
        case _:  # pragma: no cover
            raise ValueError(f'Unsupported number of parameters: {num_params}')
