import abc
import functools
import inspect
import itertools
import uuid
from contextlib import asynccontextmanager, contextmanager
from typing import Any, Callable, NamedTuple

import mlflow
import mlflow.utils.autologging_utils
from mlflow.entities.run_status import RunStatus
from mlflow.environment_variables import _MLFLOW_AUTOLOGGING_TESTING
from mlflow.exceptions import MlflowException
from mlflow.utils import gorilla, is_iterator
from mlflow.utils.autologging_utils import _logger
from mlflow.utils.autologging_utils.events import AutologgingEventLoggerWrapper
from mlflow.utils.autologging_utils.logging_and_warnings import (
    MlflowEventsAndWarningsBehaviorGlobally,
    NonMlflowWarningsBehaviorForCurrentThread,
)
from mlflow.utils.mlflow_tags import MLFLOW_AUTOLOGGING

_AUTOLOGGING_PATCHES = {}


# Function attribute used for testing purposes to verify that a given function
# has been wrapped with the `exception_safe_function_for_class` and
# `picklable_exception_safe_function` decorators
_ATTRIBUTE_EXCEPTION_SAFE = "exception_safe"


_ERROR_MSG = "Encountered unexpected error during {} autologging: {}"


def exception_safe_function_for_class(function):
    """
    Wraps the specified function with broad exception handling to guard
    against unexpected errors during autologging.
    Note this function creates an unpicklable function as `safe_function` is locally defined,
    but a class instance containing methods decorated by this function should be pickalable,
    because pickle only saves instance attributes, not methods.
    See https://docs.python.org/3/library/pickle.html#pickling-class-instances for more details.
    """
    if is_testing():
        setattr(function, _ATTRIBUTE_EXCEPTION_SAFE, True)

    def safe_function(*args, **kwargs):
        try:
            return function(*args, **kwargs)
        except Exception as e:
            if is_testing():
                raise
            else:
                _logger.warning("Encountered unexpected error during autologging: %s", e)

    return update_wrapper_extended(safe_function, function)


def _safe_function(function, *args, **kwargs):
    try:
        return function(*args, **kwargs)
    except Exception as e:
        if is_testing():
            raise
        else:
            _logger.warning("Encountered unexpected error during autologging: %s", e)


def picklable_exception_safe_function(function):
    """
    Wraps the specified function with broad exception handling to guard
    against unexpected errors during autologging while preserving picklability.
    """
    if is_testing():
        setattr(function, _ATTRIBUTE_EXCEPTION_SAFE, True)

    return update_wrapper_extended(functools.partial(_safe_function, function), function)


def _exception_safe_class_factory(base_class):
    """
    Creates an exception safe metaclass that inherits from `base_class`.
    """

    class _ExceptionSafeClass(base_class):
        """
        Metaclass that wraps all functions defined on the specified class with broad error handling
        logic to guard against unexpected errors during autlogging.

        Rationale: Patched autologging functions commonly pass additional class instances as
        arguments to their underlying original training routines; for example, Keras autologging
        constructs a subclass of `keras.callbacks.Callback` and forwards it to `Model.fit()`.
        To prevent errors encountered during method execution within such classes from disrupting
        model training, this metaclass wraps all class functions in a broad try / catch statement.

        Note: `ExceptionSafeClass` does not handle exceptions in class methods or static methods,
        as these are not always Python callables and are difficult to wrap
        """

        def __new__(cls, name, bases, dct):
            for m in dct:
                # class methods or static methods are not callable.
                if callable(dct[m]):
                    dct[m] = exception_safe_function_for_class(dct[m])
            return base_class.__new__(cls, name, bases, dct)

    return _ExceptionSafeClass


ExceptionSafeClass = _exception_safe_class_factory(type)

# `ExceptionSafeClass` causes an error when used with an abstract class.
#
# ```
# class AbstractClass(abc.ABC):
#    ...
#
# class DerivedClass(AbstractClass, metaclass=ExceptionSafeClass):
#    ...
# ```
#
# This raises:
#
# ```
# TypeError: metaclass conflict: the metaclass of a derived class must be
#            a (non-strict) subclass of the metaclasses of all its bases.
# ```
#
# To avoid this error, create `ExceptionSafeAbstractClass` that is based on `abc.ABCMeta`.
ExceptionSafeAbstractClass = _exception_safe_class_factory(abc.ABCMeta)


def with_managed_run(autologging_integration, patch_function, tags=None):
    """Given a `patch_function`, returns an `augmented_patch_function` that wraps the execution of
    `patch_function` with an active MLflow run. The following properties apply:

        - An MLflow run is only created if there is no active run present when the
          patch function is executed

        - If an active run is created by the `augmented_patch_function`, it is terminated
          with the `FINISHED` state at the end of function execution

        - If an active run is created by the `augmented_patch_function`, it is terminated
          with the `FAILED` if an unhandled exception is thrown during function execution

    Note that, if nested runs or non-fluent runs are created by `patch_function`, `patch_function`
    is responsible for terminating them by the time it terminates
    (or in the event of an exception).

    Args:
        autologging_integration: The autologging integration associated
            with the `patch_function`.
        patch_function: A function object compatible with `safe_patch`.
        tags: A dictionary of string tags to set on each managed run created during the
            execution of `patch_function`.
    """
    from mlflow.tracking.fluent import active_run
    from mlflow.utils.autologging_utils import _has_active_training_session

    def create_managed_run():
        managed_run = mlflow.start_run(tags=tags)
        _logger.info(
            "Created MLflow autologging run with ID '%s', which will track hyperparameters,"
            " performance metrics, model artifacts, and lineage information for the"
            " current %s workflow",
            managed_run.info.run_id,
            autologging_integration,
        )
        return managed_run

    def patch_with_managed_run(original, *args, **kwargs):
        managed_run = None
        # If there is an active training session but there is no active run
        # in current thread, it means the thread is spawned by `estimator.fit`
        # as a worker thread, we should disable autologging in
        # these worker threads, so skip creating managed run.
        if not active_run() and not _has_active_training_session():
            managed_run = create_managed_run()

        try:
            result = patch_function(original, *args, **kwargs)
        except (Exception, KeyboardInterrupt):
            # In addition to standard Python exceptions, handle keyboard interrupts to ensure
            # that runs are terminated if a user prematurely interrupts training execution
            # (e.g. via sigint / ctrl-c)
            if managed_run:
                mlflow.end_run(RunStatus.to_string(RunStatus.FAILED))
            raise
        else:
            if managed_run:
                mlflow.end_run(RunStatus.to_string(RunStatus.FINISHED))
            return result

    return patch_with_managed_run


def is_testing():
    """
    Indicates whether or not autologging functionality is running in test mode (as determined
    by the `MLFLOW_AUTOLOGGING_TESTING` environment variable). Test mode performs additional
    validation during autologging, including:

        - Checks for the exception safety of arguments passed to model training functions
          (i.e. all additional arguments should be "exception safe" functions or classes)
        - Disables exception handling for patched function logic, ensuring that patch code
          executes without errors during testing
    """
    return _MLFLOW_AUTOLOGGING_TESTING.get()


def _resolve_extra_tags(autologging_integration, extra_tags):
    tags = {MLFLOW_AUTOLOGGING: autologging_integration}
    if extra_tags:
        if isinstance(extra_tags, dict):
            if MLFLOW_AUTOLOGGING in extra_tags:
                extra_tags.pop(MLFLOW_AUTOLOGGING)
                _logger.warning(
                    f"Tag `{MLFLOW_AUTOLOGGING}` is ignored as it is a reserved tag by MLflow "
                    f"autologging."
                )
            tags.update(extra_tags)
        else:
            raise mlflow.exceptions.MlflowException.invalid_parameter_value(
                f"Invalid `extra_tags` type: expecting dictionary, "
                f"received `{type(extra_tags).__name__}`"
            )
    return tags


def safe_patch(
    autologging_integration,
    destination,
    function_name,
    patch_function,
    manage_run=False,
    extra_tags=None,
):
    """Patches the specified `function_name` on the specified `destination` class for autologging
    purposes, preceding its implementation with an error-safe copy of the specified patch
    `patch_function` with the following error handling behavior:
        - Exceptions thrown from the underlying / original function
          (`<destination>.<function_name>`) are propagated to the caller.
        - Exceptions thrown from other parts of the patched implementation (`patch_function`)
          are caught and logged as warnings.

    Args:
        autologging_integration: The name of the autologging integration associated with the
            patch.
        destination: The Python class on which the patch is being defined.
        function_name: The name of the function to patch on the specified `destination` class.
        patch_function: The patched function code to apply. The first argument should be reserved
            for an `original` argument representing the underlying / original function. Subsequent
            arguments should be identical to those of the original function being patched.
        manage_run: If `True`, applies the `with_managed_run` wrapper to the specified
            `patch_function`, which automatically creates & terminates an MLflow
            active run during patch code execution if necessary. If `False`,
            does not apply the `with_managed_run` wrapper to the specified
            `patch_function`.
        extra_tags: A dictionary of extra tags to set on each managed run created by autologging.
    """
    from mlflow.tracking.fluent import active_run
    from mlflow.utils.autologging_utils import autologging_is_disabled, get_autologging_config

    # NB: Checking the signature of the patch function rather than original, so that we don't
    # accidentally change the behavior of existing patches that may use sync patch function
    # for async original functions (e.g. LangChain).
    is_async_function = inspect.iscoroutinefunction(patch_function)

    if manage_run:
        if is_async_function:
            raise MlflowException("manage_run parameter is not supported for async functions.")

        tags = _resolve_extra_tags(autologging_integration, extra_tags)
        patch_function = with_managed_run(
            autologging_integration,
            patch_function,
            tags=tags,
        )

    original_fn = gorilla.get_original_attribute(
        destination, function_name, bypass_descriptor_protocol=False
    )
    # Retrieve raw attribute while bypassing the descriptor protocol
    raw_original_obj = gorilla.get_original_attribute(
        destination, function_name, bypass_descriptor_protocol=True
    )
    if original_fn != raw_original_obj:
        raise RuntimeError(f"Unsupported patch on {destination}.{function_name}")
    elif isinstance(original_fn, property):
        if is_async_function:
            raise MlflowException("Patching async property methods is not supported.")

        is_property_method = True

        # For property decorated methods (a kind of method delegation), e.g.
        # class A:
        #   @property
        #   def f1(self):
        #     ...
        #     return delegated_f1
        #
        # suppose `a1` is an instance of class `A`,
        # `A.f1.fget` will get the original `def f1(self)` method,
        # and `A.f1.fget(a1)` will be equivalent to `a1.f1()` and
        # its return value will be the `delegated_f1` function.
        # So using the `property.fget` we can construct the (delegated) "original_fn"
        def original(self, *args, **kwargs):
            # the `original_fn.fget` will get the original method decorated by `property`
            # the `original_fn.fget(self)` will get the delegated function returned by the
            # property decorated method.
            bound_delegate_method = original_fn.fget(self)
            return bound_delegate_method(*args, **kwargs)

    else:
        original = original_fn
        is_property_method = False

    def safe_patch_function(*args, **kwargs):
        """
        A safe wrapper around the specified `patch_function` implementation designed to
        handle exceptions thrown during the execution of `patch_function`. This wrapper
        distinguishes exceptions thrown from the underlying / original function
        (`<destination>.<function_name>`) from exceptions thrown from other parts of
        `patch_function`. This distinction is made by passing an augmented version of the
        underlying / original function to `patch_function` that uses nonlocal state to track
        whether or not it has been executed and whether or not it threw an exception.
        Exceptions thrown from the underlying / original function are propagated to the caller,
        while exceptions thrown from other parts of `patch_function` are caught and logged as
        warnings.

        NB: PLEASE BE SUPER CAREFUL WHEN MODIFYING THIS FUNCTION. IT IS USED IN A WIDE VARIETY
        OF CONTEXTX AND CRITICAL PATH IN DBR/MLR BY DEFAULT. ANY BUG HERE CAN BREAK USERS'
        WORKLOAD WITHOUT THEM TAKING ANY ACTION.
        """
        # Reroute warnings encountered during the patch function implementation to an MLflow event
        # logger, and enforce silent mode if applicable (i.e. if the corresponding autologging
        # integration was called with `silent=True`), hiding MLflow event logging statements and
        # hiding all warnings in the autologging preamble and postamble (i.e. the code surrounding
        # the user's original / underlying ML function). Non-MLflow warnings are enabled during the
        # execution of the original / underlying ML function
        #
        # Note that we've opted *not* to apply this context manager as a decorator on
        # `safe_patch_function` because the context-manager-as-decorator pattern uses
        # `contextlib.ContextDecorator`, which creates generator expressions that cannot be pickled
        # during model serialization by ML frameworks such as scikit-learn
        is_silent_mode = get_autologging_config(autologging_integration, "silent", False)
        with (
            MlflowEventsAndWarningsBehaviorGlobally(
                # MLflow warnings emitted during autologging training sessions are likely not
                # actionable and result from the autologging implementation invoking another MLflow
                # API. Accordingly, we reroute these warnings to the MLflow event logger with level
                # WARNING For reference, see recommended warning and event logging behaviors from
                # https://docs.python.org/3/howto/logging.html#when-to-use-logging
                reroute_warnings=True,
                disable_event_logs=is_silent_mode,
                disable_warnings=is_silent_mode,
            ),
            NonMlflowWarningsBehaviorForCurrentThread(
                # non-MLflow Warnings emitted during the autologging preamble (before the original /
                # underlying ML function is called) and postamble (after the original / underlying
                # ML function is called) are likely not actionable and result from the autologging
                # implementation invoking an API from a dependent library. Accordingly, we reroute
                # these warnings to the MLflow event logger with level WARNING. For reference, see
                # recommended warning and event logging behaviors from
                # https://docs.python.org/3/howto/logging.html#when-to-use-logging
                reroute_warnings=True,
                disable_warnings=is_silent_mode,
            ),
        ):
            if is_testing():
                preexisting_run_for_testing = active_run()

            # Whether or not to exclude autologged content from user-created fluent runs
            # (i.e. runs created manually via `mlflow.start_run()`)
            exclusive = get_autologging_config(autologging_integration, "exclusive", False)
            user_created_fluent_run_is_active = (
                active_run() and not _AutologgingSessionManager.active_session()
            )
            active_session_failed = (
                _AutologgingSessionManager.active_session() is not None
                and _AutologgingSessionManager.active_session().state == "failed"
            )

            if (
                active_session_failed
                or autologging_is_disabled(autologging_integration)
                or (user_created_fluent_run_is_active and exclusive)
                or (
                    mlflow.utils.autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED
                    and autologging_integration
                )
            ):
                # If the autologging integration associated with this patch is disabled,
                # or if the current autologging integration is in exclusive mode and a user-created
                # fluent run is active, call the original function and return. Restore the original
                # warning behavior during original function execution, since autologging is being
                # skipped
                with NonMlflowWarningsBehaviorForCurrentThread(
                    disable_warnings=False,
                    reroute_warnings=False,
                ):
                    return original(*args, **kwargs)

            # Whether or not the original / underlying function has been called during the
            # execution of patched code
            original_has_been_called = False
            # The value returned by the call to the original / underlying function during
            # the execution of patched code
            original_result = None
            # Whether or not an exception was raised from within the original / underlying function
            # during the execution of patched code
            failed_during_original = False
            # The active MLflow run (if any) associated with patch code execution
            patch_function_run_for_testing = None
            # The exception raised during executing patching function
            patch_error = None

            with _AutologgingSessionManager.start_session(autologging_integration) as session:
                event_logger = AutologgingEventLoggerWrapper(session, destination, function_name)

                def call_original_fn_with_event_logging(original_fn, og_args, og_kwargs):
                    try:
                        event_logger.log_original_function_start(og_args, og_kwargs)

                        original_fn_result = original_fn(*og_args, **og_kwargs)

                        event_logger.log_original_function_success(og_args, og_kwargs)
                        return original_fn_result
                    except Exception as e:
                        event_logger.log_original_function_error(og_args, og_kwargs, e)

                        nonlocal failed_during_original
                        failed_during_original = True
                        raise

                try:

                    def call_original(*og_args, **og_kwargs):
                        def _original_fn(*_og_args, **_og_kwargs):
                            if is_testing():
                                _validate_args(
                                    autologging_integration,
                                    function_name,
                                    args,
                                    kwargs,
                                    og_args,
                                    og_kwargs,
                                )
                                # By the time `original` is called by the patch implementation, we
                                # assume that either: 1. the patch implementation has already
                                # created an MLflow run or 2. the patch code will not create an
                                # MLflow run during the current execution. Here, we capture a
                                # reference to the active run, which we will use later on to
                                # determine whether or not the patch implementation created
                                # a run and perform validation if necessary
                                nonlocal patch_function_run_for_testing
                                patch_function_run_for_testing = active_run()

                            nonlocal original_has_been_called
                            original_has_been_called = True

                            nonlocal original_result
                            # Show all non-MLflow warnings as normal (i.e. not as event logs)
                            # during original function execution, even if silent mode is enabled
                            # (`silent=True`), since these warnings originate from the ML framework
                            # or one of its dependencies and are likely relevant to the caller
                            with NonMlflowWarningsBehaviorForCurrentThread(
                                disable_warnings=False,
                                reroute_warnings=False,
                            ):
                                original_result = original(*_og_args, **_og_kwargs)
                                return original_result

                        return call_original_fn_with_event_logging(_original_fn, og_args, og_kwargs)

                    # Apply the name, docstring, and signature of `original` to `call_original`.
                    # This is important because several autologging patch implementations inspect
                    # the signature of the `original` argument during execution
                    call_original = update_wrapper_extended(call_original, original)

                    event_logger.log_patch_function_start(args, kwargs)

                    patch_function(call_original, *args, **kwargs)

                    session.state = "succeeded"
                    event_logger.log_patch_function_success(args, kwargs)

                except Exception as e:
                    session.state = "failed"
                    patch_error = e
                    # Exceptions thrown during execution of the original function should be
                    # propagated to the caller. Additionally, exceptions encountered during test
                    # mode should be reraised to detect bugs in autologging implementations
                    if failed_during_original or is_testing():
                        raise

                if is_testing() and not preexisting_run_for_testing:
                    # If an MLflow run was created during the execution of patch code, verify that
                    # it is no longer active and that it contains expected autologging tags
                    assert not active_run(), (
                        f"Autologging integration {autologging_integration} leaked an active run"
                    )
                    if patch_function_run_for_testing:
                        _validate_autologging_run(
                            autologging_integration, patch_function_run_for_testing.info.run_id
                        )
                try:
                    if original_has_been_called:
                        return original_result
                    else:
                        return call_original_fn_with_event_logging(original, args, kwargs)
                finally:
                    # If original function succeeds, but `patch_function_exception` exists,
                    # it represent patching code unexpected failure, so we call
                    # `log_patch_function_error` in this case.
                    # If original function failed, we don't call `log_patch_function_error`
                    # even if `patch_function_exception` exists, because original function failure
                    # means there's some error in user code (e.g. user provide wrong arguments)
                    if patch_error is not None and not failed_during_original:
                        event_logger.log_patch_function_error(args, kwargs, patch_error)
                        _logger.warning(_ERROR_MSG.format(autologging_integration, patch_error))

    async def async_safe_patch_function(*args, **kwargs):
        """
        Async version of safe_patch_function.

        This code brainlessly copies the synchronous version of the function, but with async
        context managers and async functions. This is done to avoid the risk of introducing
        any bugs or regressions in the async version of the function. Note that we need to
        be really careful here, because autologging is enabled by-default in DBR/MLR, hence
        any bug here can break users' workload without them taking any action.

        That said, some long comments are omitted in this version to avoid redundancy. If
        you want to understand the context of the code better, please refer to the
        synchronous version as well.
        """
        is_silent_mode = get_autologging_config(autologging_integration, "silent", False)
        async with (
            MlflowEventsAndWarningsBehaviorGlobally(
                reroute_warnings=True,
                disable_event_logs=is_silent_mode,
                disable_warnings=is_silent_mode,
            ),
            NonMlflowWarningsBehaviorForCurrentThread(
                disable_warnings=is_silent_mode,
                reroute_warnings=True,
            ),
        ):
            if is_testing():
                preexisting_run_for_testing = active_run()

            # Whether or not to exclude autologged content from user-created fluent runs
            # (i.e. runs created manually via `mlflow.start_run()`)
            exclusive = get_autologging_config(autologging_integration, "exclusive", False)
            user_created_fluent_run_is_active = (
                active_run() and not _AutologgingSessionManager.active_session()
            )
            active_session_failed = (
                _AutologgingSessionManager.active_session() is not None
                and _AutologgingSessionManager.active_session().state == "failed"
            )

            if (
                active_session_failed
                or autologging_is_disabled(autologging_integration)
                or (user_created_fluent_run_is_active and exclusive)
                or (
                    mlflow.utils.autologging_utils._AUTOLOGGING_GLOBALLY_DISABLED
                    and autologging_integration
                )
            ):
                async with NonMlflowWarningsBehaviorForCurrentThread(False, False):
                    return await original(*args, **kwargs)

            original_has_been_called = False
            original_result = None
            failed_during_original = False
            patch_function_run_for_testing = None
            patch_error = None

            async with _AutologgingSessionManager.astart_session(
                autologging_integration
            ) as session:
                event_logger = AutologgingEventLoggerWrapper(session, destination, function_name)

                async def call_original_fn_with_event_logging(original_fn, og_args, og_kwargs):
                    try:
                        event_logger.log_original_function_start(og_args, og_kwargs)
                        original_fn_result = await original_fn(*og_args, **og_kwargs)
                        event_logger.log_original_function_success(og_args, og_kwargs)
                        return original_fn_result
                    except Exception as e:
                        event_logger.log_original_function_error(og_args, og_kwargs, e)
                        nonlocal failed_during_original
                        failed_during_original = True
                        raise

                try:

                    async def call_original(*og_args, **og_kwargs):
                        async def _original_fn(*_og_args, **_og_kwargs):
                            if is_testing():
                                _validate_args(
                                    autologging_integration,
                                    function_name,
                                    args,
                                    kwargs,
                                    og_args,
                                    og_kwargs,
                                )
                                nonlocal patch_function_run_for_testing
                                patch_function_run_for_testing = active_run()

                            nonlocal original_has_been_called
                            original_has_been_called = True

                            nonlocal original_result
                            async with NonMlflowWarningsBehaviorForCurrentThread(False, False):
                                original_result = await original(*_og_args, **_og_kwargs)
                                return original_result

                        return await call_original_fn_with_event_logging(
                            _original_fn, og_args, og_kwargs
                        )

                    # Apply the name, docstring, and signature of `original` to `call_original`.
                    # This is important because several autologging patch implementations inspect
                    # the signature of the `original` argument during execution
                    call_original = update_wrapper_extended(call_original, original)

                    event_logger.log_patch_function_start(args, kwargs)

                    await patch_function(call_original, *args, **kwargs)

                    session.state = "succeeded"
                    event_logger.log_patch_function_success(args, kwargs)

                except Exception as e:
                    session.state = "failed"
                    patch_error = e
                    # Exceptions thrown during execution of the original function should be
                    # propagated to the caller. Additionally, exceptions encountered during test
                    # mode should be reraised to detect bugs in autologging implementations
                    if failed_during_original or is_testing():
                        raise

                if is_testing() and not preexisting_run_for_testing:
                    # If an MLflow run was created during the execution of patch code, verify that
                    # it is no longer active and that it contains expected autologging tags
                    assert not active_run(), (
                        f"Autologging integration {autologging_integration} leaked an active run"
                    )
                    if patch_function_run_for_testing:
                        _validate_autologging_run(
                            autologging_integration, patch_function_run_for_testing.info.run_id
                        )
                try:
                    if original_has_been_called:
                        return original_result
                    else:
                        return await call_original_fn_with_event_logging(original, args, kwargs)
                finally:
                    if patch_error is not None and not failed_during_original:
                        event_logger.log_patch_function_error(args, kwargs, patch_error)
                        _logger.warning(_ERROR_MSG.format(autologging_integration, patch_error))

    if is_property_method:
        # Create a patched function (also property decorated)
        # like:
        #
        # class A:
        # @property
        # def get_bound_safe_patch_fn(self):
        #   original_fn.fget(self) # do availability check
        #   return bound_safe_patch_fn
        #
        # Suppose `a1` is instance of class A,
        # then `a1.get_bound_safe_patch_fn(*args, **kwargs)` will be equivalent to
        # `bound_safe_patch_fn(*args, **kwargs)`
        def get_bound_safe_patch_fn(self):
            # This `original_fn.fget` call is for availability check, if it raise error
            # then `hasattr(obj, {func_name})` will return False
            # so it mimic the original property behavior.
            original_fn.fget(self)

            def bound_safe_patch_fn(*args, **kwargs):
                return safe_patch_function(self, *args, **kwargs)

            # Make bound method `instance.target_method` keep the same doc and signature.
            # Here return the bound safe patch function because user call property decorated
            # method will like `instance.property_decorated_method(...)`, and internally it will
            # call the `bound_safe_patch_fn`, the argument list don't include the `self` argument,
            # so return bound function here.
            return update_wrapper_extended(bound_safe_patch_fn, original_fn.fget)

        # Make unbound method `class.target_method` keep the same doc and signature
        get_bound_safe_patch_fn = update_wrapper_extended(get_bound_safe_patch_fn, original_fn.fget)
        safe_patch_obj = property(get_bound_safe_patch_fn)
    elif is_async_function:
        safe_patch_obj = update_wrapper_extended(async_safe_patch_function, original)
    else:
        safe_patch_obj = update_wrapper_extended(safe_patch_function, original)

    new_patch = _wrap_patch(destination, function_name, safe_patch_obj)
    _store_patch(autologging_integration, new_patch)


def revert_patches(autologging_integration):
    """Reverts all patches on the specified destination class for autologging disablement purposes.

    Args:
        autologging_integration: The name of the autologging integration associated with the
            patch. Note: If called via fluent api (`autologging_integration="mlflow"`), then revert
            all patches for all active autologging integrations.

    """
    for patch in _AUTOLOGGING_PATCHES.get(autologging_integration, []):
        gorilla.revert(patch)

    _AUTOLOGGING_PATCHES.pop(autologging_integration, None)


# Represents an active autologging session using two fields:
# - integration: the name of the autologging integration corresponding to the session
# - id: a unique session identifier (e.g., a UUID)
# - state: the state of AutologgingSession, will be one of running/succeeded/failed
class AutologgingSession:
    def __init__(self, integration, id_):
        self.integration = integration
        self.id = id_
        self.state = "running"


class _AutologgingSessionManager:
    _session = None

    @classmethod
    @contextmanager
    def start_session(cls, integration):
        try:
            prev_session = cls._session
            if prev_session is None:
                session_id = uuid.uuid4().hex
                cls._session = AutologgingSession(integration, session_id)
            yield cls._session
        finally:
            # Only end the session upon termination of the context if we created
            # the session; otherwise, leave the session open for later termination
            # by its creator
            if prev_session is None:
                cls._end_session()

    @classmethod
    @asynccontextmanager
    async def astart_session(cls, integration):
        try:
            prev_session = cls._session
            if prev_session is None:
                session_id = uuid.uuid4().hex
                cls._session = AutologgingSession(integration, session_id)
            yield cls._session
        finally:
            if prev_session is None:
                cls._end_session()

    @classmethod
    def active_session(cls):
        return cls._session

    @classmethod
    def _end_session(cls):
        cls._session = None


def update_wrapper_extended(wrapper, wrapped):
    """Update a `wrapper` function to look like the `wrapped` function. This is an extension of
    `functools.update_wrapper` that applies the docstring *and* signature of `wrapped` to
    `wrapper`, producing a new function.

    Returns:
        A new function with the same implementation as `wrapper` and the same docstring
        & signature as `wrapped`.
    """
    updated_wrapper = functools.update_wrapper(wrapper, wrapped)
    # Assign the signature of the `wrapped` function to the updated wrapper function.
    # Certain frameworks may disallow signature inspection, causing `inspect.signature()` to throw.
    # One such example is the `tensorflow.estimator.Estimator.export_savedmodel()` function
    try:
        updated_wrapper.__signature__ = inspect.signature(wrapped)
    except Exception:
        _logger.debug("Failed to restore original signature for wrapper around %s", wrapped)
    return updated_wrapper


def _wrap_patch(destination, name, patch_obj, settings=None):
    """Apply a patch.

    Args:
        destination: Patch destination.
        name: Name of the attribute at the destination.
        patch_obj: Patch object, it should be a function or a property decorated function
            to be assigned to the patch point {destination}.{name}.
        settings: Settings for gorilla.Patch.

    """
    if settings is None:
        settings = gorilla.Settings(allow_hit=True, store_hit=True)

    patch = gorilla.Patch(destination, name, patch_obj, settings=settings)
    gorilla.apply(patch)
    return patch


def _store_patch(autologging_integration, patch):
    """
    Stores a patch for a specified autologging_integration class. Later to be used for being able
    to revert the patch when disabling autologging.

    Args:
        autologging_integration: The name of the autologging integration associated with the
            patch.
        patch: The patch to be stored.
    """
    if autologging_integration in _AUTOLOGGING_PATCHES:
        _AUTOLOGGING_PATCHES[autologging_integration].add(patch)
    else:
        _AUTOLOGGING_PATCHES[autologging_integration] = {patch}


def _validate_autologging_run(autologging_integration, run_id):
    """
    For testing purposes, verifies that an MLflow run produced by an `autologging_integration`
    satisfies the following properties:

        - The run has an autologging tag whose value is the name of the autologging integration
        - The run has a terminal status (e.g., KILLED, FAILED, FINISHED)
    """
    from mlflow.tracking.client import MlflowClient

    client = MlflowClient()
    run = client.get_run(run_id)
    autologging_tag_value = run.data.tags.get(MLFLOW_AUTOLOGGING)
    assert autologging_tag_value == autologging_integration, (
        f"Autologging run with id {run_id} failed to set autologging tag with expected value. "
        f"Expected: '{autologging_integration}', Actual: '{autologging_tag_value}'"
    )
    assert RunStatus.is_terminated(RunStatus.from_string(run.info.status)), (
        f"Autologging run with id {run_id} has a non-terminal status '{run.info.status}'"
    )


class ValidationExemptArgument(NamedTuple):
    """
    A NamedTuple representing the properties of an argument that is exempt from validation

    autologging_integration: The name of the autologging integration.
    function_name: The name of the function that is being validated.
    type_function: A Callable that accepts an object and returns True if the given object matches
                   the argument type. Returns False otherwise.
    positional_argument_index: The index of the argument in the function signature.
    keyword_argument_name: The name of the argument in the function signature.
    """

    autologging_integration: str
    function_name: str
    type_function: Callable[..., Any]
    positional_argument_index: int | None = None
    keyword_argument_name: str | None = None

    def matches(
        self,
        autologging_integration,
        function_name,
        value,
        argument_index=None,
        argument_name=None,
    ):
        """
        This method checks if the properties provided through the function arguments matches the
        properties defined in the NamedTuple.

        Args:
            autologging_integration: The name of an autologging integration.
            function_name: The name of the function that is being matched.
            value: The value of the argument.
            argument_index: The index of the argument, if it is passed as a positional
                argument. Otherwise it is None.
            argument_name: The name of the argument, if it is passed as a keyword
                argument. Otherwise it is None.

        Returns:
            Returns True if the given function properties matches the exempt argument's
            properties. Returns False otherwise.
        """
        return (
            self.autologging_integration == autologging_integration
            and self.function_name == function_name
            and (
                self.positional_argument_index == argument_index
                or self.keyword_argument_name == argument_name
            )
            and self.type_function(value)
        )


# WARNING: Exemptions should NOT be introduced unless absolutely necessary. If deemed necessary,
#          clear reasons must be provided as comment in addition to thorough integration tests.
_VALIDATION_EXEMPT_ARGUMENTS = [
    # When extracting implicitly defined `batch_size` in the case that `x` is a generator or a
    # generator class, we need to consume and restore the first element back to the generator to
    # calculate the `batch_size`. This means that:
    # 1. The type of `x` will become 'generator' regardless if user provided `x` as a generator or a
    #    custom generator class.
    # 2. The instance of `x` will be different, since we reconstructed the generator after consuming
    #    the first element.
    ValidationExemptArgument("tensorflow", "fit", is_iterator, 1, "x"),
    ValidationExemptArgument("keras", "fit", is_iterator, 1, "x"),
]


def _is_arg_exempt_from_validation(
    autologging_integration,
    function_name,
    argument,
    argument_index=None,
    argument_name=None,
):
    """This function is responsible for determining whether or not an argument is exempt from
    autolog safety validations. This includes both type checking and immutable checking.

    Args:
        autologging_integration: The name of the autologging integration.
        function_name: The name of the function that is being validated.
        argument: The actual argument.
        argument_index: The index of the argument, if it is passed as a positional
            argument. Otherwise it is None.
        argument_name: The name of the argument, if it is passed as a keyword argument.
            Otherwise it is None.

    Returns:
        True or False
    """
    return any(
        exemption.matches(
            autologging_integration,
            function_name,
            argument,
            argument_index,
            argument_name,
        )
        for exemption in _VALIDATION_EXEMPT_ARGUMENTS
    )


def _validate_args(
    autologging_integration,
    function_name,
    user_call_args,
    user_call_kwargs,
    autologging_call_args,
    autologging_call_kwargs,
):
    """
    Used for testing purposes to verify that, when a patched ML function calls its underlying
    / original ML function, the following properties are satisfied:

        - All arguments supplied to the patched ML function are forwarded to the
          original ML function
        - Any additional arguments supplied to the original function are exception safe (i.e.
          they are either functions decorated with the `@exception_safe_function_for_class` or
          `@pickalable_exception_safe_function` decorators, or classes / instances of classes with
          type `ExceptionSafeClass`
    """

    def _validate_new_input(inp):
        """
        Validates a new input (arg or kwarg) introduced to the underlying / original ML function
        call during the execution of a patched ML function. The new input is valid if:

            - The new input is a function that has been decorated with
              `exception_safe_function_for_class` or `pickalable_exception_safe_function`
            - OR the new input is a class with the `ExceptionSafeClass` metaclass
            - OR the new input is a list and each of its elements is valid according to the
              these criteria
        """
        if type(inp) == list:
            for item in inp:
                _validate_new_input(item)
        elif isinstance(inp, dict) and "callbacks" in inp:
            _validate_new_input(inp["callbacks"])
        elif callable(inp):
            assert getattr(inp, _ATTRIBUTE_EXCEPTION_SAFE, False), (
                f"New function argument '{inp}' passed to original function is not exception-safe."
                " Please decorate the function with `exception_safe_function` or "
                "`pickalable_exception_safe_function`"
            )
        else:
            assert hasattr(inp, "__class__") and type(inp.__class__) in [
                ExceptionSafeClass,
                ExceptionSafeAbstractClass,
            ], (
                f"Invalid new input '{inp}'. New args / kwargs introduced to `original` function "
                "calls by patched code must either be functions decorated with "
                "`exception_safe_function_for_class`, instances of classes with the "
                "`ExceptionSafeClass` or `ExceptionSafeAbstractClass` metaclass safe or lists of "
                "such exception safe functions / classes."
            )

    def _assert_autologging_input_positional_args_are_superset(
        autologging_call_input, user_call_input
    ):
        length_diff = len(autologging_call_input) - len(user_call_input)
        assert length_diff >= 0, (
            f"{length_diff} expected inputs are missing from the call to the original function."
        )

    def _assert_autologging_input_kwargs_are_superset(autologging_call_input, user_call_input):
        assert set(user_call_input.keys()).issubset(set(autologging_call_input.keys())), (
            "Keyword or dictionary arguments to original function omit"
            " one or more expected keys: '{}'".format(
                set(user_call_input.keys()) - set(autologging_call_input.keys())
            )
        )

    def _validate(autologging_call_input, user_call_input=None):
        """
        Validates that the specified `autologging_call_input` and `user_call_input`
        are compatible. If `user_call_input` is `None`, then `autologging_call_input`
        is regarded as a new input added by autologging and is validated using
        `_validate_new_input`. Otherwise, the following properties must hold:

            - `autologging_call_input` and `user_call_input` must have the same type
              (referred to as "input type")
            - if the input type is a tuple, list or dictionary, then `autologging_call_input` must
              be equivalent to `user_call_input` or be a superset of `user_call_input`
            - for all other input types, `autologging_call_input` and `user_call_input`
              must be equivalent by reference equality or by object equality

        Args:
            autologging_call_input: call input from autologging.
            user_call_input: call input from user.
        """

        if user_call_input is None and autologging_call_input is not None:
            _validate_new_input(autologging_call_input)
            return

        assert type(autologging_call_input) == type(user_call_input), (
            "Type of input to original function '{}' does not match expected type '{}'".format(
                type(autologging_call_input), type(user_call_input)
            )
        )

        if type(autologging_call_input) in [list, tuple]:
            _assert_autologging_input_positional_args_are_superset(
                autologging_call_input, user_call_input
            )
            # If the autologging call input is longer than the user call input, we `zip_longest`
            # will pad the user call input with `None` values to ensure that the subsequent calls
            # to `_validate` identify new inputs added by the autologging call
            for a, u in itertools.zip_longest(autologging_call_input, user_call_input):
                _validate(a, u)
        elif type(autologging_call_input) == dict:
            _assert_autologging_input_kwargs_are_superset(autologging_call_input, user_call_input)
            for key in autologging_call_input.keys():
                _validate(autologging_call_input[key], user_call_input.get(key, None))

        else:
            assert (
                autologging_call_input is user_call_input
                or autologging_call_input == user_call_input
            ), (
                "Input to original function does not match expected input."
                f" Original: '{autologging_call_input}'. Expected: '{user_call_input}'"
            )

    # Similar validation logic found in _validate, unraveling the list of arguments to exclude
    # checks for any validation exempt positional arguments.
    _assert_autologging_input_positional_args_are_superset(autologging_call_args, user_call_args)
    for index, autologging_call_arg, user_call_arg in itertools.zip_longest(
        range(len(user_call_args)), autologging_call_args, user_call_args
    ):
        if not _is_arg_exempt_from_validation(
            autologging_integration,
            function_name,
            user_call_arg,
            argument_index=index,
        ):
            _validate(autologging_call_arg, user_call_arg)

    # Similar validation logic found in _validate, unraveling the dictionary of arguments to exclude
    # checks for any validation exempt keyword arguments.
    _assert_autologging_input_kwargs_are_superset(autologging_call_kwargs, user_call_kwargs)
    for key in autologging_call_kwargs.keys():
        if not _is_arg_exempt_from_validation(
            autologging_integration,
            function_name,
            user_call_kwargs.get(key, None),
            argument_name=key,
        ):
            _validate(
                autologging_call_kwargs[key],
                user_call_kwargs.get(key, None),
            )
