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')


@overload
def split_when(data: Iterable[T], predicate: Callable[[T], bool], /) -> tuple[list[T], list[T]]: ...


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


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


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


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


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


@make_data_last
def split_when(
    iterable: Iterable[T],
    predicate: Callable[[T], bool] | Callable[[T, int], bool] | Callable[[T, int, Sequence[T]], bool],
    /,
) -> tuple[list[T], list[T]]:
    r"""
    Splits the given iterable at the first element for whom the predicate returns True.

    Parameters
    ----------
    iterable: Iterable[T]
        Iterable to split (positional-only).
    predicate: Callable[[T], bool] | Callable[[T, int], bool] | Callable[[T, int, Sequence[T]], bool]
        Predicate to split by (positional-only).

    Returns
    -------
    tuple[list[T], list[T]]
        Tuple of two sequences.

    Examples
    --------
    Data first:
    >>> R.split_when([1, 2, 3], R.eq(2))
    ([1], [2, 3])
    >>> R.split_when([1, 2, 3, 4], lambda x, i: x % 2 == 0 and i > 1)
    ([1, 2, 3], [4])
    >>> R.split_when([1, 2, 3, 4], lambda x, i, arr: x % 2 == 0 and i > 1)
    ([1, 2, 3], [4])

    Data last:
    >>> R.split_when(R.eq(2))([1, 2, 3])
    ([1], [2, 3])

    """
    # TODO: refactor
    sig = inspect.signature(predicate)
    num_params = len(sig.parameters)
    result: tuple[list[T], list[T]] = ([], [])
    result_index = 0

    match num_params:
        case 1:
            predicate = cast(Callable[[T], bool], predicate)
            for item in iterable:
                if predicate(item):
                    result_index = 1
                result[result_index].append(item)
        case 2:
            predicate = cast(Callable[[T, int], bool], predicate)
            for i, item in enumerate(iterable):
                if predicate(item, i):
                    result_index = 1
                result[result_index].append(item)
        case 3:
            predicate = cast(Callable[[T, int, Sequence[T]], bool], predicate)
            data_list = list(iterable)  # TODO: make conditional and configurable
            for i, item in enumerate(data_list):
                if predicate(item, i, data_list):
                    result_index = 1
                result[result_index].append(item)
        case _:  # pragma: no cover
            raise ValueError(f'Unsupported number of parameters: {num_params}')
    return result
