import unittest.mock as mock


class _ANY_NOT_NONE(object):
    """A helper object similar to unittest.mock.ANY
    that compares equal to everything that isn't None.
    """

    def __eq__(self, other):
        return other is not None

    def __ne__(self, other):
        return other is None

    def __repr__(self):
        return "<ANY_NOT_NONE>"


ANY = mock.ANY
ANY_NOT_NONE = _ANY_NOT_NONE()


class MockFunction(mock.MagicMock):
    def __init__(self, fn_name):
        super().__init__(side_effect=self._side_effect)
        self._fn_name = fn_name
        self._expected_invocations = []
        self._reset_temp()

    def on(self, *args, **kwargs):
        self._temp_args = args
        self._temp_kwargs = kwargs
        return self

    def return_val(self, return_value):
        self._expected_invocations.append(
            {
                "args": self._temp_args,
                "kwargs": self._temp_kwargs,
                "return_value": return_value,
                "invoked": False,
            }
        )
        self._reset_temp()

    def raise_exception(self, exception):
        self._expected_invocations.append(
            {
                "args": self._temp_args,
                "kwargs": self._temp_kwargs,
                "exception": exception,
                "invoked": False,
            }
        )

    def assert_expectations(self):
        for expected in self._expected_invocations:
            args, kwargs = expected["args"], expected["kwargs"]
            assert expected[
                "invoked"
            ], f"Expected an invocation of function {self._fn_name} with args {args} and kwargs {kwargs}"

    def _side_effect(self, *args, **kwargs):
        for expected in self._expected_invocations:
            if self._args_match(expected["args"], args) and self._kwargs_match(
                expected["kwargs"], kwargs
            ):
                expected["invoked"] = True
                if "exception" in expected:
                    raise expected["exception"]
                else:
                    return expected["return_value"]
        raise Exception(
            f"Unexpected invocation of function '{self._fn_name}' with args: {list(args)} and kwargs: {kwargs}"
        )

    def _args_match(self, expected, actual):
        expected_lst, actual_lst = list(expected), list(actual)
        if len(expected) != len(actual):
            return False

        for index, expected_val in enumerate(expected_lst):
            if expected_val != actual_lst[index]:
                return False
        return True

    def _kwargs_match(self, expected, actual):
        for key in actual:
            if key not in expected:
                return False

        for key, expected_val in expected.items():
            if expected_val != actual.get(key, None):
                return False
        return True

    def _reset_temp(self):
        self._temp_args = ()
        self._temp_kwargs = {}


class UnexpectedFunctionInvocation(mock.MagicMock):
    def __init__(self, fn_name):
        super().__init__(side_effect=self._side_effect)
        self._fn_name = fn_name

    def _side_effect(self, *args, **kwargs):
        raise AttributeError(
            f"Unexpected call to function '{self._fn_name}' with args {list(args)} and kwargs {kwargs}"
        )


class MockApi:
    def on(self, fn_name, *args, **kwargs):
        if not self._hasattr(fn_name):
            setattr(self, fn_name, MockFunction(fn_name))
        mock_function = getattr(self, fn_name, None)
        return mock_function.on(*args, **kwargs)

    def assert_expectations(self):
        for attr in dir(self):
            if isinstance(getattr(self, attr), MockFunction):
                mock_function = getattr(self, attr)
                mock_function.assert_expectations()

    def _hasattr(self, name) -> bool:
        attr = getattr(self, name)
        return not isinstance(attr, UnexpectedFunctionInvocation)

    def __getattr__(self, name: str):
        return self.__dict__.get(name, UnexpectedFunctionInvocation(name))
