"""Async SQLAlchemy database setup for Smart City Digital Twin."""

from collections.abc import AsyncGenerator
from typing import NoReturn

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

from app.config import settings


# ── Declarative base ──────────────────────────────────────────────────
class Base(DeclarativeBase):
    pass


# ── Engine ───────────────────────────────────────────────────────────
engine = create_async_engine(
    settings.DATABASE_URL,
    echo=False,
    pool_size=10,
    max_overflow=20,
)

# ── Session factory ──────────────────────────────────────────────────
AsyncSessionLocal = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
)


# ── Dependency: per-request session ──────────────────────────────────
async def get_session() -> AsyncGenerator[AsyncSession, None]:
    """Yield an async session and close it when the request is done."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()


# ── Startup: create tables ───────────────────────────────────────────
async def init_db() -> None:
    """Create all tables defined via ``Base`` models (idempotent)."""
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


# ── Shutdown: release connections ────────────────────────────────────
async def close_db() -> None:
    """Dispose the engine and release all pooled connections."""
    await engine.dispose()