"""Database session management."""

from typing import AsyncGenerator

from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import NullPool

from apex.core.config import get_settings

settings = get_settings()


def get_database_url() -> str:
    """Convert PostgreSQL URL to async-compatible format."""
    url = str(settings.DATABASE_URL)
    if url.startswith("postgresql://"):
        return url.replace("postgresql://", "postgresql+asyncpg://", 1)
    return url


# Create async engine
engine = create_async_engine(
    get_database_url(),
    echo=settings.DB_ECHO,
    poolclass=NullPool,
    future=True,
)

# Create async session factory
AsyncSessionLocal = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False,
)


# For synchronous operations (migrations, CLI)
def get_sync_database_url() -> str:
    """Get synchronous database URL."""
    return str(settings.DATABASE_URL)


# Note: For synchronous operations, use create_engine instead
# This is kept for compatibility but users should use async operations
from sqlalchemy import create_engine

sync_engine = create_engine(
    get_sync_database_url(),
    echo=settings.DB_ECHO,
    future=True,
)

SessionLocal = sessionmaker(
    bind=sync_engine,
    autocommit=False,
    autoflush=False,
)


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """Dependency for getting database session."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()

