from flask_jwt_extended import JWTManager
from iti.applications.common.utils import fail

jwt = JWTManager()


def init_jwt(app) -> None:
    """
    初始化 JWT
    """
    jwt.init_app(app)

    # 自定义错误消息
    @jwt.unauthorized_loader
    def unauthorized_loader(_callback):
        return fail("缺少令牌参数 Authorization Bearer", code=401), 401

    @jwt.invalid_token_loader
    def invalid_token_loader(_callback):
        return fail("无效的令牌", code=401, data=str(_callback)), 401

    @jwt.expired_token_loader
    def expired_token_loader(_header, _payload):
        return fail("令牌已过期", code=401), 401

    @jwt.user_identity_loader
    def user_identity_loader(user):
        if user is None or not hasattr(user, "id"):
            return None
        return user.id