import jwt
from functools import wraps
from flask import request, jsonify, current_app, g


def token_required(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        token = None

        # Get token from Authorization header
        auth_header = request.headers.get('Authorization')
        if auth_header and auth_header.startswith("Bearer "):
            token = auth_header.split(" ")[1]

        if not token:
            return jsonify(success=False, error="Token is missing", code="TOKEN_MISSING"), 401

        try:
            # Decode
            query_results = current_app.extensions['query_results']

            decoded = jwt.decode(token, current_app.config['JWT_SECRET_KEY'], algorithms=["HS256"])
            user_id = decoded.get("userid")

            if not user_id:
                raise jwt.InvalidTokenError("Invalid subject in token.")

            # Fetch user from MySQL
            statement = "SELECT * FROM users WHERE {} = %s"
            params = (user_id,)
            result = query_results.fetch_query_results(statement, params=params,column_name=("id",))

            if not result:
                return jsonify(success=False, error="User not found", code="USER_NOT_FOUND"), 404

            g.current_user = result[0]  # Assuming single user

        except jwt.ExpiredSignatureError:
            return jsonify(success=False, error="Token expired", code="TOKEN_EXPIRED"), 401
        except jwt.InvalidTokenError as e:
            return jsonify(success=False, error=str(e), code="INVALID_TOKEN"), 401
        except Exception as e:
            return jsonify(success=False, error="Server error: " + str(e), code="SERVER_ERROR"), 500

        return fn(*args, **kwargs)

    return wrapper