from typing import Type

from afnio.autodiff.function import Function


def evaluator(cls: Type[Function]) -> Type[Function]:
    """
    Decorator to mark a class as an evaluator within the `afnio` framework.

    This decorator is intended to be applied to classes inheriting from
    [`Function`][afnio.autodiff.function.Function]. It sets an internal attribute
    `_is_evaluator` on the class to `True`, allowing the autodiff engine to recognize
    the class as an evaluator. This designation enables the evaluator's
    [`backward`][afnio.autodiff.function.Function.backward] function
    to be called without any input arguments when computing gradients.

    Evaluators are responsible for assessing predictions using either deterministic
    or language model-based approaches. By using this decorator, users can build
    their own evaluator classes without worrying about internal setup details.

    Args:
        cls: The class being decorated,
            which should inherit from [`Function`][afnio.autodiff.function.Function].

    Returns:
        The same class with the `_is_evaluator` attribute set to `True`.

    Examples:
        >>> @evaluator
        >>> class MyCustomEvaluator(Function):
        >>>     @staticmethod
        >>>     def forward(ctx, prediction, target):
        >>>         # Evaluation logic
        >>>         pass
        >>>     @staticmethod
        >>>     def backward(ctx, grad_output):
        >>>         # Backpropagation logic
        >>>         pass
    """
    if not isinstance(cls, type):
        raise TypeError(
            "`evaluator` decorator must be applied to a class, not an instance."
        )
    if not issubclass(cls, Function):
        raise TypeError(
            "`evaluator` decorator must be applied to a subclass of Function."
        )

    cls._is_evaluator = True
    return cls
