"""Unit of Work — wraps a SQLAlchemy session for explicit transaction control.

Usage in routers::

    with UnitOfWork(db) as uow:
        service_a(db, ...)
        service_b(db, ...)
        uow.commit()          # single commit for the entire operation

If an exception propagates, ``__exit__`` issues a rollback automatically.
Services should **never** call ``db.commit()``; they use ``db.add()`` /
``db.flush()`` to stage work and let the caller decide when to commit.
"""

from __future__ import annotations

from sqlalchemy.orm import Session


class UnitOfWork:
    """Lightweight transaction wrapper around an existing SQLAlchemy session."""

    def __init__(self, session: Session) -> None:
        self._session = session

    # -- context manager -----------------------------------------------------

    def __enter__(self) -> "UnitOfWork":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if exc_type is not None:
            self.rollback()

    # -- public API ----------------------------------------------------------

    def commit(self) -> None:
        """Flush pending changes and commit the transaction."""
        self._session.commit()

    def rollback(self) -> None:
        """Roll back the current transaction."""
        self._session.rollback()

    def flush(self) -> None:
        """Flush pending changes without committing (useful for getting IDs)."""
        self._session.flush()