from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from src.repository.auth import AuthRepo
from src.repository.tasks import TasksRepo
from src.repository.users import UsersRepo


class DBManager:
    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
        self.session_factory = session_factory

    async def __aenter__(self) -> "DBManager":
        self.session: AsyncSession = self.session_factory()
        self.user = UsersRepo(self.session)
        self.task = TasksRepo(self.session)
        self.auth = AuthRepo(self.session)
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        await self.session.rollback()
        await self.session.close()

    async def commit(self) -> None:
        await self.session.commit()