# AUTOGENERATED - modify shared_anyscale_util in root directory to make changes
import asyncio
from asyncio import Future, isfuture
import logging
from typing import Awaitable, List, Optional, TypeVar, Union


_T = TypeVar("_T")
logger = logging.getLogger(__file__)


def create_future() -> Future:
    """Creates instance of `asyncio.Future` using the currently running event-loop"""
    loop = asyncio.get_running_loop()
    return loop.create_future()


def chain_future(source, destination):
    """NOTE: THIS METHOD IS COPIED FROM STDLIB VERBATIM, WITH NO MODIFICATIONS, PLEASE KEEP IT THAT WAY.
             THIS WAS COPIED BECAUSE ITS PRIVATE IN STANDARD LIBRARY MAKING IT IMPOSSIBLE TO USE IT
             OUTSIDE OF IT.

    Chain two futures so that when one completes, so does the other.

    The result (or exception) of source will be copied to destination.
    If destination is cancelled, source gets cancelled too.
    Compatible only with `asyncio.Future`
    """
    if not isfuture(source):
        raise TypeError("A future is required for source argument")
    if not isfuture(destination):
        raise TypeError("A future is required for destination argument")
    source_loop = _get_loop(source) if isfuture(source) else None
    dest_loop = _get_loop(destination) if isfuture(destination) else None

    def _set_state(future, other):
        _copy_future_state(other, future)

    def _call_check_cancel(destination):
        if destination.cancelled():
            if source_loop is None or source_loop is dest_loop:
                source.cancel()
            else:
                source_loop.call_soon_threadsafe(source.cancel)

    def _call_set_state(source):
        if destination.cancelled() and dest_loop is not None and dest_loop.is_closed():
            return
        if dest_loop is None or dest_loop is source_loop:
            _set_state(destination, source)
        else:
            dest_loop.call_soon_threadsafe(_set_state, destination, source)

    destination.add_done_callback(_call_check_cancel)
    source.add_done_callback(_call_set_state)


def _get_loop(fut):
    """NOTE: THIS METHOD IS COPIED FROM STDLIB VERBATIM, WITH NO MODIFICATIONS, PLEASE KEEP IT THAT WAY.
             THIS WAS COPIED BECAUSE ITS PRIVATE IN STANDARD LIBRARY MAKING IT IMPOSSIBLE TO USE IT
             OUTSIDE OF IT.
    """
    # Tries to call Future.get_loop() if it's available.
    # Otherwise fallbacks to using the old '_loop' property.
    try:
        get_loop = fut.get_loop
    except AttributeError:
        pass
    else:
        return get_loop()
    return fut._loop  # noqa: SLF001


def _copy_future_state(source, dest):
    """NOTE: THIS METHOD IS COPIED FROM STDLIB VERBATIM, WITH NO MODIFICATIONS, PLEASE KEEP IT THAT WAY.
             THIS WAS COPIED BECAUSE ITS PRIVATE IN STANDARD LIBRARY MAKING IT IMPOSSIBLE TO USE IT
             OUTSIDE OF IT.

    Internal helper to copy state from another Future.

    The other Future may be a concurrent.futures.Future.
    """
    assert source.done()
    if dest.cancelled():
        return
    assert not dest.done()
    if source.cancelled():
        dest.cancel()
    else:
        exception = source.exception()
        if exception is not None:
            dest.set_exception(exception)
        else:
            result = source.result()
            dest.set_result(result)


async def gather_in_batches(
    n: int,
    *tasks: Awaitable[_T],
    per_task_timeout_secs: Optional[float] = None,
    return_exceptions: bool = False,
    log_exceptions: bool = False,
) -> List[Union[_T, BaseException]]:
    # n cannot be less than 1
    n = max(n, 1)
    semaphore = asyncio.Semaphore(n)

    async def sem_task(task: Awaitable[_T]) -> _T:
        try:
            async with semaphore:
                return await asyncio.wait_for(task, per_task_timeout_secs)
        except Exception:
            if log_exceptions:
                logger.exception("Error while gathering task")
            raise  # re-raise exception

    return await asyncio.gather(
        *(sem_task(task) for task in tasks), return_exceptions=return_exceptions
    )
