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

from remedapy.local_types import SupportsGtLt
from remedapy.util import LeaftistHeapNode, heap_maybe_insert, heapify, indexes_in_heap

from .decorator import make_data_last

T = TypeVar('T')


@overload
def drop_first_by(it: Iterable[T], n: int, fn: Callable[[T], SupportsGtLt], /) -> Iterable[T]: ...


@overload
def drop_first_by(n: int, fn: Callable[[T], SupportsGtLt], /) -> Callable[[Iterable[T]], Iterable[T]]: ...


@make_data_last
def drop_first_by(iterable: Iterable[T], n: int, function: Callable[[T], SupportsGtLt], /) -> Iterable[T]:
    """
    Yields all elements of the iterable except the first* n.

    *Where "first" means first if the iterable were sorted by the function provided.

    The elements are yielded in the order they appear in the iterable.

    Parameters
    ----------
    iterable : Iterable[T]
        Input iterable (positional-only).
    n : int
        Number of elements to skip (positional-only).
    function: Callable[[T], SupportsGtLt]
        Function to "sort" the iterable by (positional-only).

    Returns
    -------
    Iterable[T]
        All elements of the iterable without the first* n.

    Examples
    --------
    Data first:
    >>> list(R.drop_first_by(['aa', 'aaaa', 'a', 'aaa'], 2, R.length))
    ['aaaa', 'aaa']

    Data last:
    >>> list(R.pipe(['aa', 'aaaa', 'a', 'aaa'], R.drop_first_by(2, R.length)))
    ['aaaa', 'aaa']

    """
    list_ = list(iterable)
    original_heap: list[T] = list_[:n]
    heap = heapify(original_heap, function)
    if heap is None:
        yield from list_
        return
    for i, x in enumerate(list_[n:], n):
        heap = heap_maybe_insert(heap, LeaftistHeapNode(value=x, key=function(x), index=i))
    indexes = indexes_in_heap(heap)
    yield from (x for i, x in enumerate(list_) if i not in indexes)
