from contextlib import asynccontextmanager
from contextvars import ContextVar

from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.pool import NullPool

_current_db_session: ContextVar[AsyncSession | None] = ContextVar(
    "current_db_session",
    default=None,
)

class TenantDatabaseManager:
    def __init__(self, config):
        self.config = config
        self._engines = {}
        self._sessionmakers = {}

    def _build_connect_args(self, tenant: str) -> dict:
        return {
            "options": f"-c application_name={self.config.app_name_prefix}:{tenant}"
        }

    def _engine_args(self):
        if self.config.mode == "testing":
            return {
                "echo": self.config.echo_sql,
                "poolclass": NullPool,
            }

        return {
            "echo": self.config.echo_sql,
            "pool_pre_ping": self.config.pool_pre_ping,
            "pool_size": self.config.pool_size,
            "max_overflow": self.config.max_overflow,
            "pool_timeout": self.config.pool_timeout,
        }

    def get_engine(self, tenant: str):
        if tenant not in self._engines:
            db_url = self.config.service.get(tenant).db.url.get_secret_value()

            engine = create_async_engine(
                db_url,
                connect_args=self._build_connect_args(tenant),
                **self._engine_args(),
            )

            self._engines[tenant] = engine

        return self._engines[tenant]

    def get_sessionmaker(self, tenant: str):
        if tenant not in self._sessionmakers:
            self._sessionmakers[tenant] = async_sessionmaker(
                self.get_engine(tenant),
                class_=AsyncSession,
                expire_on_commit=False,
            )
        return self._sessionmakers[tenant]

    @asynccontextmanager
    async def db_session_context(self, tenant: str):
        sessionmaker_factory = self.get_sessionmaker(tenant)

        async with sessionmaker_factory() as db_session:
            token = _current_db_session.set(db_session)
            try:
                yield db_session
            finally:
                _current_db_session.reset(token)

    def get_current_db_session(self) -> AsyncSession | None:
        return _current_db_session.get()

    async def dispose_all(self):
        for engine in self._engines.values():
            await engine.dispose()