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

from eda_agent.config import settings

# Convert database URL based on type
database_url = settings.eda_database_url

# For PostgreSQL, use asyncpg driver
if database_url.startswith("postgresql://"):
    database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1)

# Engine configuration
engine_kwargs = {
    "echo": settings.eda_api_debug,
}

# SQLite specific settings
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 all EDA database models."""
    pass


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()
