from sqlalchemy.exc import SQLAlchemyError

from app.dao.chat import ChatDAO
from app.dao.message import MessageDAO
from app.database import async_session_maker, mongo_db
from app.dao.user import UserDAO
from app.exceptions import BlackPhoenixException


class UnitOfWork:
	def __init__(self):
		self.session_factory = async_session_maker
		self.mongo_db = mongo_db

	async def __aenter__(self):
		self.session = self.session_factory()

		self.user = UserDAO(self.session)
		self.chat = ChatDAO(self.session)
		self.message = MessageDAO(self.mongo_db)

	async def __aexit__(self, exc_type, exc, tb):
		try:
			await self.rollback()
			await self.session.close()
		except SQLAlchemyError:
			raise BlackPhoenixException
		if isinstance(exc, SQLAlchemyError):
			raise BlackPhoenixException

	async def commit(self):
		await self.session.commit()

	async def rollback(self):
		await self.session.rollback()