"""Async SQLAlchemy engine, session factory, and FastAPI dependency.

Driver-agnostic: works with both aiosqlite (dev) and asyncpg (prod).
"""

from __future__ import annotations

from collections.abc import AsyncGenerator
from typing import Any

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

from app.core.config import get_settings


class Base(DeclarativeBase):
    """Declarative base for all ORM models.

    Re-exported here so Alembic env.py can import it from a single location.
    """


# === Engine & Session Factory ===

_settings = get_settings()

# SQLite needs check_same_thread=False even for async — handled by aiosqlite.
# echo=False in production; controlled by env if needed later.
_engine_kwargs: dict[str, Any] = {"echo": False, "future": True}

# For PostgreSQL asyncpg, pool sizing matters; for SQLite it's irrelevant.
if not _settings.DATABASE_URL.startswith("sqlite"):
    _engine_kwargs["pool_size"] = 5
    _engine_kwargs["max_overflow"] = 10
    _engine_kwargs["pool_pre_ping"] = True

engine: AsyncEngine = create_async_engine(_settings.DATABASE_URL, **_engine_kwargs)

AsyncSessionLocal: async_sessionmaker[AsyncSession] = async_sessionmaker(
    bind=engine,
    expire_on_commit=False,
    class_=AsyncSession,
    autoflush=False,
)


# === Dependency ===


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI dependency that yields an AsyncSession and ensures cleanup."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
        except Exception:
            await session.rollback()
            raise
        # No explicit close needed — async context manager handles it.


async def dispose_engine() -> None:
    """Dispose of the engine on application shutdown."""
    await engine.dispose()