import inspect
import functools


def wrap_decorator(fn, index=None):
    """
    Returns a version of the function wrapped in such a way as to allow both decorator and non-decorator syntax.

    If the first argument of the wrapped function is a callable, the wrapped function is called as-is.

    Otherwise, returns a decorator

    :param fn: Function to decorate.
    :param index: Argument index of the function.  Defaults to 0.
    """
    # Determine the name of the first argument, in case it is specified in kwargs instead.
    signature = inspect.signature(fn)
    arg = None
    param = next(iter(signature.parameters.values()), None)
    assert param
    if param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
        arg = param.name
    assert arg

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        if (args and callable(args[0])) or (arg in kwargs and callable(kwargs[arg])):
            return fn(*args, **kwargs)

        def decorator(_fn):
            return fn(_fn, *args, **kwargs)
        return decorator
    return wrapper