from collections.abc import Iterator
from contextlib import contextmanager
from typing import Any, Optional, Type

from sqlalchemy.engine import Engine
from sqlalchemy.pool import Pool
from sqlmodel import Session, create_engine

from infrastructure.postgres.config import PostgresConfig


def make_engine(
    config: PostgresConfig, poolclass: Optional[Type[Pool]] = None
) -> Engine:
    kwargs: dict[str, Any] = {"pool_pre_ping": config.pool_pre_ping}
    if poolclass is None:
        kwargs["pool_size"] = config.pool_size
        kwargs["max_overflow"] = config.max_overflow
        kwargs["pool_recycle"] = config.pool_recycle
    else:
        kwargs["poolclass"] = poolclass
    return create_engine(config.url(), **kwargs)


def make_session(engine: Engine) -> Session:
    return Session(engine)


@contextmanager  # pyright: ignore[reportDeprecated]
def transactional_session(engine: Engine) -> Iterator[Session]:
    """Yield a session whose lifecycle owns the transaction.

    On clean exit the session commits; on any exception it rolls back and
    re-raises. Either way the session is closed. Callers in the application
    layer can do their work inside the ``with`` block without ever invoking
    ``.commit()`` / ``.rollback()`` themselves -- transaction semantics stay
    in the infrastructure layer.
    """
    session = Session(engine)
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager  # pyright: ignore[reportDeprecated]
def commit_scope(session: Session) -> Iterator[Session]:
    """Commit a caller-owned session on clean exit; roll back on error.

    Like ``transactional_session`` but for a session the caller already holds
    and will close itself. Use it to keep slow, non-DB work *outside* the
    transaction: build the session, run the slow work, then enter
    ``commit_scope`` only for the persistence -- so a connection is checked out
    (SQLModel sessions are lazy) for the shortest possible window.
    """
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise