"""Engine, session factory, and the FastAPI ``get_db`` dependency.

A single sync engine is built from ``DATABASE_URL``. SQLite (used by the test
suite) gets the connection args it needs; everything else uses sane pool
defaults.
"""

from __future__ import annotations

from collections.abc import Iterator

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker

from legate.config import get_settings


def _make_engine(url: str) -> Engine:
    connect_args: dict = {}
    if url.startswith("sqlite"):
        connect_args["check_same_thread"] = False
    return create_engine(url, pool_pre_ping=True, future=True, connect_args=connect_args)


engine: Engine = _make_engine(get_settings().database_url)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)


def get_db() -> Iterator[Session]:
    """Yield a session and guarantee it is closed."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


def get_session_factory() -> sessionmaker:
    """Return the session factory.

    Exposed as a dependency so background work (e.g. agent runs that outlive the
    request) can open its own session, and so tests can point it at their
    in-memory database.
    """
    return SessionLocal
