from collections.abc import Awaitable, Callable
from functools import wraps
from typing import Concatenate, ParamSpec, Protocol, TypeVar

from .base import Database


class _HasDatabase(Protocol):
    db: Database


HasDatabase = TypeVar("HasDatabase", bound=_HasDatabase)
P = ParamSpec("P")
R = TypeVar("R")


def session_context(
    in_transaction: bool,
) -> Callable[
    [Callable[Concatenate[HasDatabase, P], Awaitable[R]]],
    Callable[Concatenate[HasDatabase, P], Awaitable[R]],
]:
    def decorator(
        func: Callable[Concatenate[HasDatabase, P], Awaitable[R]]
    ) -> Callable[Concatenate[HasDatabase, P], Awaitable[R]]:
        @wraps(func)
        async def wrapper(self: HasDatabase, /, *args: P.args, **kwargs: P.kwargs) -> R:
            async with self.db.session_context(in_transaction=in_transaction):
                return await func(self, *args, **kwargs)  # noqa

        return wrapper

    return decorator