import jwt
from flask import request, g, abort
from . import db, app
from .model import User
from functools import wraps

def get_jwt(self):
    """
    Extension method for the User class to
    compute the user's JSON Web Token
    """
    return jwt.encode({'id': self.id}, app.secret_key, algorithm='HS256')

User.get_jwt = get_jwt

def jwt_required(f):
    """
    Decorator for routes in the views module,
    returning a bad request error if the "Authorization"
    header is not set with a valid authentication token.
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'Authorization' not in request.headers:
            abort(400)
        authorization = request.headers['Authorization'].split(' ')
        if len(authorization) < 2:
            abort(400)

        try:
            decoded = jwt.decode(
                    authorization[1].encode(),
                    app.secret_key,
                    algorithms = ['HS256'])
            g.user = User.query.filter_by(id=decoded['id']).first()
            if g.user is None: abort(400)
            return f(*args, **kwargs)
        except Exception as e:
            print(e)
            abort(500)
    return decorated_function