from typing import (
    Callable,
    Optional,
)

from sqlalchemy.orm import Session

from messageprocessor import MessageProcessorCallback


class SessionProcessorCallback(MessageProcessorCallback):
    def __init__(
        self,
        session: Session,
        isolation_level_retriever: Optional[Callable[[str, str], str]]=None,
    ) -> None:

        self.session = session
        self.isolation_level_retriever = isolation_level_retriever

    def process_begun(self, event_name, request_id, data):
        self._set_isolation_level_on_connection("process_begun", event_name)

    def input_parsed(self, event_name, request_id, input_obj):
        self._set_isolation_level_on_connection("input_parsed", event_name)

    def business_called(self, event_name, request_id, input_obj, output_obj):
        self.session.commit()

    def process_finished(self, event_name, request_id, input_obj, output_obj):
        return

    def parse_exception(self, event_name, request_id, data, exception):
        self.session.rollback()

    def unhandled_exception(self, event_name, request_id, input_obj, exception):
        self.session.rollback()

    def _set_isolation_level_on_connection(self, step, event_name):
        if self.isolation_level_retriever is not None:
            isolation_level = self.isolation_level_retriever(step, event_name)
            if isolation_level is not None:
                self.session.connection(
                    execution_options={
                        "isolation_level": isolation_level
                    }
                )
