"""Selection operators."""

from __future__ import annotations

import asyncio
import builtins
import collections

from typing import Awaitable, Callable, TypeVar, AsyncIterable, AsyncIterator

from . import transform
from ..aiter_utils import aiter, anext
from ..core import streamcontext, pipable_operator

__all__ = [
    "take",
    "takelast",
    "skip",
    "skiplast",
    "getitem",
    "filter",
    "until",
    "dropwhile",
    "takewhile",
]

T = TypeVar("T")


@pipable_operator
async def take(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
    """Forward the first ``n`` elements from an asynchronous sequence.

    If ``n`` is negative, it simply terminates before iterating the source.
    """
    enumerated = transform.enumerate.raw(source)
    async with streamcontext(enumerated) as streamer:
        if n <= 0:
            return
        async for i, item in streamer:
            yield item
            if i >= n - 1:
                return


@pipable_operator
async def takelast(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
    """Forward the last ``n`` elements from an asynchronous sequence.

    If ``n`` is negative, it simply terminates after iterating the source.

    Note: it is required to reach the end of the source before the first
    element is generated.
    """
    queue: collections.deque[T] = collections.deque(maxlen=n if n > 0 else 0)
    async with streamcontext(source) as streamer:
        async for item in streamer:
            queue.append(item)
        for item in queue:
            yield item


@pipable_operator
async def skip(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
    """Forward an asynchronous sequence, skipping the first ``n`` elements.

    If ``n`` is negative, no elements are skipped.
    """
    enumerated = transform.enumerate.raw(source)
    async with streamcontext(enumerated) as streamer:
        async for i, item in streamer:
            if i >= n:
                yield item


@pipable_operator
async def skiplast(source: AsyncIterable[T], n: int) -> AsyncIterator[T]:
    """Forward an asynchronous sequence, skipping the last ``n`` elements.

    If ``n`` is negative, no elements are skipped.

    Note: it is required to reach the ``n+1`` th element of the source
    before the first element is generated.
    """
    queue: collections.deque[T] = collections.deque(maxlen=n if n > 0 else 0)
    async with streamcontext(source) as streamer:
        async for item in streamer:
            if n <= 0:
                yield item
                continue
            if len(queue) == n:
                yield queue[0]
            queue.append(item)


@pipable_operator
async def filterindex(
    source: AsyncIterable[T], func: Callable[[int], bool]
) -> AsyncIterator[T]:
    """Filter an asynchronous sequence using the index of the elements.

    The given function is synchronous, takes the index as an argument,
    and returns ``True`` if the corresponding should be forwarded,
    ``False`` otherwise.
    """
    enumerated = transform.enumerate.raw(source)
    async with streamcontext(enumerated) as streamer:
        async for i, item in streamer:
            if func(i):
                yield item


@pipable_operator
def slice(source: AsyncIterable[T], *args: int | None) -> AsyncIterator[T]:
    """Slice an asynchronous sequence.

    The arguments are the same as the builtin type slice.

    There are two limitations compare to regular slices:
    - Positive stop index with negative start index is not supported
    - Negative step is not supported
    """
    s = builtins.slice(*args)
    start, stop, step = s.start or 0, s.stop, s.step or 1
    aiterator = aiter(source)
    # Filter the first items
    if start < 0:
        aiterator = takelast.raw(aiterator, abs(start))
    elif start > 0:
        aiterator = skip.raw(aiterator, start)
    # Filter the last items
    if stop is not None:
        if stop >= 0 and start < 0:
            raise ValueError("Positive stop with negative start is not supported")
        elif stop >= 0:
            aiterator = take.raw(aiterator, stop - start)
        else:
            aiterator = skiplast.raw(aiterator, abs(stop))
    # Filter step items
    if step is not None:
        if step > 1:
            aiterator = filterindex.raw(aiterator, lambda i: i % step == 0)
        elif step < 0:
            raise ValueError("Negative step not supported")
    # Return
    return aiterator


@pipable_operator
async def item(source: AsyncIterable[T], index: int) -> AsyncIterator[T]:
    """Forward the ``n``th element of an asynchronous sequence.

    The index can be negative and works like regular indexing.
    If the index is out of range, and ``IndexError`` is raised.
    """
    # Prepare
    if index >= 0:
        source = skip.raw(source, index)
    else:
        source = takelast(source, abs(index))
    async with streamcontext(source) as streamer:
        # Get first item
        try:
            result = await anext(streamer)
        except StopAsyncIteration:
            raise IndexError("Index out of range")
        # Check length
        if index < 0:
            count = 1
            async for _ in streamer:
                count += 1
            if count != abs(index):
                raise IndexError("Index out of range")
        # Yield result
        yield result


@pipable_operator
def getitem(source: AsyncIterable[T], index: int | builtins.slice) -> AsyncIterator[T]:
    """Forward one or several items from an asynchronous sequence.

    The argument can either be a slice or an integer.
    See the slice and item operators for more information.
    """
    if isinstance(index, builtins.slice):
        return slice.raw(source, index.start, index.stop, index.step)
    if isinstance(index, int):
        return item.raw(source, index)
    raise TypeError("Not a valid index (int or slice)")


@pipable_operator
async def filter(
    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
    """Filter an asynchronous sequence using an arbitrary function.

    The function takes the item as an argument and returns ``True``
    if it should be forwarded, ``False`` otherwise.
    The function can either be synchronous or asynchronous.
    """
    iscorofunc = asyncio.iscoroutinefunction(func)
    async with streamcontext(source) as streamer:
        async for item in streamer:
            result = func(item)
            if iscorofunc:
                assert isinstance(result, Awaitable)
                result = await result
            if result:
                yield item


@pipable_operator
async def until(
    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
    """Forward an asynchronous sequence until a condition is met.

    Contrary to the ``takewhile`` operator, the last tested element is included
    in the sequence.

    The given function takes the item as an argument and returns a boolean
    corresponding to the condition to meet. The function can either be
    synchronous or asynchronous.
    """
    iscorofunc = asyncio.iscoroutinefunction(func)
    async with streamcontext(source) as streamer:
        async for item in streamer:
            result = func(item)
            if iscorofunc:
                assert isinstance(result, Awaitable)
                result = await result
            yield item
            if result:
                return


@pipable_operator
async def takewhile(
    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
    """Forward an asynchronous sequence while a condition is met.

    Contrary to the ``until`` operator, the last tested element is not included
    in the sequence.

    The given function takes the item as an argument and returns a boolean
    corresponding to the condition to meet. The function can either be
    synchronous or asynchronous.
    """
    iscorofunc = asyncio.iscoroutinefunction(func)
    async with streamcontext(source) as streamer:
        async for item in streamer:
            result = func(item)
            if iscorofunc:
                assert isinstance(result, Awaitable)
                result = await result
            if not result:
                return
            yield item


@pipable_operator
async def dropwhile(
    source: AsyncIterable[T], func: Callable[[T], bool | Awaitable[bool]]
) -> AsyncIterator[T]:
    """Discard the elements from an asynchronous sequence
    while a condition is met.

    The given function takes the item as an argument and returns a boolean
    corresponding to the condition to meet. The function can either be
    synchronous or asynchronous.
    """
    iscorofunc = asyncio.iscoroutinefunction(func)
    async with streamcontext(source) as streamer:
        async for item in streamer:
            result = func(item)
            if iscorofunc:
                assert isinstance(result, Awaitable)
                result = await result
            if not result:
                yield item
                break
        async for item in streamer:
            yield item
