# Copyright (c) 2008-2016 Szczepan Faber, Serhiy Oplakanets, Herr Kaste
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from __future__ import annotations
import inspect
import operator
import types
import functools
from collections import deque
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, AsyncIterator, Callable, Iterable, Iterator, cast

from . import invocation, sameish, signature, utils
from . import verification as verificationModule
from .mock_registry import mock_registry
from .patching import Patch, patcher


__all__ = ['mock']

__tracebackhide__ = operator.methodcaller(
    "errisinstance",
    invocation.InvocationError
)
SUPPORTS_MARKCOROUTINEFUNCTION = hasattr(inspect, "markcoroutinefunction")

_CONFIG_ASYNC_PREFIX = "async "
_ASYNC_BY_PROTOCOL_METHODS = {"__aenter__", "__aexit__", "__anext__"}


class _Dummy:
    # We spell out `__call__` here for convenience. All other magic methods
    # must be configured before use, but we want `mock`s to be callable by
    # default.
    def __call__(self, *args, **kwargs):
        return self.__getattr__('__call__')(*args, **kwargs)  # type: ignore[attr-defined]  # noqa: E501


def remembered_invocation_builder(
    mock: Mock,
    method_name: str,
    discard_first_arg: bool,
    *args,
    **kwargs
):
    invoc = invocation.RememberedInvocation(
        mock, method_name, discard_first_arg=discard_first_arg
    )
    return invoc(*args, **kwargs)


ANSWER_SELECTOR_METHODS = {
    'thenReturn',
    'thenRaise',
    'thenAnswer',
    'thenCallOriginalImplementation',
}


@dataclass(frozen=True)
class Chain:
    theMock: Mock
    options: dict[str, Any]
    segments: tuple[Segment, ...] = ()

    def __add__(self, segment: Segment) -> "Chain":
        return Chain(self.theMock, self.options, self.segments + (segment,))

    def new_segment(self, name: str):
        return _chain_segment(self, name)

    def rollback(self) -> None:
        for segment in reversed(self.segments):
            segment.invoc.forget_self()


@dataclass(frozen=True)
class Segment:
    name: str
    invoc: invocation.StubbedInvocation
    answer_selector: invocation.AnswerSelector


def _chain_segment(chain: Chain, name: str):
    class SegmentFacade:
        def __call__(self, *args, **kwargs):
            if chain.segments and name in ANSWER_SELECTOR_METHODS:
                getattr(chain.segments[-1].answer_selector, name)(
                    *args,
                    **kwargs,
                )
                return _wait_for_attr(chain)

            segment = _materialize_method_segment(chain, name, args, kwargs)
            next_chain = chain + segment
            return _wait_for_attr(next_chain)

        def __getattr__(self, attr_name):
            try:
                segment = _materialize_property_segment(chain, name)
            except invocation.InvocationError:
                chain.rollback()
                raise

            next_chain = chain + segment
            return next_chain.new_segment(attr_name)

        def __enter__(self):
            chain.rollback()
            raise AttributeError('__enter__')

        def __exit__(self, *exc_info):
            chain.rollback()
            raise AttributeError('__exit__')

    return SegmentFacade()


def _wait_for_attr(chain: Chain):
    class WaitForAttr:
        def __getattr__(self, attr_name):
            return chain.new_segment(attr_name)

        __enter__ = chain.segments[-1].answer_selector.__enter__
        __exit__ = chain.segments[-1].answer_selector.__exit__

    return WaitForAttr()


def _materialize_method_segment(
    chain: Chain,
    name: str,
    args: tuple[Any, ...],
    kwargs: dict[str, Any],
) -> Segment:
    if not chain.segments:
        # Pre-checks here are strict-independent and validate fluent-DSL
        # *shape* (e.g. call-vs-attribute intent) before we patch anything.
        # Strict/spec validation remains in StubbedInvocation.__call__
        # (`ensure_mocked_object_has_method` + signature checks).
        _ensure_target_is_callable(chain.theMock, name)

        invoc = invocation.StubbedInvocation(chain.theMock, name, **chain.options)
        answer_selector = invoc(*args, **kwargs)
        return Segment(name, invoc, answer_selector)

    chain_mock, parent_invocation, verification = _transition_to_child_chain(
        chain.segments
    )
    invoc = invocation.StubbedInvocation(
        chain_mock,
        name,
        verification=verification,
        parent_invocation=parent_invocation,
    )
    answer_selector = invoc(*args, **kwargs)
    return Segment(name, invoc, answer_selector)


def _materialize_property_segment(
    chain: Chain,
    name: str,
) -> Segment:
    if not chain.segments:
        # Same rationale as method materialization above: these pre-checks are
        # strict-independent API-shape guards for property-style stubbing.
        # Strict/spec checks for property stubs live in
        # StubbedPropertyAccess.__call__ (`ensure_mocked_object_has_attribute`).
        _ensure_target_is_not_callable(chain.theMock, name)

        if not inspect.isclass(chain.theMock.mocked_obj):
            raise invocation.InvocationError(
                "Cannot stub property '%s' on an instance. "
                "Use class-level stubbing instead: "
                "when(%s).%s.thenReturn(...)."
                % (
                    name,
                    type(chain.theMock.mocked_obj).__name__,
                    name,
                )
            )

        invoc = invocation.StubbedPropertyAccess(
            chain.theMock,
            name,
            **chain.options,
        )
        answer_selector = invoc()
        return Segment(name, invoc, answer_selector)

    chain_mock, parent_invocation, verification = _transition_to_child_chain(
        chain.segments
    )
    invoc = invocation.StubbedPropertyAccess(
        chain_mock,
        name,
        verification=verification,
        parent_invocation=parent_invocation,
    )
    answer_selector = invoc()
    return Segment(name, invoc, answer_selector)


def _transition_to_child_chain(
    segments: tuple[Segment, ...],
) -> tuple[
    Mock,
    invocation.StubbedInvocation,
    verificationModule.VerificationMode | None,
]:
    parent_invocation = segments[-1].invoc
    continuation = parent_invocation.transition_to_chain()
    verification = parent_invocation.pop_verification()
    return continuation.chain_mock, continuation.invocation, verification


def _ensure_target_is_callable(theMock: Mock, method_name: str) -> None:
    target, was_in_spec = theMock._get_original_method_before_stub(method_name)

    # Missing attributes can still be added in loose mode.
    if not was_in_spec and target is None:
        return

    if _should_continue_with_stubbed_invocation(
        target,
        allow_classes=True,
        spec=theMock.spec,
    ):
        return

    raise invocation.InvocationError("'%s' is not callable." % method_name)


def _ensure_target_is_not_callable(theMock: Mock, method_name: str) -> None:
    spec = theMock.spec
    if spec is None:
        return

    try:
        value = inspect.getattr_static(spec, method_name)
    except AttributeError:
        if inspect.isclass(spec):
            try:
                value = getattr(spec, method_name)
            except AttributeError:
                return
        else:
            return

    if _should_continue_with_stubbed_invocation(value, spec=spec):
        raise invocation.InvocationError(
            f"expected an invocation of '{method_name}'"
        )


def _should_continue_with_stubbed_invocation(
    value: object,
    allow_classes: bool = False,
    spec: object | None = None,
) -> bool:
    if (
        inspect.isfunction(value)
        or inspect.ismethod(value)
        or inspect.isbuiltin(value)
        or isinstance(value, staticmethod)
        or isinstance(value, classmethod)
        or isinstance(value, functools.partialmethod)
        or isinstance(value, types.MethodDescriptorType)
        or isinstance(value, types.WrapperDescriptorType)
        or isinstance(value, types.ClassMethodDescriptorType)
    ):
        return True

    # For class specs, callable descriptors (objects implementing both
    # `__call__` and `__get__`) are generally meant to be stubbed through
    # the property path. For non-class specs (e.g. module attributes such as
    # `numpy.vstack`), `__get__` should not disqualify callable targets.
    treat_callable_descriptors_as_non_callable = inspect.isclass(spec)

    # Generic callable fallback, with optional handling for callable
    # descriptor-like objects (`__call__` + `__get__`).
    return (
        callable(value)
        and (allow_classes or not inspect.isclass(value))
        and (
            not treat_callable_descriptors_as_non_callable
            or not hasattr(value, '__get__')
        )
    )


class _mocked_property:
    def __init__(self, mock, method_name):
        self.mock = mock
        self.method_name = method_name

    def __get__(self, obj, type):
        # For property/descriptors, `thenCallOriginalImplementation()` must
        # call `original_descriptor.__get__(obj, type)`. Unlike regular method
        # stubs, these `obj/type` binding values are only available during this
        # attribute access path, so we keep them in temporary context.
        with self.mock.property_access_context(self.method_name, obj, type):
            return invocation.RememberedPropertyAccess(
                self.mock, self.method_name)()

    def __set__(self, obj, value):
        # Keep this wrapper a data descriptor so it wins over instance
        # __dict__ during reads.
        return None

    def __delete__(self, obj):
        return None


class Mock:
    def __init__(
        self,
        mocked_obj: object,
        strict: bool = True,
        spec: object | None = None
    ) -> None:
        self.mocked_obj = mocked_obj
        self.strict = strict
        self.spec = spec

        self.invocations: list[invocation.RealInvocation] = []
        self.stubbed_invocations: deque[invocation.StubbedInvocation] = deque()

        self._original_methods: dict[str, object | None] = {}
        self._methods_to_unstub: dict[str, Patch] = {}
        self._signatures_store: dict[str, signature.Signature | None] = {}
        self._property_access_context: \
            list[tuple[str, object | None, object]] = []

        self._continuations: dict[
            invocation.StubbedInvocation,
            invocation.ConfiguredContinuation,
        ] = {}

        self._observers: list = []
        self._methods_marked_as_coroutine: set[str] = set()

    def attach(self, observer) -> None:
        if observer not in self._observers:
            self._observers.append(observer)

    def detach(self, observer) -> None:
        try:
            self._observers.remove(observer)
        except ValueError:
            pass

    def remember(self, invocation: invocation.RealInvocation) -> None:
        self.invocations.append(invocation)
        for observer in self._observers:
            observer.update(invocation)

    def finish_stubbing(
        self, stubbed_invocation: invocation.StubbedInvocation
    ) -> None:
        self.stubbed_invocations.appendleft(stubbed_invocation)

    def clear_invocations(self) -> None:
        self.invocations = []

    def continuation_for(
        self, invoc: invocation.StubbedInvocation
    ) -> invocation.Continuation:
        continuation = self._continuations.get(invoc)
        if continuation is not None:
            return continuation

        for other in self._sameish_invocations(invoc):
            configured_continuation = self._continuations.get(other)
            if isinstance(configured_continuation, invocation.ValueContinuation):
                return configured_continuation

            if isinstance(configured_continuation, invocation.ChainContinuation):
                # We do not keep mixed continuation modes (`Value` and `Chain`)
                # alive at the same time. So the first chain continuation
                # can be returned immediately.
                return configured_continuation

        return invocation.UnconfiguredContinuation()

    def set_continuation(self, continuation: invocation.ConfiguredContinuation) -> None:
        self._continuations[continuation.invocation] = continuation

    def _sameish_invocations(
        self, same: invocation.StubbedInvocation
    ) -> list[invocation.StubbedInvocation]:
        """Find prior stubs that are signature-compatible.

        This is used only for continuation bookkeeping (value-vs-chain mode),
        not for runtime call dispatch. The comparison is structural and avoids
        executing matcher predicates, so `arg_that(...)` and other custom
        matchers cannot crash internal equivalence probing.

        Why this exists: repeated selectors such as

            when(cat).meow().purr()
            when(cat).meow().roll()

        should share the same root continuation for `meow()`.
        """
        sameish: list[invocation.StubbedInvocation] = []
        for invoc in self.stubbed_invocations:
            if invoc is same:
                continue

            if invoc.method_name != same.method_name:
                continue

            if self._invocations_are_sameish(invoc, same):
                sameish.append(invoc)

        return sameish

    def _invocations_are_sameish(
        self,
        left: invocation.StubbedInvocation,
        right: invocation.StubbedInvocation,
    ) -> bool:
        return sameish.invocations_are_sameish(left, right)

    def get_original_method(self, method_name: str) -> object | None:
        return self._original_methods.get(method_name, None)

    def peek_original_method(self, method_name: str) -> object | None:
        try:
            return self._original_methods[method_name]
        except KeyError:
            original_method, _ = self._get_original_method_before_stub(method_name)
            return original_method

    @contextmanager
    def property_access_context(
        self, method_name: str, obj: object | None, type_: object
    ):
        self._property_access_context.append((method_name, obj, type_))
        try:
            yield
        finally:
            self._property_access_context.pop()

    def get_current_property_access(
        self, method_name: str
    ) -> tuple[object | None, object]:
        for accessed_method_name, obj, type_ in reversed(
            self._property_access_context
        ):
            if accessed_method_name == method_name:
                return obj, type_

        raise RuntimeError(
            "Could not resolve property access context for '%s'."
            % method_name
        )

    # STUBBING

    def _get_original_method_before_stub(
        self, method_name: str
    ) -> tuple[object | None, bool]:
        """
        Looks up the original method on the `spec` object and returns it
        together with an indication of whether the method is found
        "directly" on the `spec` object.

        This is used to decide whether the method should be stored as an
        original_method and should therefore be replaced when unstubbing.
        """
        if self.spec is None:
            return None, False

        return utils.get_original_attribute(self.spec, method_name, default=None)

    def replace_method(
        self,
        method_name: str,
        original_method: object | None,
    ) -> Patch:
        discard_first_arg = self._takes_implicit_self_or_cls(original_method)

        def new_mocked_method(*args, **kwargs):
            return remembered_invocation_builder(
                self, method_name, discard_first_arg, *args, **kwargs
            )

        new_mocked_method.__name__ = method_name
        if original_method:
            new_mocked_method.__doc__ = original_method.__doc__
            new_mocked_method.__wrapped__ = original_method  # type: ignore[attr-defined]
            try:
                new_mocked_method.__module__ = original_method.__module__
            except AttributeError:
                pass

        if (
            self.method_expects_awaitable(method_name, original_method)
            and SUPPORTS_MARKCOROUTINEFUNCTION
        ):
            new_mocked_method = inspect.markcoroutinefunction(new_mocked_method)

        if inspect.ismethod(original_method):
            new_mocked_method = utils.newmethod(new_mocked_method, self.mocked_obj)

        if isinstance(original_method, staticmethod):
            new_mocked_method = staticmethod(new_mocked_method)
        elif isinstance(original_method, classmethod):
            new_mocked_method = classmethod(new_mocked_method)  # type: ignore[assignment]  # noqa: E501
        elif (
            inspect.isclass(self.mocked_obj)
            and inspect.isclass(original_method)  # TBC: Inner classes
        ):
            new_mocked_method = staticmethod(new_mocked_method)

        return patcher.patch_attribute(
            self.mocked_obj,
            method_name,
            new_mocked_method,
            allow_unstub_by_replacement=False,
        )

    def stub(self, method_name: str) -> None:
        try:
            self._methods_to_unstub[method_name]
        except KeyError:
            original_method, _ = self._get_original_method_before_stub(method_name)
            self._original_methods[method_name] = original_method
            self._methods_to_unstub[method_name] = self.replace_method(
                method_name,
                original_method,
            )

    def stub_property(self, method_name: str) -> None:
        try:
            self._methods_to_unstub[method_name]
        except KeyError:
            original_method, _ = self._get_original_method_before_stub(method_name)
            self._original_methods[method_name] = original_method
            self._methods_to_unstub[method_name] = patcher.patch_attribute(
                self.mocked_obj,
                method_name,
                _mocked_property(self, method_name),
                allow_unstub_by_replacement=False,
            )


    def forget_stubbed_invocation(
        self, invocation: invocation.StubbedInvocation
    ) -> None:
        assert invocation in self.stubbed_invocations

        self.stubbed_invocations.remove(invocation)
        self._continuations.pop(invocation, None)

        if not any(
            inv.method_name == invocation.method_name
            for inv in self.stubbed_invocations
        ):
            patch = self._methods_to_unstub.pop(invocation.method_name)
            patch.restore_and_unregister()

        if self.stubbed_invocations:
            return

        mock_registry.unstub(self.mocked_obj)

    def unstub_method(self, method_name: str) -> None:
        invocations = [
            invoc
            for invoc in self.stubbed_invocations
            if invoc.method_name == method_name
        ]
        if not invocations:
            return

        for invoc in invocations:
            invoc.forget_self()

        self.invocations = [
            invocation
            for invocation in self.invocations
            if invocation.method_name != method_name
        ]

    def unstub(self) -> None:
        while self._methods_to_unstub:
            _, patch = self._methods_to_unstub.popitem()
            patch.restore_and_unregister()
        self.stubbed_invocations = deque()
        self.invocations = []
        self._methods_marked_as_coroutine = set()
        self._continuations = {}

    # SPECCING

    def mark_as_coroutine(self, method_name: str) -> None:
        self._methods_marked_as_coroutine.add(method_name)

    def is_marked_as_coroutine(self, method_name: str) -> bool:
        return method_name in self._methods_marked_as_coroutine

    def method_expects_awaitable(
        self,
        method_name: str,
        original_method: object | None = None,
    ) -> bool:
        if original_method is None:
            original_method = self.peek_original_method(method_name)

        return (
            _is_coroutine_method(original_method)
            or self.is_marked_as_coroutine(method_name)
            or (
                self.spec is None
                and method_name in _ASYNC_BY_PROTOCOL_METHODS
            )
        )

    def has_method(self, method_name: str) -> bool:
        if self.spec is None:
            return True

        return hasattr(self.spec, method_name)

    def get_signature(self, method_name: str) -> signature.Signature | None:
        if self.spec is None:
            return None

        try:
            return self._signatures_store[method_name]
        except KeyError:
            sig = signature.get_signature(self.spec, method_name)
            self._signatures_store[method_name] = sig
            return sig

    def will_have_self_or_cls(self, method_name: str) -> bool:
        original_method = self.peek_original_method(method_name)
        return self._takes_implicit_self_or_cls(original_method)

    def _takes_implicit_self_or_cls(
        self, original_method: object | None
    ) -> bool:
        # If original_method is None, we *added* it to mocked_obj
        # and thus, it will eat self iff mocked_obj is a class.
        return (
            inspect.ismethod(original_method)
            or (
                inspect.isclass(self.mocked_obj)
                and not isinstance(original_method, staticmethod)
                and not inspect.isclass(original_method)
            )
        )


def _is_coroutine_method(method: object | None) -> bool:
    if isinstance(method, (staticmethod, classmethod)):
        method = method.__func__
    elif inspect.ismethod(method):
        method = method.__func__

    return inspect.iscoroutinefunction(method)


class _OMITTED(object):
    def __repr__(self):
        return 'OMITTED'


OMITTED = _OMITTED()

def mock(config_or_spec=None, spec=None, strict=OMITTED):  # noqa: C901
    """Create 'empty' objects ('Mocks').

    Will create an empty unconfigured object, that you can pass
    around. All interactions (method calls) will be recorded and can be
    verified using :func:`verify` et.al.

    A plain `mock()` will be not `strict`, and thus all methods regardless
    of the arguments will return ``None``.

    .. note:: Technically all attributes will return an internal interface.
        Because of that a simple ``if mock().foo:`` will surprisingly pass.

    If you set strict to ``True``: ``mock(strict=True)`` all unexpected
    interactions will raise an error instead.

    You configure a mock using :func:`when`, :func:`when2` or :func:`expect`.
    You can also very conveniently just pass in a dict here::

        response = mock({'text': 'ok', 'raise_for_status': lambda: None})

    You can also create an empty Mock which is specced against a given
    `spec`: ``mock(requests.Response)``. These mock are by default strict,
    thus they raise if you want to stub a method, the spec does not implement.
    Mockito will also match the function signature.

    You can pre-configure a specced mock as well::

        response = mock({'json': lambda: {'status': 'Ok'}},
                        spec=requests.Response)

    Mocks are by default callable. Configure the callable behavior using
    `when`::

        dummy = mock()
        when(dummy).__call__(1).thenReturn(2)

    All other magic methods must be configured this way or they will raise an
    AttributeError.


    See :func:`verify` to verify your interactions after usage.

    """

    if type(config_or_spec) is dict:
        config, spec = config_or_spec, spec
    else:
        config, spec = {}, spec or config_or_spec

    if strict is OMITTED:
        strict = False if spec is None else True


    class Dummy(_Dummy):
        if spec:
            __class__ = spec  # make isinstance work

        def __getattr__(self, method_name):
            if strict:
                __tracebackhide__ = operator.methodcaller(
                    "errisinstance", AttributeError
                )
                # deepcopy catches a possible AttributeError, fallback
                # to an arbitrary RuntimeError
                error_type = (
                    RuntimeError
                    if method_name == '__deepcopy__'
                    else AttributeError
                )
                raise error_type(
                    "'Dummy' has no attribute %r configured" % method_name)

            if (
                method_name != "__call__"
                and method_name == "__{}__".format(method_name[2:-2])
            ):
                raise AttributeError(method_name)

            # If a descriptor exists on the dummy class, resolve it here so
            # InvocationError from descriptor-backed stubs is not converted
            # into a dynamic fallback attribute.
            if method_name != "__call__":
                try:
                    class_attr = inspect.getattr_static(type(self), method_name)
                except AttributeError:
                    pass
                else:
                    if hasattr(class_attr, "__get__"):
                        try:
                            return class_attr.__get__(self, type(self))
                        except AttributeError as error:
                            if isinstance(error, invocation.InvocationError):
                                raise
                            # Keep dynamic-attribute behavior for descriptors that
                            # deliberately signal missing via AttributeError.

            def ad_hoc_function(*args, **kwargs):
                return remembered_invocation_builder(
                    theMock, method_name, False, *args, **kwargs
                )
            ad_hoc_function.__name__ = method_name
            ad_hoc_function.__self__ = obj  # type: ignore[attr-defined]
            if spec:
                try:
                    original_method = getattr(spec, method_name)
                    ad_hoc_function.__wrapped__ = original_method  # type: ignore[attr-defined]  # noqa: E501
                    ad_hoc_function.__doc__ = original_method.__doc__
                except AttributeError:
                    pass

            return ad_hoc_function

        def __repr__(self):
            name = 'Dummy'
            if spec:
                name += spec.__name__
            return "<%s id=%s>" % (name, id(self))


    # That's a tricky one: The object we will return is an *instance* of our
    # Dummy class, but the mock we register will point and patch the class.
    # T.i. so that magic methods (`__call__` etc.) can be configured.
    obj = Dummy()
    theMock = Mock(Dummy, strict=strict, spec=spec)

    normalized_names = {
        _normalize_config_key(raw_name)[0]
        for raw_name in config
    }

    for raw_name, value in config.items():
        _configure_mock_from_shorthand(
            theMock,
            Dummy,
            obj,
            raw_name,
            value,
            normalized_names,
        )

    mock_registry.register(obj, theMock)
    return obj


def _configure_mock_from_shorthand(
    theMock: Mock,
    Dummy: type,
    obj: object,
    raw_name: str,
    value: object,
    configured_names: set[str],
) -> None:
    method_name, marked_async = _normalize_config_key(raw_name)
    should_be_async = marked_async or method_name in _ASYNC_BY_PROTOCOL_METHODS

    if method_name in {"__enter__", "__aenter__"} and value is Ellipsis:
        _stub_from_shorthand(
            theMock,
            method_name,
            return_value=obj,
            force_async=(method_name == "__aenter__"),
        )

        companion_exit = "__aexit__" if method_name == "__aenter__" else "__exit__"
        if companion_exit not in configured_names:
            _stub_from_shorthand(
                theMock,
                companion_exit,
                return_value=False,
                force_async=(companion_exit == "__aexit__"),
            )
        return

    if method_name == "__iter__":
        iter_answer = _normalize_iter_answer(value)
        _stub_from_shorthand(theMock, method_name, answer=iter_answer)
        return

    if method_name == "__aiter__":
        aiter_answer = _normalize_aiter_answer(value)
        _stub_from_shorthand(theMock, method_name, answer=aiter_answer)
        return

    if inspect.isfunction(value):
        function_answer = _widen_zero_arg_callable(value)
        _stub_from_shorthand(
            theMock,
            method_name,
            answer=function_answer,
            force_async=should_be_async,
        )
        return

    if should_be_async:
        if value is Ellipsis:
            _stub_from_shorthand(theMock, method_name, force_async=True)
            return

        raise TypeError(
            "Async shorthand '%s' expects a function value or Ellipsis. "
            "Use `lambda: value` for fixed async return values."
            % raw_name
        )

    setattr(Dummy, method_name, value)


def _normalize_config_key(raw_name: str) -> tuple[str, bool]:
    if raw_name.startswith(_CONFIG_ASYNC_PREFIX):
        return raw_name[len(_CONFIG_ASYNC_PREFIX):], True
    return raw_name, False


def _stub_from_shorthand(
    theMock: Mock,
    method_name: str,
    *,
    answer: object = OMITTED,
    return_value: object = OMITTED,
    force_async: bool = False,
) -> None:
    if force_async:
        theMock.mark_as_coroutine(method_name)

    stubbed = invocation.StubbedInvocation(theMock, method_name)(Ellipsis)

    if answer is not OMITTED:
        stubbed.thenAnswer(answer)  # type: ignore[arg-type]
    elif return_value is not OMITTED:
        stubbed.thenReturn(return_value)


def _widen_zero_arg_callable(function: object):
    if not inspect.isfunction(function):
        return function

    try:
        params = inspect.signature(function).parameters
    except Exception:
        return function

    if params:
        return function

    def widened(*args, **kwargs):
        return function()

    widened.__name__ = function.__name__
    widened.__doc__ = function.__doc__
    return widened


def _normalize_iter_answer(value) -> Callable[..., Iterator[object]]:
    def answer(*args, **kwargs) -> Iterator[object]:
        result = value(*args, **kwargs) if callable(value) else value
        return iter(cast(Iterable[object], result))

    return answer


def _normalize_aiter_answer(value) -> Callable[..., AsyncIterator[object]]:
    def answer(*args, **kwargs) -> AsyncIterator[object]:
        result = value(*args, **kwargs) if callable(value) else value
        return _normalize_aiter_result(result)

    return answer


def _normalize_aiter_result(value) -> AsyncIterator[object]:
    if hasattr(value, "__anext__"):
        return cast(AsyncIterator[object], value)

    aiter = getattr(value, "__aiter__", None)
    if callable(aiter):
        candidate = aiter()
        if hasattr(candidate, "__anext__"):
            return cast(AsyncIterator[object], candidate)
        raise TypeError(
            "__aiter__() must return an async iterator implementing __anext__"
        )

    iterator = iter(cast(Iterable[object], value))

    async def generator() -> AsyncIterator[object]:
        for item in iterator:
            yield item

    return generator()
