"""Auth module Unit of Work."""

from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from identity_plan_kit.auth.repositories.token_repo import RefreshTokenRepository
from identity_plan_kit.auth.repositories.user_repo import UserRepository
from identity_plan_kit.plans.repositories.plan_repo import PlanRepository
from identity_plan_kit.rbac.repositories.rbac_repo import RBACRepository
from identity_plan_kit.shared.uow import BaseUnitOfWork


class AuthUnitOfWork(BaseUnitOfWork):
    """
    Unit of Work for authentication operations.

    Provides access to user, token, plan, and RBAC repositories within a transaction.
    The plan and RBAC repositories are included to support atomic user registration
    with default plan/role assignment.

    Usage with internal session (IPK manages transaction):
        ```python
        async with AuthUnitOfWork(session_factory) as uow:
            user = await uow.users.get_by_email(email)
            await uow.tokens.create(user.id, token_hash, expires_at)
            # Commits automatically on successful exit
        ```

    Usage with external session (your app manages transaction):
        ```python
        async with your_session.begin():
            async with AuthUnitOfWork(session_factory, session=your_session) as uow:
                user = await uow.users.get_by_email(email)
                # Your other operations in same transaction
            # You control commit/rollback
        ```
    """

    users: UserRepository
    tokens: RefreshTokenRepository
    plans: PlanRepository
    rbac: RBACRepository

    def __init__(
        self,
        session_factory: async_sessionmaker[AsyncSession],
        session: AsyncSession | None = None,
    ) -> None:
        super().__init__(session_factory, session=session)
        self.users: UserRepository
        self.tokens: RefreshTokenRepository
        self.plans: PlanRepository
        self.rbac: RBACRepository

    def _init_repositories(self) -> None:
        """Initialize auth repositories."""
        self.users = UserRepository(self.session)
        self.tokens = RefreshTokenRepository(self.session)
        self.plans = PlanRepository(self.session)
        self.rbac = RBACRepository(self.session)
