# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Track queued publication callbacks on their originating event loop."""

from __future__ import annotations

import inspect
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine
from contextvars import ContextVar
from typing import Any


class _CallbackState:
    """Shared liveness for contexts copied from one publication callback."""

    __slots__ = ("active",)

    def __init__(self) -> None:
        self.active = True


_ACTIVE: ContextVar[_CallbackState | None] = ContextVar("nemo_relay_event_sanitizer_active", default=None)


def callback_active() -> bool:
    """Return whether the current context is running queued publication work."""
    state = _ACTIVE.get()
    return state is not None and state.active


async def _await_result(result: Awaitable[Any], state: _CallbackState, owner: bool) -> Any:
    token = _ACTIVE.set(state)
    try:
        return await result
    finally:
        if owner:
            state.active = False
        _ACTIVE.reset(token)


async def await_result(result: Awaitable[Any]) -> Any:
    """Await an arbitrary awaitable without changing sanitizer context."""
    return await result


async def drive_stream_future(driver: Any) -> Any:
    """Poll a native streaming future locally, awaiting only pending work."""
    try:
        while True:
            ready, result = driver.poll()
            if ready:
                return result
            await result
    finally:
        driver.cancel()


def loop_affine(callback: Callable[..., Any], *, sanitizer: bool = False) -> Callable[..., Awaitable[Any]]:
    """Defer a callback's synchronous prelude to the awaiting event-loop task."""

    async def wrapped(*args: Any) -> Any:
        result = invoke(callback, *args) if sanitizer else callback(*args)
        if inspect.isawaitable(result):
            return await result
        return result

    return wrapped


def async_iter_next(iterator: Any) -> Coroutine[Any, Any, Any]:
    """Validate the iterator then await its next item on event-loop thread."""
    if not isinstance(iterator, AsyncIterator):
        raise TypeError(f"expected an asynchronous iterator, got {type(iterator).__name__!r}.")

    return _anext_on_event_loop(iterator)


async def _anext_on_event_loop(iterator: AsyncIterator[Any]) -> Any:
    return await anext(iterator)


async def async_iter_close(iterator: Any) -> None:
    """Invoke and await ``aclose`` on the current event-loop thread when present."""
    close = getattr(iterator, "aclose", None)
    if close is not None:
        await close()


def invoke(callback: Callable[..., Any], *args: Any) -> Any:
    """Invoke queued publication work while marking its execution context."""
    state = _ACTIVE.get()
    owner = state is None or not state.active
    if owner:
        state = _CallbackState()
    assert state is not None
    token = _ACTIVE.set(state)
    try:
        result = callback(*args)
    except BaseException:
        if owner:
            state.active = False
        raise
    finally:
        _ACTIVE.reset(token)
    if inspect.isawaitable(result):
        return _await_result(result, state, owner)
    if owner:
        state.active = False
    return result