"""
Provides a Golang-like ``defer()`` API using decorators, which allows grouping
resource initialization and cleanup in one place without extra indentations.

Example:

.. code-block:: python3

   async def init(x):
       ...

   async def cleanup(x):
       ...

   @aiotools.adefer
   async def do(defer):  # <-- be aware of defer argument!
       x = SomeResource()
       await init(x)
       defer(cleanup(x))
       ...
       ...

This is equivalent to:

.. code-block:: python3

   async def do():
       x = SomeResource()
       await init(x)
       try:
           ...
           ...
       finally:
           await cleanup(x)

Note that :class:`aiotools.context.AsyncContextGroup` or
:class:`contextlib.AsyncExitStack` serves well for the same purpose, but for simple
cleanups, this defer API makes your codes simple because it steps aside the main
execution context without extra indentations.

.. warning::

   Any exception in the deferred functions is raised transparently, and may block
   execution of the remaining deferred functions.
   This behavior may be changed in the future versions, though.
"""

from __future__ import annotations

import functools
import inspect
from collections import deque
from collections.abc import Awaitable, Callable
from typing import Any, Concatenate, ParamSpec, TypeAlias, TypeVar

P = ParamSpec("P")
R = TypeVar("R")

__all__ = (
    "defer",
    "adefer",
    "DeferFunc",
    "AsyncDeferrable",
    "AsyncDeferFunc",
)

DeferFunc: TypeAlias = Callable[[Callable[[], None]], None]
AsyncDeferrable: TypeAlias = (
    Callable[[], Awaitable[Any]] | Callable[[], None] | Awaitable[None]
)
AsyncDeferFunc: TypeAlias = Callable[[AsyncDeferrable], None]


def defer(func: Callable[Concatenate[DeferFunc, P], R]) -> Callable[P, R]:
    """
    A synchronous version of the defer API.
    It can only defer normal functions.
    """
    assert not inspect.iscoroutinefunction(func), (
        "the decorated function must not be async"
    )

    @functools.wraps(func)
    def _wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
        deferreds: deque[Callable[[], None]] = deque()

        # the registration function passed to the original function
        def defer(f: Callable[[], None]) -> None:
            assert not inspect.iscoroutinefunction(f), (
                "the deferred function must not be async"
            )
            assert not inspect.iscoroutine(f), (
                "the deferred object must not be a coroutine"
            )
            deferreds.append(f)

        try:
            return func(defer, *args, **kwargs)
        finally:
            while deferreds:
                f = deferreds.pop()
                f()

    return _wrapped


def adefer(
    func: Callable[Concatenate[AsyncDeferFunc, P], Awaitable[R]],
) -> Callable[P, Awaitable[R]]:
    """
    An asynchronous version of the defer API.
    It can defer coroutine functions, coroutines, and normal functions.
    """
    assert inspect.iscoroutinefunction(func), "the decorated function must be async"

    @functools.wraps(func)
    async def _wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
        deferreds: deque[AsyncDeferrable] = deque()

        # the registration function passed to the original function
        def defer(f: AsyncDeferrable) -> None:
            deferreds.append(f)

        try:
            return await func(defer, *args, **kwargs)
        finally:
            while deferreds:
                f = deferreds.pop()
                if inspect.iscoroutinefunction(f):
                    await f()
                elif inspect.iscoroutine(f):
                    await f
                else:
                    f()  # type: ignore[operator]

    return _wrapped
