from __future__ import annotations

from collections.abc import Callable
from types import TracebackType
from typing import Optional

from sqlmodel import Session

from repositories.magic_plan.magic_plan_postgres_repository import (
    MagicPlanPostgresRepository,
)
from repositories.uploaded_file.uploaded_file_postgres_repository import (
    UploadedFilePostgresRepository,
)


class AuditGeneratorUnitOfWork:
    def __init__(self, session_factory: Callable[[], Session]) -> None:
        self._session_factory = session_factory

    def __enter__(self) -> "AuditGeneratorUnitOfWork":
        self._session = self._session_factory()
        self.uploaded_file = UploadedFilePostgresRepository(self._session)
        self.magic_plan = MagicPlanPostgresRepository(self._session)
        return self

    def __exit__(
        self,
        exc_type: Optional[type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> None:
        try:
            self._session.rollback()
        finally:
            self._session.close()

    def commit(self) -> None:
        self._session.commit()