import inspect
from contextvars import ContextVar
from contextlib import asynccontextmanager, AbstractAsyncContextManager
from functools import wraps

from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker

__all__ = ['RepoBase']



class RepoBase(AbstractAsyncContextManager):
    def __init__(self, db: AsyncEngine):
        self.db = db
        self._asession_maker = async_sessionmaker(self.db, expire_on_commit=False)
        self._current_session: ContextVar = ContextVar('current_session', default=None)

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        for name, method in cls.__dict__.items():
            if inspect.iscoroutinefunction(method) and not name.startswith('_'):
                def make_wrapper(m):
                    @wraps(m)
                    async def wrapper(self, *args, **kwargs):
                        async with self._session():
                            return await m(self, *args, **kwargs)
                    return wrapper
                setattr(cls, name, make_wrapper(method))

    @property
    def _s(self):
        return self._current_session.get()

    @asynccontextmanager
    async def _session(self):
        if self._current_session.get() is not None:
            yield
            return
        async with self._asession_maker.begin() as session:
            token = self._current_session.set(session)
            try:
                yield
            finally:
                self._current_session.reset(token)

    async def __aenter__(self):
        if self._current_session.get() is not None:
            raise RuntimeError("Nested transactions are not supported")
        session = self._asession_maker()
        try:
            await session.begin()
        except:
            await session.close()
            raise
        self._current_session.set(session)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        session = self._current_session.get()
        self._current_session.set(None)
        try:
            if exc_type is None:
                await session.commit()
            else:
                await session.rollback()
        finally:
            await session.close()
