Module livekit.agents.utils.aio

async def cancel_and_wait(*futures: asyncio.Future[Any]) -> None:
    loop = asyncio.get_running_loop()
    waiters = []

    for fut in futures:
        waiter = loop.create_future()
        cb = functools.partial(_release_waiter, waiter)
        waiters.append((waiter, cb))
        fut.add_done_callback(cb)
        fut.cancel()

    try:
        for waiter, _ in waiters:
            await waiter
    finally:
        for i, fut in enumerate(futures):
            _, cb = waiters[i]
            fut.remove_done_callback(cb)

async def cancel_and_wait(*futures: asyncio.Future[Any]) -> None:
    loop = asyncio.get_running_loop()
    waiters = []

    for fut in futures:
        waiter = loop.create_future()
        cb = functools.partial(_release_waiter, waiter)
        waiters.append((waiter, cb))
        fut.add_done_callback(cb)
        fut.cancel()

    try:
        for waiter, _ in waiters:
            await waiter
    finally:
        for i, fut in enumerate(futures):
            _, cb = waiters[i]
            fut.remove_done_callback(cb)

def interval(interval: float) -> Interval:
    return Interval(interval)

def sleep(delay: float) -> Sleep:
    return Sleep(delay)

class AsyncAtomicCounter:
    """Async atomic counter implementation."""

    def __init__(self, initial: int = 0):
        self._value = initial
        self._lock = asyncio.Lock()

    async def increment(self, n: int = 1) -> int:
        async with self._lock:
            self._value += n
            return self._value

    async def decrement(self, n: int = 1) -> int:
        async with self._lock:
            self._value -= n
            return self._value

    async def get(self) -> int:
        async with self._lock:
            return self._value

    async def set(self, value: int) -> None:
        async with self._lock:
            self._value = value

    async def compare_and_swap(self, expected: int, new: int) -> bool:
        async with self._lock:
            if self._value == expected:
                self._value = new
                return True
            return False

    async def get_and_reset(self, reset_value: int = 0) -> int:
        """Atomically read the current value and reset it."""
        async with self._lock:
            prev = self._value
            self._value = reset_value
            return prev

    def get_nowait(self) -> int:
        """Best-effort non-async read — safe if no await between read and use."""
        return self._value

class Chan(Generic[T]):
    def __init__(
        self,
        maxsize: int = 0,
        loop: asyncio.AbstractEventLoop | None = None,
    ) -> None:
        self._loop = loop or asyncio.get_event_loop()
        self._maxsize = max(maxsize, 0)
        #        self._finished_ev = asyncio.Event()
        self._close_ev = asyncio.Event()
        self._closed = False
        self._gets: deque[asyncio.Future[T | None]] = deque()
        self._puts: deque[asyncio.Future[T | None]] = deque()
        self._queue: deque[T] = deque()

    def _wakeup_next(self, waiters: deque[asyncio.Future[T | None]]) -> None:
        while waiters:
            waiter = waiters.popleft()
            if not waiter.done():
                waiter.set_result(None)
                break

    async def send(self, value: T) -> None:
        while self.full() and not self._close_ev.is_set():
            p = self._loop.create_future()
            self._puts.append(p)
            try:
                await p
            except ChanClosed:
                raise
            except:
                p.cancel()
                with contextlib.suppress(ValueError):
                    self._puts.remove(p)

                if not self.full() and not p.cancelled():
                    self._wakeup_next(self._puts)
                raise

        self.send_nowait(value)

    def send_nowait(self, value: T) -> None:
        if self._close_ev.is_set():
            raise ChanClosed

        if self.full():
            raise ChanFull

        self._queue.append(value)
        self._wakeup_next(self._gets)

    async def recv(self) -> T:
        while self.empty() and not self._close_ev.is_set():
            g = self._loop.create_future()
            self._gets.append(g)

            try:
                await g
            except ChanClosed:
                raise
            except BaseException:
                g.cancel()
                with contextlib.suppress(ValueError):
                    self._gets.remove(g)

                if not self.empty() and not g.cancelled():
                    self._wakeup_next(self._gets)

                raise

        return self.recv_nowait()

    def recv_nowait(self) -> T:
        if self.empty():
            if self._close_ev.is_set():
                raise ChanClosed
            else:
                raise ChanEmpty
        item = self._queue.popleft()
        #        if self.empty() and self._close_ev.is_set():
        #            self._finished_ev.set()
        self._wakeup_next(self._puts)
        return item

    def close(self) -> None:
        self._closed = True
        self._close_ev.set()
        while self._puts:
            putter = self._puts.popleft()
            if not putter.done():
                putter.set_exception(ChanClosed())

        while len(self._gets) > self.qsize():
            getter = self._gets.pop()
            if not getter.cancelled():
                getter.set_exception(ChanClosed())

        while self._gets:
            self._wakeup_next(self._gets)

    #        if self.empty():
    #            self._finished_ev.set()

    @property
    def closed(self) -> bool:
        return self._closed

    #    async def join(self) -> None:
    #        await self._finished_ev.wait()

    def qsize(self) -> int:
        """the number of elements queued (unread) in the channel buffer"""
        return len(self._queue)

    def full(self) -> bool:
        if self._maxsize <= 0:
            return False
        else:
            return self.qsize() >= self._maxsize

    def empty(self) -> bool:
        return not self._queue

    def __aiter__(self) -> AsyncIterator[T]:
        return self

    async def __anext__(self) -> T:
        try:
            return await self.recv()
        except ChanClosed:
            raise StopAsyncIteration from None

class ChanClosed(Exception):
    pass

class ChanReceiver(Protocol[T_co]):
    async def recv(self) -> T_co: ...

    def recv_nowait(self) -> T_co: ...

    def close(self) -> None: ...

    def __aiter__(self) -> AsyncIterator[T_co]: ...

    async def __anext__(self) -> T_co: ...

class ChanSender(Protocol[T_contra]):
    async def send(self, value: T_contra) -> None: ...

    def send_nowait(self, value: T_contra) -> None: ...

    def close(self) -> None: ...

class Interval:
    def __init__(self, interval: float) -> None:
        self._interval = interval
        self._last_sleep = 0.0
        self._i = 0
        self._handler: asyncio.TimerHandle | None = None
        self._fut: asyncio.Future[Any] | None = None

    def reset(self) -> None:
        if self._fut and self._handler and not self._handler.cancelled():
            self._handler.cancel()
            loop = asyncio.get_event_loop()
            self._handler = loop.call_later(self._interval, _finish_fut, self._fut)
        else:
            self._last_sleep = 0

    async def tick(self) -> int:
        loop = asyncio.get_event_loop()

        if self._last_sleep:
            self._fut = loop.create_future()
            delay = self._last_sleep - loop.time() + self._interval
            self._handler = loop.call_later(delay, _finish_fut, self._fut)
            try:
                await self._fut
            finally:
                self._handler.cancel()
            self._i += 1

        self._last_sleep = loop.time()
        return self._i

    def __aiter__(self) -> Interval:
        return self

    async def __anext__(self) -> int:
        return await self.tick()

class Sleep:
    """Same as asyncio.sleep except it is resettable"""

    def __init__(self, delay: float) -> None:
        self._delay = delay
        self._handler: asyncio.TimerHandle | None = None

    def reset(self, new_delay: float | None = None) -> None:
        if new_delay is None:
            new_delay = self._delay

        self._delay = new_delay

        if self._handler is None:
            return

        if self._handler.cancelled() or self._fut.done():
            raise SleepFinished

        self._handler.cancel()
        loop = asyncio.get_event_loop()
        self._handler = loop.call_later(new_delay, _finish_fut, self._fut)

    def cancel(self) -> None:
        if self._handler is None:
            return

        self._handler.cancel()
        self._fut.cancel()

    async def _sleep(self) -> None:
        if self._delay <= 0:
            self._fut = asyncio.Future[None]()
            self._fut.set_result(None)
            return

        loop = asyncio.get_event_loop()
        self._fut = loop.create_future()
        self._handler = loop.call_later(self._delay, _finish_fut, self._fut)

        try:
            await asyncio.shield(self._fut)
        finally:
            self._handler.cancel()

    def __await__(self) -> Generator[Any, Any, None]:
        return self._sleep().__await__()

class SleepFinished(Exception):
    pass

class TaskSet:
    """Small utility to create tasks in a fire-and-forget fashion."""

    def __init__(self, loop: asyncio.AbstractEventLoop | None = None) -> None:
        self._loop = loop or asyncio.get_event_loop()
        self._set = set[asyncio.Task[Any]]()
        self._closed = False

    def create_task(
        self, coro: Coroutine[Any, Any, _T], name: str | None = None
    ) -> asyncio.Task[_T]:
        if self._closed:
            raise RuntimeError("TaskSet is closed")

        task = self._loop.create_task(coro, name=name)
        self._set.add(task)
        task.add_done_callback(self._set.remove)
        return task

    @property
    def tasks(self) -> set[asyncio.Task[Any]]:
        return self._set.copy()

class WaitGroup:
    """
    asyncio wait group implementation (similar to sync.WaitGroup in go)
    """

    def __init__(self) -> None:
        self._counter = 0
        self._zero_event = asyncio.Event()
        self._zero_event.set()

    def add(self, delta: int = 1) -> None:
        new_value = self._counter + delta
        if new_value < 0:
            raise ValueError("WaitGroup counter cannot go negative.")

        self._counter = new_value

        if self._counter == 0:
            self._zero_event.set()
        else:
            self._zero_event.clear()

    def done(self) -> None:
        self.add(-1)

    async def wait(self) -> None:
        await self._zero_event.wait()

class WebSocketMonitor:
    """Receive WebSocket messages with an idle probe and a bounded silence timeout.

    The socket must be connected with ``autoping=False`` and without aiohttp's
    heartbeat. Incoming control frames are handled here so they count as activity.
    Use one monitor per active connection, including across cancelled receive calls.
    """

    def __init__(
        self,
        ws: aiohttp.ClientWebSocketResponse,
        *,
        liveness_timeout: float = DEFAULT_API_CONNECT_OPTIONS.liveness_timeout,
    ) -> None:
        self._ws = ws
        self._liveness_timeout = liveness_timeout
        self._ping_interval = liveness_timeout / 3
        self._loop = asyncio.get_running_loop()
        self._last_received = self._loop.time()
        self._ping_sent = False

    async def receive(self) -> aiohttp.WSMessage:
        receiver = asyncio.create_task(self._ws.receive())
        ping: asyncio.Task[None] | None = None
        try:
            while True:
                deadline = self._last_received + (
                    self._liveness_timeout if self._ping_sent else self._ping_interval
                )
                tasks: list[asyncio.Task[aiohttp.WSMessage] | asyncio.Task[None]] = [receiver]
                if ping is not None:
                    tasks.append(ping)
                done, _ = await asyncio.wait(
                    tasks,
                    timeout=max(0.0, deadline - self._loop.time()),
                    return_when=asyncio.FIRST_COMPLETED,
                )
                if receiver in done:
                    msg = receiver.result()
                    if msg.type in (
                        aiohttp.WSMsgType.TEXT,
                        aiohttp.WSMsgType.BINARY,
                        aiohttp.WSMsgType.PING,
                        aiohttp.WSMsgType.PONG,
                    ):
                        self._last_received = self._loop.time()
                        self._ping_sent = False
                        if ping is not None:
                            ping.cancel()
                            await asyncio.gather(ping, return_exceptions=True)
                            ping = None

                    if msg.type == aiohttp.WSMsgType.PING:
                        await asyncio.wait_for(self._ws.pong(msg.data), self._liveness_timeout)
                    elif msg.type != aiohttp.WSMsgType.PONG:
                        return msg
                    receiver = asyncio.create_task(self._ws.receive())
                    continue

                if ping is not None and ping in done:
                    ping.result()
                    ping = None
                    continue

                if self._loop.time() - self._last_received >= self._liveness_timeout:
                    raise aiohttp.ServerTimeoutError(
                        f"No WebSocket messages received for {self._liveness_timeout} seconds"
                    )
                self._ping_sent = True
                ping = asyncio.create_task(self._ws.ping())
        except (aiohttp.ClientError, ConnectionError, asyncio.TimeoutError) as exc:
            # A dead peer cannot complete a close handshake. Bound cleanup so the
            # original failure reaches retry/fallback without another close timeout.
            with contextlib.suppress(Exception):
                await asyncio.wait_for(self._ws.close(), _CLOSE_TIMEOUT)
            raise APIConnectionError("WebSocket connection lost") from exc
        finally:
            receiver.cancel()
            if ping is not None:
                ping.cancel()
            await asyncio.gather(
                receiver, *([ping] if ping is not None else []), return_exceptions=True
            )

    def __aiter__(self) -> WebSocketMonitor:
        return self

    async def __anext__(self) -> aiohttp.WSMessage:
        msg = await self.receive()
        if msg.type in (
            aiohttp.WSMsgType.CLOSE,
            aiohttp.WSMsgType.CLOSED,
            aiohttp.WSMsgType.CLOSING,
        ):
            raise StopAsyncIteration
        return msg