"""
session.py

This module implements the Unit of Work pattern for asynchronous SQLAlchemy sessions.

Key features:
    - Centralized management of async database sessions and transactions.
    - Explicit commit and rollback handling with custom exception translation.
    - Designed to be used with 'async with' for safe and automatic transaction boundaries.
    - Suitable for working with repository patterns and batch/bulk operations.
"""

from contextlib import AbstractAsyncContextManager
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from sqlalchemy.exc import IntegrityError, DataError, SQLAlchemyError
from sqlalchemy.orm.exc import FlushError
from .utils import ensure_active_session

from ..errors import (
    UnknownTransactionError,
    ErrorMapper,
)


class UnitOfWork(AbstractAsyncContextManager):
    """
    Async Unit of Work for SQLAlchemy sessions.

    Handles session/transaction lifecycle for one business operation, with robust error mapping.

    Args:
        session_factory (async_sessionmaker): Factory for creating AsyncSession instances.

    Attributes:
        session (AsyncSession): The current database session within the Unit of Work.

    Key Features:
        - Centralized transaction and session management for async workflows.
        - Safe usage via 'async with' blocks for automatic commit/rollback on exit.
        - Explicit 'flush()' method to synchronize ORM state with the database and fetch autogenerated values (e.g., primary keys) before commit.
        - Integrated exception translation via ErrorMapper: all common SQLAlchemy errors are mapped to domain-specific exceptions for unified error handling.
        - Suitable for integration with repository patterns and batch/bulk operations.

    Usage:
        async with UnitOfWork(session_factory) as uow:
            # Use uow.session for all DB operations within this transaction
            await uow.flush()    # (optional) Force sync with DB
            await uow.commit()   # (optional, usually not needed, handled automatically)

    Raises:
        UniqueViolation: If a unique constraint is violated on flush or commit.
        ForeignKeyViolation: If a foreign key constraint is violated on flush or commit.
        DataValidationError: If provided data is invalid for the model.
        CreateError: For generic integrity errors.
        TransactionError: For other SQLAlchemy transaction errors.
        UnknownTransactionError: For any other unexpected errors during flush or commit.
    """

    def __init__(self, session_factory: async_sessionmaker):
        self._session_factory = session_factory
        self.session: AsyncSession | None = None

    async def __aenter__(self):
        self.session = self._session_factory()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        try:
            if exc_type:
                await self.session.rollback()
            else:
                await self.commit()
        finally:
            await self.session.close()

    @ensure_active_session
    async def flush(self):
        try:
            await self.session.flush()
        except (IntegrityError, DataError, SQLAlchemyError, FlushError) as exc:
            await self.session.rollback()
            raise ErrorMapper.map(exc) from exc
        except Exception as exc:
            await self.session.rollback()
            raise UnknownTransactionError(details="", original=str(exc)) from exc

    @ensure_active_session
    async def commit(self):
        try:
            await self.session.commit()
        except (IntegrityError, DataError, SQLAlchemyError) as exc:
            await self.session.rollback()
            raise ErrorMapper.map(exc) from exc
        except Exception as exc:
            await self.session.rollback()
            raise UnknownTransactionError(details="", original=str(exc)) from exc

    async def rollback(self):
        await self.session.rollback()
