"""EDA Agent database connection module."""
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase

from eda_agent.config import settings

database_url = settings.eda_database_url
if database_url.startswith("postgresql://"):
    database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)

engine_kwargs = {"echo": settings.eda_api_debug}
if "sqlite" in database_url:
    engine_kwargs["connect_args"] = {"check_same_thread": False}
else:
    engine_kwargs["pool_pre_ping"] = True
    engine_kwargs["pool_size"] = 10
    engine_kwargs["max_overflow"] = 20

engine = create_async_engine(database_url, **engine_kwargs)
async_session_maker = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
)


class Base(DeclarativeBase):
    """Base class for EDA Agent ORM models."""


async def get_db() -> AsyncSession:
    """Dependency for getting EDA database session."""
    async with async_session_maker() as session:
        yield session


async def init_db():
    """Initialize EDA database tables."""
    # Import models so metadata is populated before create_all.
    from eda_agent.models import persistence  # noqa: F401

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


async def close_db():
    """Close EDA database connections."""
    await engine.dispose()
