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 flat_map(iterable: Iterable[T], callbackfn: Callable[[], Iterable[R]], /) -> Iterable[R]: ...


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


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


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


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


@make_data_last
def flat_map(
    iterable: Iterable[T],
    callbackfn: Callable[[], Iterable[R]]
    | Callable[[T], Iterable[R]]
    | Callable[[T, int], Iterable[R]]
    | Callable[[T, int, Sequence[T]], Iterable[R]],
    /,
) -> Iterable[R]:
    """
    Yields the elements resulting from applying given function to the given iterable's elements.

    Parameters
    ----------
    iterable: Iterable[T]
        The iterable to flatten (positional-only).
    callbackfn: Callable[[], Iterable[R]] |
                Callable[[T], Iterable[R]] |
                Callable[[T, int], Iterable[R]] |
                Callable[[T, int, Sequence[T]], Iterable[R]]
        The function to apply to each element of the iterable (positional-only).

    Yields
    ------
    T
        The flattened elements.

    Examples
    --------
    Data first:
    >>> list(R.flat_map([1, 2, 3], lambda x: [x, x * 10]))
    [1, 10, 2, 20, 3, 30]
    >>> list(R.flat_map([1, 2, 3], lambda: [0, 1]))
    [0, 1, 0, 1, 0, 1]
    >>> list(R.flat_map([1, 2, 3], lambda x, i: [x, x * 10 + i]))
    [1, 10, 2, 21, 3, 32]
    >>> list(R.flat_map([1, 2, 3], lambda x, i, arr: [x + i, x * 10 + len(arr)]))
    [1, 13, 3, 23, 5, 33]

    Data last:
    >>> R.pipe([1, 2, 3], R.flat_map(lambda x: [x, x * 10]), list)
    [1, 10, 2, 20, 3, 30]

    """
    sig = inspect.signature(callbackfn)
    num_params = len(sig.parameters)
    match num_params:
        case 0:
            callbackfn = cast(Callable[[], Iterable[R]], callbackfn)
            for _ in iterable:
                yield from callbackfn()
        case 1:
            callbackfn = cast(Callable[[T], Iterable[R]], callbackfn)
            for x in iterable:
                yield from callbackfn(x)
        case 2:
            callbackfn = cast(Callable[[T, int], Iterable[R]], callbackfn)
            for i, x in enumerate(iterable):
                yield from callbackfn(x, i)
        case 3:
            callbackfn = cast(Callable[[T, int, Iterable[T]], Iterable[R]], callbackfn)
            data_list = list(iterable)
            for i, x in enumerate(data_list):
                yield from callbackfn(x, i, data_list)
        case _:  # pragma: no cover
            raise ValueError(f'Unsupported number of parameters: {num_params}')
