from typing import Any, Callable

from ariadne import SchemaDirectiveVisitor
from graphql import GraphQLError


def _wrap_resolver(field: Any, check: Callable[[Any], None]) -> None:
    original = field.resolve

    def resolve(source: Any, info: Any, **kwargs: Any) -> Any:
        check(info.context)
        return original(source, info, **kwargs)

    field.resolve = resolve


# при применении к полю оборачивает оригинальный резолвер:
# если в контексте нет `user`, бросает исключение с сообщением «Authentication required»;
class IsAuthenticatedDirective(SchemaDirectiveVisitor):
    def visit_field_definition(self, field: Any, object_type: Any) -> Any:
        def check(context: Any) -> None:
            if not context or not context.get("user"):
                raise GraphQLError("Authentication required")

        _wrap_resolver(field, check)
        return field


# аналогично проверяет, что `user.is_admin == True`, иначе «Admin only».
class IsAdminDirective(SchemaDirectiveVisitor):
    def visit_field_definition(self, field: Any, object_type: Any) -> Any:
        def check(context: Any) -> None:
            user = context.get("user") if context else None
            if user is None:
                raise GraphQLError("Authentication required")
            if not user.is_admin:
                raise GraphQLError("Admin only")

        _wrap_resolver(field, check)
        return field