"""LangGraph auth context propagation for wl-apdp SDK.

Provides a context manager that propagates authorization metadata
via Python's contextvars, making it available to any @authorized_tool
or @create_authorized_node decorators within a LangGraph execution.
"""

from __future__ import annotations

import contextvars
from typing import Any

# Context variable for propagating auth context through LangGraph execution
_auth_context_var: contextvars.ContextVar[LangGraphAuthContext | None] = (
    contextvars.ContextVar("wl_apdp_auth_context", default=None)
)


def get_current_auth_context() -> LangGraphAuthContext | None:
    """Get the current auth context from contextvars.

    Returns:
        The active LangGraphAuthContext, or None if not in an auth scope.
    """
    return _auth_context_var.get()


class LangGraphAuthContext:
    """Context manager that propagates auth metadata via contextvars.

    Use this to wrap graph.invoke() calls so that authorization info
    is available to all nodes and tools within the graph execution.

    Attributes:
        principal_id: The agent slug.
        principal_type: Type of principal (e.g., "Agent").
        tenant_id: Tenant ID for isolation.
        session_id: Request/job ID for correlation.
        attributes: Additional auth attributes (trust_level, etc.).

    Usage:
        with LangGraphAuthContext(
            principal_id="chat-assistant-v1",
            principal_type="Agent",
            tenant_id=family_id,
            session_id=request_id,
            attributes={"trust_level": 5},
        ):
            graph.invoke(state)
    """

    def __init__(
        self,
        principal_id: str = "",
        principal_type: str = "Agent",
        tenant_id: str = "",
        session_id: str = "",
        attributes: dict[str, Any] | None = None,
    ) -> None:
        self.principal_id = principal_id
        self.principal_type = principal_type
        self.tenant_id = tenant_id
        self.session_id = session_id
        self.attributes = attributes or {}
        self._token: contextvars.Token[LangGraphAuthContext | None] | None = None

    def __enter__(self) -> LangGraphAuthContext:
        self._token = _auth_context_var.set(self)
        return self

    def __exit__(self, *args: Any) -> None:
        if self._token is not None:
            _auth_context_var.reset(self._token)
            self._token = None
