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

from .decorator import make_data_last

T = TypeVar('T')


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


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


def sorted_index_with_binary_search(data: Sequence[T], predicate: Callable[[T], bool], start: int, end: int, /) -> int:
    if start == end:
        return start
    mid = (start + end) // 2
    if predicate(data[mid]):
        return sorted_index_with_binary_search(data, predicate, mid + 1, end)
    return sorted_index_with_binary_search(data, predicate, start, mid)


@make_data_last
def sorted_index_with(data: Sequence[T], predicate: Callable[[T], bool], /) -> int:
    """
    Given a sorted sequence and a predicate, returns the index where the predicate becomes false.

    "sorted" means that elements for whom the predicate is true are first in the sequence.

    The result is also the number of elements for whom the predicate is true.

    Parameters
    ----------
    data: Sequence[T]
        The data.
    predicate: Callable[[T], bool]
        Predicate to check.

    Returns
    -------
    int
        The index. Will be non-negative.

    Examples
    --------
    Data first:
    >>> R.sorted_index_with(['a', 'ab', 'abc'], R.piped(R.length, R.lt(2)))
    1
    >>> R.sorted_index_with(['a', 'ab', 'abc'], R.piped(R.length, R.lt(5)))
    3
    >>> R.sorted_index_with(['a', 'ab', 'abc'], R.piped(R.length, R.gt(0)))
    3
    >>> R.sorted_index_with(['a', 'ab', 'abc'], R.piped(R.length, R.lt(0)))
    0

    Data last:
    >>> R.pipe(['a', 'ab', 'abc'], R.sorted_index_with(R.piped(R.length, R.lt(2))))
    1

    """
    return sorted_index_with_binary_search(data, predicate, 0, len(data))
