"""Combination operators."""
from __future__ import annotations

import asyncio
import builtins

from typing import (
    Awaitable,
    Protocol,
    TypeVar,
    AsyncIterable,
    AsyncIterator,
    Callable,
    cast,
)
from typing_extensions import ParamSpec

from ..aiter_utils import AsyncExitStack, anext
from ..core import streamcontext, pipable_operator

from . import create
from . import select
from . import advanced
from . import aggregate

__all__ = ["chain", "zip", "map", "merge", "ziplatest", "amap", "smap"]

T = TypeVar("T")
U = TypeVar("U")
K = TypeVar("K")
P = ParamSpec("P")


@pipable_operator
async def chain(
    source: AsyncIterable[T], *more_sources: AsyncIterable[T]
) -> AsyncIterator[T]:
    """Chain asynchronous sequences together, in the order they are given.

    Note: the sequences are not iterated until it is required,
    so if the operation is interrupted, the remaining sequences
    will be left untouched.
    """
    sources = source, *more_sources
    for source in sources:
        async with streamcontext(source) as streamer:
            async for item in streamer:
                yield item


@pipable_operator
async def zip(
    source: AsyncIterable[T], *more_sources: AsyncIterable[T]
) -> AsyncIterator[tuple[T, ...]]:
    """Combine and forward the elements of several asynchronous sequences.

    Each generated value is a tuple of elements, using the same order as
    their respective sources. The generation continues until the shortest
    sequence is exhausted.

    Note: the different sequences are awaited in parrallel, so that their
    waiting times don't add up.
    """
    sources = source, *more_sources

    # One sources
    if len(sources) == 1:
        (source,) = sources
        async with streamcontext(source) as streamer:
            async for item in streamer:
                yield (item,)
        return

    # N sources
    async with AsyncExitStack() as stack:
        # Handle resources
        streamers = [
            await stack.enter_async_context(streamcontext(source)) for source in sources
        ]
        # Loop over items
        while True:
            try:
                coros = builtins.map(anext, streamers)
                items = await asyncio.gather(*coros)
            except StopAsyncIteration:
                break
            else:
                yield tuple(items)


X = TypeVar("X", contravariant=True)
Y = TypeVar("Y", covariant=True)


class SmapCallable(Protocol[X, Y]):
    def __call__(self, arg: X, /, *args: X) -> Y:
        ...


class AmapCallable(Protocol[X, Y]):
    async def __call__(self, arg: X, /, *args: X) -> Y:
        ...


class MapCallable(Protocol[X, Y]):
    def __call__(self, arg: X, /, *args: X) -> Awaitable[Y] | Y:
        ...


@pipable_operator
async def smap(
    source: AsyncIterable[T],
    func: SmapCallable[T, U],
    *more_sources: AsyncIterable[T],
) -> AsyncIterator[U]:
    """Apply a given function to the elements of one or several
    asynchronous sequences.

    Each element is used as a positional argument, using the same order as
    their respective sources. The generation continues until the shortest
    sequence is exhausted. The function is treated synchronously.

    Note: if more than one sequence is provided, they're awaited concurrently
    so that their waiting times don't add up.
    """
    stream = zip(source, *more_sources)
    async with streamcontext(stream) as streamer:
        async for item in streamer:
            yield func(*item)


@pipable_operator
def amap(
    source: AsyncIterable[T],
    corofn: AmapCallable[T, U],
    *more_sources: AsyncIterable[T],
    ordered: bool = True,
    task_limit: int | None = None,
) -> AsyncIterator[U]:
    """Apply a given coroutine function to the elements of one or several
    asynchronous sequences.

    Each element is used as a positional argument, using the same order as
    their respective sources. The generation continues until the shortest
    sequence is exhausted.

    The results can either be returned in or out of order, depending on
    the corresponding ``ordered`` argument.

    The coroutines run concurrently but their amount can be limited using
    the ``task_limit`` argument. A value of ``1`` will cause the coroutines
    to run sequentially.

    If more than one sequence is provided, they're also awaited concurrently,
    so that their waiting times don't add up.
    """

    async def func(arg: T, *args: T) -> AsyncIterable[U]:
        yield await corofn(arg, *args)

    if ordered:
        return advanced.concatmap.raw(
            source, func, *more_sources, task_limit=task_limit
        )
    return advanced.flatmap.raw(source, func, *more_sources, task_limit=task_limit)


@pipable_operator
def map(
    source: AsyncIterable[T],
    func: MapCallable[T, U],
    *more_sources: AsyncIterable[T],
    ordered: bool = True,
    task_limit: int | None = None,
) -> AsyncIterator[U]:
    """Apply a given function to the elements of one or several
    asynchronous sequences.

    Each element is used as a positional argument, using the same order as
    their respective sources. The generation continues until the shortest
    sequence is exhausted. The function can either be synchronous or
    asynchronous (coroutine function).

    The results can either be returned in or out of order, depending on
    the corresponding ``ordered`` argument. This argument is ignored if the
    provided function is synchronous.

    The coroutines run concurrently but their amount can be limited using
    the ``task_limit`` argument. A value of ``1`` will cause the coroutines
    to run sequentially. This argument is ignored if the provided function
    is synchronous.

    If more than one sequence is provided, they're also awaited concurrently,
    so that their waiting times don't add up.

    It might happen that the provided function returns a coroutine but is not
    a coroutine function per se. In this case, one can wrap the function with
    ``aiostream.async_`` in order to force ``map`` to await the resulting
    coroutine. The following example illustrates the use ``async_`` with a
    lambda function::

        from aiostream import stream, async_
        ...
        ys = stream.map(xs, async_(lambda ms: asyncio.sleep(ms / 1000)))
    """
    if asyncio.iscoroutinefunction(func):
        return amap.raw(
            source, func, *more_sources, ordered=ordered, task_limit=task_limit
        )
    sync_func = cast("SmapCallable[T, U]", func)
    return smap.raw(source, sync_func, *more_sources)


@pipable_operator
def merge(
    source: AsyncIterable[T], *more_sources: AsyncIterable[T]
) -> AsyncIterator[T]:
    """Merge several asynchronous sequences together.

    All the sequences are iterated simultaneously and their elements
    are forwarded as soon as they're available. The generation continues
    until all the sequences are exhausted.
    """
    sources = [source, *more_sources]
    source_stream: AsyncIterable[AsyncIterable[T]] = create.iterate.raw(sources)
    return advanced.flatten.raw(source_stream)


@pipable_operator
def ziplatest(
    source: AsyncIterable[T],
    *more_sources: AsyncIterable[T],
    partial: bool = True,
    default: T | None = None,
) -> AsyncIterator[tuple[T | None, ...]]:
    """Combine several asynchronous sequences together, producing a tuple with
    the lastest element of each sequence whenever a new element is received.

    The value to use when a sequence has not procuded any element yet is given
    by the ``default`` argument (defaulting to ``None``).

    The producing of partial results can be disabled by setting the optional
    argument ``partial`` to ``False``.

    All the sequences are iterated simultaneously and their elements
    are forwarded as soon as they're available. The generation continues
    until all the sequences are exhausted.
    """
    sources = source, *more_sources
    n = len(sources)

    # Custom getter
    def getter(dct: dict[int, T]) -> Callable[[int], T | None]:
        return lambda key: dct.get(key, default)

    # Add source index to the items
    def make_func(i: int) -> SmapCallable[T, dict[int, T]]:
        def func(x: T, *_: object) -> dict[int, T]:
            return {i: x}

        return func

    new_sources = [smap.raw(source, make_func(i)) for i, source in enumerate(sources)]

    # Merge the sources
    merged = merge.raw(*new_sources)

    # Accumulate the current state in a dict
    accumulated = aggregate.accumulate.raw(merged, lambda x, e: {**x, **e})

    # Filter partial result
    filtered = (
        accumulated
        if partial
        else select.filter.raw(accumulated, lambda x: len(x) == n)
    )

    # Convert the state dict to a tuple
    def dict_to_tuple(x: dict[int, T], *_: object) -> tuple[T | None, ...]:
        return tuple(builtins.map(getter(x), range(n)))

    return smap.raw(filtered, dict_to_tuple)