from typing import Callable, Any, TypeVar, List
import enum 
import dataclasses

T = TypeVar("T")


class DeferType(enum.Enum):
    ON_SUCCESS = enum.auto()
    ON_FAILURE = enum.auto()
    ON_EXIT = enum.auto()

@dataclasses.dataclass
class DeferFunc:
    f: Callable[[], None]
    defer_type: DeferType


class DeferException(RuntimeError):
    def __init__(self, stack):
        super().__init__(
            "exceptions were raised while calling defered functions", 
            stack
        )
        self.stack = stack


class Defer:
    def __init__(self, default_exit_success: bool=True):
        self._defer_funcs: List[DeferFunc] = []
        self._success = default_exit_success

    def __enter__(self):
        return self

    def __exit__(self, exception_type, exception_value, exception_traceback):
        if exception_type is not None:
            self._success = False

        exceptions: List[Exception] = []
        for defer_func in reversed(self._defer_funcs):
            try:
                if defer_func.defer_type == DeferType.ON_EXIT:
                    defer_func.f()
                elif defer_func.defer_type == DeferType.ON_SUCCESS and self._success:
                    defer_func.f()
                elif defer_func.defer_type == DeferType.ON_FAILURE and not self._success:
                    defer_func.f()
            except Exception as e: # pylint: disable=broad-exception-caught
                exceptions.append(e)

        if len(exceptions) > 0:
            raise DeferException(exceptions)

    def exit_success(self, any_ret_val: T) -> T:
        self._success = True
        return any_ret_val

    def exit_failure(self, any_ret_val: T) -> T:
        self._success = False
        return any_ret_val

    def push(self, f: Callable[[], Any], defer_type: DeferType = DeferType.ON_EXIT):
        self._defer_funcs.append(DeferFunc(f, defer_type))
