"""Strawberry GraphQL schema extension for PyAuthz authorization.

The PyAuthzExtension resolves the Subject at the start of each GraphQL
request and stores it in the Strawberry context for use by resolvers
and field-level authorization.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

from strawberry.extensions import SchemaExtension

from pyauthz import PyAuthz, Subject

# Type for the subject resolver: takes a Strawberry context dict, returns Subject.
SubjectResolver = Callable[[dict[str, Any]], Subject]


class PyAuthzExtension(SchemaExtension):
    """Strawberry schema extension that integrates PyAuthz authorization.

    Resolves the Subject from the GraphQL context and makes it available
    to resolvers via get_subject(info.context).

    Example:
        ext = PyAuthzExtension(pyauthz=rp, subject_resolver=resolve_user)
        schema = strawberry.Schema(query=Query, extensions=[ext])
    """

    def __init__(
        self,
        *,
        pyauthz: PyAuthz,
        subject_resolver: SubjectResolver,
        execution_context: Any = None,
    ) -> None:
        self.pyauthz = pyauthz
        self.subject_resolver = subject_resolver
        super().__init__(execution_context=execution_context)

    def resolve(self, _next: Callable, root: Any, info: Any, **kwargs: Any) -> Any:
        """Extension hook — resolves subject and attaches to context."""
        context = info.context
        if context is None:
            return _next(root, info, **kwargs)

        # Only resolve once per request
        if "pyauthz" not in context:
            context["pyauthz"] = self.pyauthz
            context["pyauthz_subject"] = self.subject_resolver(context)

        return _next(root, info, **kwargs)


def get_subject(context: dict[str, Any]) -> Subject:
    """Get the current Subject from the Strawberry context.

    Args:
        context: The Strawberry info.context dict.

    Returns:
        The resolved Subject.
    """
    return context["pyauthz_subject"]


def get_pyauthz(context: dict[str, Any]) -> PyAuthz:
    """Get the PyAuthz engine from the Strawberry context.

    Args:
        context: The Strawberry info.context dict.

    Returns:
        The PyAuthz engine instance.
    """
    return context["pyauthz"]
