from contextlib import asynccontextmanager, contextmanager

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlmodel import Session, create_engine


class DBManager:
    sync_url = ''
    async_url = ''

    def __init__(self, database_url: str):
        self.sync_url, self.async_url = self._initialize_urls(database_url)

        self.sync_engine = create_engine(self.sync_url)
        self.async_engine = create_async_engine(self.async_url)

    def _initialize_urls(self, database_url: str):
        url_parts = database_url.split('://')

        return (
            f'postgresql+psycopg://{url_parts[1]}',
            f'postgresql+asyncpg://{url_parts[1]}',
        )

    def sync_session(self):
        with Session(self.sync_engine) as session:
            yield session

    async def async_session(self):
        async with AsyncSession(self.async_engine) as session:
            yield session

    @contextmanager
    def sync_context_session(self):
        with Session(self.sync_engine) as session:
            yield session

    @asynccontextmanager
    async def async_context_session(self):
        async with AsyncSession(self.async_engine) as session:
            yield session
