import asyncio
from typing import Dict, Any
from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.exc import OperationalError, IntegrityError, DBAPIError, InvalidRequestError
from Osdental.Exception.ControlledException import OSDException, DatabaseException
from Osdental.Shared.Enums.Code import Code

class UnitOfWork:
    def __init__(self, session: AsyncSession, repositories: Dict[str, Any] = None):
        self.session = session
        self._repositories = repositories or {}
        self._instances = {}

    def __getattr__(self, item):
        if item in self._instances:
            return self._instances[item]
        
        raise OSDException(error=f'UnitOfWork object has no attribute {item}')


    @asynccontextmanager
    async def __call__(self):
        try:
            async with self.session.begin():
                self._instances = {name: repo_cls(self.session) for name, repo_cls in self._repositories.items()}
                try:
                    yield self
                except Exception:
                    await self.session.rollback()
                    raise
                else:
                    await self.session.commit()

        except asyncio.CancelledError:
            raise 
        except IntegrityError as ierr:
            raise DatabaseException(status_code=Code.INTEGRITY_DB_ERROR_CODE, error=str(ierr))
        except OperationalError as oerr:
            raise DatabaseException(status_code=Code.OPERATIONAL_DB_ERROR_CODE, error=str(oerr))
        except (DBAPIError, InvalidRequestError) as db_err:
            raise DatabaseException(status_code=Code.DB_UNKNOWN_ERROR_CODE, error=str(db_err))
        except DatabaseException:
            raise
        except Exception as err:
            raise OSDException(error=str(err))
        finally:
            self._instances.clear()
