from collections.abc import AsyncGenerator, Callable

from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlmodel import SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession

# Import all models to ensure SQLModel metadata discovery
import app.models  # noqa: F401
from app.core.config import settings
from app.core.logging import get_logger

engine: AsyncEngine = create_async_engine(
    settings.DATABASE_URL,
    echo=settings.DATABASE_ECHO,
)


async def get_db() -> AsyncGenerator[AsyncSession]:
    """Get a database session for dependency injection."""
    logger = get_logger(__name__)
    async with AsyncSession(engine) as session:
        try:
            yield session
        except Exception:
            logger.exception("Database session error")
            await session.rollback()
            raise
        finally:
            await session.close()


def get_session_factory() -> Callable[[], AsyncSession]:
    """Get a session factory function for services."""

    def session_factory() -> AsyncSession:
        return AsyncSession(engine)

    return session_factory


async def init_db() -> None:
    """Initialize the database using Alembic migrations."""
    logger = get_logger(__name__)
    try:
        logger.info("Running database migrations")
        # Run Alembic migrations programmatically
        from alembic import command
        from alembic.config import Config

        # Get the alembic config
        alembic_cfg = Config("alembic.ini")

        # Run migrations to the latest revision
        command.upgrade(alembic_cfg, "head")
        logger.info("Database migrations completed successfully")

    except Exception:
        logger.exception("Failed to initialize database")
        raise