# Importações de bibliotecas padrão
import contextlib
from typing import AsyncIterator, Optional

from fastapi import Depends
# Importações de bibliotecas de terceiros
from sqlalchemy import MetaData
from sqlalchemy.ext.asyncio import (
    AsyncConnection,
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
    AsyncAttrs,
)
from sqlalchemy.orm import DeclarativeBase
from app.config import ECHO
import traceback

EXTENSIONS = ["uuid-ossp", "postgis", "postgis_topology"]

naming_convention = {
    "ix": "ix_ct_%(table_name)s_%(column_0_N_name)s",
    "uq": "uq_ct_%(table_name)s_%(column_0_N_name)s",
    "ck": "ck_ct_%(table_name)s_%(constraint_name)s",
    "fk": "fk_ct_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "pk": "pk_ct_%(table_name)s",
}


class Base(DeclarativeBase, AsyncAttrs):
    metadata = MetaData(naming_convention=naming_convention)


class DatabaseSessionManager:
    def __init__(self):
        self._engine: AsyncEngine | None = None
        self._sessionmaker: async_sessionmaker | None = None

    def init(self, host: str):
        self._engine = create_async_engine(host, echo=ECHO)
        self._sessionmaker = async_sessionmaker(autocommit=False, bind=self._engine)

    async def close(self):
        if self._engine is None:
            raise Exception("DatabaseSessionManager is not initialized")
        await self._engine.dispose()
        self._engine = None
        self._sessionmaker = None

    def is_initialized(self) -> bool:
        """Verifica se o engine foi inicializado"""
        return self._engine is not None

    @contextlib.asynccontextmanager
    async def connect(self) -> AsyncIterator[AsyncConnection]:
        if self._engine is None:
            raise Exception("DatabaseSessionManager is not initialized")

        async with self._engine.begin() as connection:
            try:
                yield connection
            except Exception:
                await connection.rollback()
                raise

    @contextlib.asynccontextmanager
    async def session(self) -> AsyncIterator[AsyncSession]:
        if self._sessionmaker is None:
            raise Exception("DatabaseSessionManager is not initialized")

        session = self._sessionmaker()
        try:
            yield session
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()

    @contextlib.asynccontextmanager
    async def session_with_tenant(self, tenant_schema: Optional[str]) -> AsyncIterator[AsyncSession]:

        if self._engine is None:
            raise Exception("DatabaseSessionManager is not initialized")

        # Verifica e configura o schema_translate_map
        if tenant_schema:

            schema_engine = self._engine.execution_options(
                schema_translate_map={None: str(tenant_schema)}
            )
        else:

            schema_engine = self._engine.execution_options(
                schema_translate_map=None
            )

        # Cria a sessão vinculada ao schema_engine
        try:
            session = self._sessionmaker(bind=schema_engine)


        except Exception as e:
            print(f"Erro ao configurar o bind da sessão: {e}")
            raise

        try:
            yield session
        except Exception as e:
            await session.rollback()
            print(f"Erro encontrado: {e}. Traceback: {traceback.format_exc()}")
            raise
        finally:
            await session.close()
        #     print("Sessão encerrada")

    async def create_all(self, connection: AsyncConnection):
        await connection.run_sync(Base.metadata.create_all)

    async def drop_all(self, connection: AsyncConnection):
        await connection.run_sync(Base.metadata.drop_all)

    # Multi Tenant
    def get_engine(self):
        if not self.is_initialized():
            raise Exception("DatabaseSessionManager is not initialized")
        return self._engine

    def get_sessionmaker(self):
        if self._sessionmaker is None:
            raise Exception("DatabaseSessionManager is not initialized")
        return self._sessionmaker

    # Multi Tenant


sessionmanager = DatabaseSessionManager()


async def get_db():
    async with sessionmanager.session() as session:
        yield session