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

from typing_extensions import TypeIs

from remedapy.decorator import make_data_last

T = TypeVar('T')
U = TypeVar('U')


# TypeIs overloads:
@overload
def filter(data: Iterable[T], predicate: Callable[[T], TypeIs[U]], /) -> Iterable[U]: ...


@overload
def filter(data: Iterable[T], predicate: Callable[[T, int], TypeIs[U]], /) -> Iterable[U]: ...


@overload
def filter(data: Iterable[T], predicate: Callable[[T, int, Sequence[T]], TypeIs[U]], /) -> Iterable[U]: ...


@overload
def filter(predicate: Callable[[T], TypeIs[U]], /) -> Callable[[Iterable[T]], Iterable[U]]: ...


@overload
def filter(predicate: Callable[[T, int], TypeIs[U]], /) -> Callable[[Iterable[T]], Iterable[U]]: ...


@overload
def filter(predicate: Callable[[T, int, Sequence[T]], TypeIs[U]], /) -> Callable[[Iterable[T]], Iterable[U]]: ...
# End of TypeIs overloads.
# TypeGuard overloads:
@overload
def filter(data: Iterable[T], predicate: Callable[[T], TypeGuard[U]], /) -> Iterable[U]: ...


@overload
def filter(data: Iterable[T], predicate: Callable[[T, int], TypeGuard[U]], /) -> Iterable[U]: ...


@overload
def filter(data: Iterable[T], predicate: Callable[[T, int, Sequence[T]], TypeGuard[U]], /) -> Iterable[U]: ...


@overload
def filter(predicate: Callable[[T], TypeGuard[U]], /) -> Callable[[Iterable[T]], Iterable[U]]: ...


@overload
def filter(predicate: Callable[[T, int], TypeGuard[U]], /) -> Callable[[Iterable[T]], Iterable[U]]: ...


@overload
def filter(predicate: Callable[[T, int, Sequence[T]], TypeGuard[U]], /) -> Callable[[Iterable[T]], Iterable[U]]: ...
# End of TypeGuard overloads.
@overload
def filter(data: Iterable[T], predicate: Callable[[T], bool], /) -> Iterable[T]: ...


@overload
def filter(data: Iterable[T], predicate: Callable[[T, int], bool], /) -> Iterable[T]: ...


@overload
def filter(data: Iterable[T], predicate: Callable[[T, int, Sequence[T]], bool], /) -> Iterable[T]: ...


@overload
def filter(predicate: Callable[[T], bool], /) -> Callable[[Iterable[T]], Iterable[T]]: ...


@overload
def filter(predicate: Callable[[T, int], bool], /) -> Callable[[Iterable[T]], Iterable[T]]: ...


@overload
def filter(predicate: Callable[[T, int, Sequence[T]], bool], /) -> Callable[[Iterable[T]], Iterable[T]]: ...


@make_data_last
def filter(
    iterable: Iterable[T],
    predicate: Callable[[T], bool] | Callable[[T, int], bool] | Callable[[T, int, Sequence[T]], bool],
    /,
) -> Iterable[T]:
    """
    Given an iterable and a predicate yields the elements satisfying the predicate.

    Predicate can appept a value, a value and index, or a value, index, and the whole sequence.

    If iterable is not a sequence it is collected into a list first.

    Parameters
    ----------
    iterable: Iterable[T]
        The iterable to partition.
    predicate: Callable[[T], bool] | Callable[[T, int], bool] | Callable[[T, int, Sequence[T]], bool]
        The predicate to use for filtering.

    Yields
    ------
    T
        Elements satisfying the predicate.

    See Also
    --------
    partition

    Examples
    --------
    Data first:
    >>> list(R.filter([1, 2, 3], lambda x: x % 2 == 1))
    [1, 3]

    Data last:
    >>> R.pipe([1, 2, 3], R.filter(R.is_odd), list)
    [1, 3]

    """
    sig = inspect.signature(predicate)
    num_params = len(sig.parameters)
    match num_params:
        case 1:
            predicate = cast(Callable[[T], bool], predicate)
            return (item for item in iterable if predicate(item))
        case 2:
            predicate = cast(Callable[[T, int], bool], predicate)
            return (item for index, item in enumerate(iterable) if predicate(item, index))
        case 3:
            predicate = cast(Callable[[T, int, Sequence[T]], bool], predicate)
            data_list = list(iterable)  # TODO: make conditional and configurable

            return (item for index, item in enumerate(data_list) if predicate(item, index, data_list))
        case _:  # pragma: no cover
            raise ValueError(f'Unsupported number of parameters: {num_params}')
