"""
Database engine and session utilities for Chrono Temporal.
"""
from typing import AsyncGenerator
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    create_async_engine,
    async_sessionmaker,
    AsyncEngine,
)
from sqlalchemy.orm import DeclarativeBase


class ChronoBase(DeclarativeBase):
    """Base class for all Chrono Temporal models."""
    pass


def get_engine(database_url: str, echo: bool = False, **kwargs) -> AsyncEngine:
    """
    Create an async SQLAlchemy engine.

    Args:
        database_url: PostgreSQL connection string (must use asyncpg driver)
                      e.g. postgresql+asyncpg://user:pass@localhost/dbname
        echo: Log SQL statements (default False)
        **kwargs: Additional engine options

    Returns:
        AsyncEngine instance
    """
    return create_async_engine(
        database_url,
        echo=echo,
        pool_pre_ping=True,
        **kwargs,
    )


def get_session(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
    """
    Create an async session factory.

    Args:
        engine: AsyncEngine instance from get_engine()

    Returns:
        async_sessionmaker configured for the given engine
    """
    return async_sessionmaker(
        engine,
        class_=AsyncSession,
        expire_on_commit=False,
    )


async def create_tables(engine: AsyncEngine) -> None:
    """Create all Chrono Temporal tables in the database."""
    from chrono_temporal.models import TemporalRecord  # noqa
    async with engine.begin() as conn:
        await conn.run_sync(ChronoBase.metadata.create_all)
