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, heap_to_list, heapify

from .decorator import make_data_last

T = TypeVar('T')


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


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


@make_data_last
def take_first_by(iterable: Iterable[T], n: int, function: Callable[[T], SupportsGtLt], /) -> list[T]:
    """
    Yields the first n elements of the iterable, as if it were sorted by the function provided.

    They are yielded in the order they appear in the iterable.

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

    Yields
    ------
    T
        First n elements of the iterable if sorted by the function provided.

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

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

    """
    iterable_ = iter(iterable)
    original_heap: list[T] = []
    for _ in range(n):
        original_heap.append(next(iterable_))
    heap = heapify(original_heap, function)
    if heap is None:
        return []
    for i, x in enumerate(iterable_, n):
        heap = heap_maybe_insert(heap, LeaftistHeapNode(value=x, key=function(x), index=i))
    return heap_to_list(heap)
