from contextlib import contextmanager
from typing import Optional
from app.database.session import sessionmanager
from sqlalchemy import MetaData
from app.database.models import Base


@contextmanager
def with_db(tenant_schema: Optional[str] = None):
    """
    Gerencia uma sessão do banco de dados com suporte a schema_translate_map.

    Args:
        tenant_schema (Optional[str]): Nome do schema do tenant (opcional).

    Yields:
        AsyncSession: Sessão configurada para o schema especificado.
    """
    if not sessionmanager.is_initialized():
        raise Exception("DatabaseSessionManager is not initialized")

    # Configura o schema_translate_map para o inquilino
    if tenant_schema:
        schema_translate_map = {None: tenant_schema}
    else:
        schema_translate_map = None

    # Configura a conexão com o schema correto
    connectable = sessionmanager.get_engine().execution_options(schema_translate_map=schema_translate_map)

    # Cria uma sessão vinculada ao connectable configurado
    session = sessionmanager.get_sessionmaker(bind=connectable)
    try:
        yield session
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def get_tenant_specific_metadata():
    meta = MetaData(schema=None)
    for table in Base.metadata.tables.values():
        if table.schema is None:
            table.to_metadata(meta)
    return meta