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

from .decorator import make_data_last

T = TypeVar('T')


@overload
def take_while(it: Iterable[T], n: Callable[[T], bool], /) -> Iterable[T]: ...


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


@make_data_last
def take_while(iterable: Iterable[T], predicate: Callable[[T], bool], /) -> Iterable[T]:
    """
    Yields the elements of the iterable until encountering the element that does not satisfy the predicate.

    Doesn't yield the element that does not satisfy the predicate.

    Parameters
    ----------
    iterable : Iterable[T]
        Input iterable (positional-only).
    predicate: Callable[[T], bool]
        Predicate to check the elements of the iterable (positional-only).

    Returns
    -------
    Iterable[T]
        Elements of the iterable until encountering the element that does not satisfy the predicate.

    Examples
    --------
    Data first:
    >>> list(R.take_while([1, 2, 3, 4, 3, 2, 1], R.neq(4)))
    [1, 2, 3]

    Data last:
    >>> R.pipe([1, 2, 3, 4, 3, 2, 1], R.take_while(R.neq(4)), list)
    [1, 2, 3]

    """
    for x in iterable:
        if predicate(x):
            yield x
        else:
            break
