from functools import wraps
from flask import g, jsonify
from jwtauth import JWTAuth

jwt_auth = JWTAuth()

AUTHORIZED_IDENTITIES = {"internal-sync-service", "migration-worker"}


def only_for_identity(*allowed_identities):
    """Restrict an endpoint to specific JWT identity IDs."""
    def decorator(f):
        @wraps(f)
        def wrapper(*args, **kwargs):
            identity_id = getattr(g.request_context, "jwt_identity_id", None)
            if not identity_id or identity_id not in allowed_identities:
                return jsonify({"error": "forbidden"}), 403
            return f(*args, **kwargs)
        return wrapper
    return decorator