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 partition(data: Iterable[T], predicate: Callable[[T], bool], /) -> tuple[list[T], list[T]]: ...


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


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


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


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


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


@make_data_last
def partition(
    iterable: Iterable[T],
    predicate: Callable[[T], bool] | Callable[[T, int], bool] | Callable[[T, int, Sequence[T]], bool],
    /,
) -> tuple[list[T], list[T]]:
    """
    Given an iterable and a predicate returns 2 lists: one with items that satisfy the predicate and one with the rest.

    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 partitioning.

    Returns
    -------
    tuple[list[T], list[T]]
        A tuple of two lists: the first with items that satisfy the predicate and the second with items that don't.

    See Also
    --------
    filter

    Examples
    --------
    Data first:
    >>> R.partition(
    ...     ['one', 'two', 'forty two'],
    ...     lambda x: len(x) == 3,
    ... )
    (['one', 'two'], ['forty two'])
    >>> R.partition(['one', 'two', 'forty two'], lambda x, i: len(x) == 3 and i % 2 == 0)
    (['one'], ['two', 'forty two'])
    >>> R.partition(['one', 'two', 'forty two'], lambda x, i, data: len(x) == 3 and i % 2 == 0 and data[0] != x)
    ([], ['one', 'two', 'forty two'])

    Data last:
    >>> R.pipe(['one', 'two', 'forty two'], R.partition(lambda x: len(x) == 3))
    (['one', 'two'], ['forty two'])

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