from functools import wraps
from typing import Callable

from flask import jsonify, Response

from atheneum.api.model import APIResponse


def return_json(func: Callable) -> Callable:
    """
    If an Response object is not returned, jsonify the result and return it
    :param func:
    :return:
    """

    @wraps(func)
    def decorate(*args, **kwargs):
        result = func(*args, **kwargs)
        if isinstance(result, Response):
            return result
        if isinstance(result, APIResponse):
            return jsonify(result.payload), result.status
        return jsonify(result)

    return decorate