from functools import wraps
from flask_jwt_extended import get_jwt_identity
from application_tpl.models import User


def roles_required(*roles):
    def decorator(func):
        @wraps(func)
        def wrapped(*args, **kwargs):
            current_user = get_jwt_identity()
            user: User = User.get_by_email(current_user)
            try:
                if user.has_roles(*roles):
                    return func(*args, **kwargs)
            except AttributeError:
                return (
                    {"code": 400, "message": "Missing access token."},
                    400,
                )
            return {"code": 403, "status": "Forbidden"}, 403

        return wrapped

    return decorator


def roles_accepted(*roles):
    def decorator(func):
        @wraps(func)
        def wrapped(*args, **kwargs):
            current_user = get_jwt_identity()
            user: User = User.get_by_email(current_user)
            try:
                if user.has_any_roles(*roles):
                    return func(*args, **kwargs)
            except AttributeError:
                return (
                    {"code": 400, "message": "Missing access token."},
                    400,
                )
            return {"code": 403, "status": "Forbidden"}, 403

        return wrapped

    return decorator
