from collections.abc import Callable

from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncEngine, AsyncSession

class AsyncSessionManager:
	"""
	Asynchronous context manager for handling commitment and rollback of `AsyncSession` objects.

	The session is commited on successful exit.
	If an exception occurs, the transaction is rolled back.
	Finally, the session is closed.

	Example::
	
		AsyncSession = async_sessionmaker(engine)

		async with AsyncSessionManager(AsyncSession) as session:
			await session.execute(update(...))

			# On success, transaction is commited.
			# On exception, it is rolled back.
			# Session is closed regardless.

	:param session_factory: Callable that returns an `AsyncSession`
	"""
	
	def __init__(self, session_factory: Callable[..., AsyncSession]):
		self.session_factory = session_factory
		self.session: AsyncSession | None = None
	
	async def __aenter__(self):
		self.session = self.session_factory()

		return self.session
	
	async def __aexit__(self, exc_type, exc, tb):
		try:
			if exc_type is None:
				await self.session.commit()
			else:
				await self.session.rollback()
		finally:
			await self.session.close()

			self.session = None

class AsyncSessionManagerFactory:
	def __init__(self, engine: str | AsyncEngine=None):
		if isinstance(engine, str):
			if not engine.startswith("sqlite+aiosqlite:///"):
				engine = "sqlite+aiosqlite:///" + engine

			self.engine = create_async_engine(engine)
		elif isinstance(engine, AsyncEngine):
			self.engine = engine

		self.session_factory = async_sessionmaker(self.engine, expire_on_commit=False)

	def __call__(self):
		return AsyncSessionManager(self.session_factory)

__all__ = ["AsyncSessionManager", "AsyncSessionManagerFactory"]