"""Database engine and session factory for the memory SDK.

Creates an async SQLAlchemy engine from a connection string provided by the consumer.
The SDK owns its own DeclarativeBase — no dependency on external models.
"""

from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase


class Base(DeclarativeBase):
    """Base class for all memory SDK models."""

    pass


def create_engine(database_url: str) -> AsyncEngine:
    """Create an async SQLAlchemy engine from a connection string.

    Automatically converts ``postgresql://`` to ``postgresql+psycopg://``
    if the async driver prefix is missing.
    """
    if database_url.startswith("postgresql://"):
        database_url = database_url.replace("postgresql://", "postgresql+psycopg://", 1)

    return create_async_engine(database_url, echo=False)


def create_session_factory(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
    """Create an async session factory bound to the given engine."""
    return async_sessionmaker(engine, expire_on_commit=False)


async def init_db(engine: AsyncEngine) -> None:
    """Create all memory tables in the consumer's database.

    Uses ``Base.metadata.create_all`` — safe to call multiple times
    (creates only tables that don't already exist).
    """
    from arandu.models import Base as _  # type: ignore[attr-defined]  # noqa: F401 — ensure models registered

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