"""Context propagation utilities for wl-apdp.

This module provides context management for authorization throughout async/sync
execution. It uses Python's contextvars for proper propagation across async
boundaries.

Example:
    >>> from wl_apdp import AuthorizationContext, authorization_context
    >>>
    >>> ctx = AuthorizationContext(
    ...     principal="agent-123",
    ...     principal_type="Agent",
    ...     tenant_id="tenant-abc"
    ... )
    >>>
    >>> with authorization_context(ctx):
    ...     # All code here has access to the context
    ...     current = get_current_context()
    ...     print(current.to_cedar_principal())  # 'Agent::"agent-123"'
"""

from __future__ import annotations

from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from typing import Any

from .models import PrincipalType

# Context variable for async propagation
_authorization_context: ContextVar[AuthorizationContext | None] = ContextVar(
    "authorization_context", default=None
)


@dataclass
class AuthorizationContext:
    """Context for authorization throughout execution.

    This dataclass holds all information about the current principal and their
    context. It can be used with the `authorization_context` context manager
    to propagate through async execution.

    Attributes:
        principal: The unique identifier for the principal.
        principal_type: The type of principal (User, Agent, Service, Bot).
        tenant_id: Optional tenant identifier for multi-tenant isolation.
        session_id: Optional session identifier for tracking.
        attributes: Additional attributes for policy evaluation.

    Example:
        >>> ctx = AuthorizationContext(
        ...     principal="agent-123",
        ...     principal_type="Agent",
        ...     tenant_id="tenant-abc",
        ...     attributes={"trust_level": 5}
        ... )
        >>> ctx.to_cedar_principal()
        'Agent::"agent-123"'
    """

    principal: str
    principal_type: str | PrincipalType = "Agent"
    tenant_id: str | None = None
    session_id: str | None = None
    attributes: dict[str, Any] = field(default_factory=dict)

    def to_cedar_principal(self) -> str:
        """Format as Cedar entity reference.

        Returns:
            Cedar principal reference (e.g., 'Agent::"agent-123"').
        """
        type_str = (
            self.principal_type.value
            if isinstance(self.principal_type, PrincipalType)
            else self.principal_type
        )
        return f'{type_str}::"{self.principal}"'

    def to_context_dict(self) -> dict[str, Any]:
        """Build context dict for authorization requests.

        This combines tenant_id, session_id, and any additional attributes
        into a single dict suitable for passing to the authorization API.

        Returns:
            Dictionary of context values.
        """
        ctx: dict[str, Any] = {**self.attributes}
        if self.tenant_id:
            ctx["tenant_id"] = self.tenant_id
        if self.session_id:
            ctx["session_id"] = self.session_id
        return ctx

    def with_attributes(self, **kwargs: Any) -> AuthorizationContext:
        """Create a new context with additional attributes.

        This is useful for adding request-specific context without modifying
        the original context.

        Args:
            **kwargs: Additional attributes to add.

        Returns:
            New AuthorizationContext with merged attributes.

        Example:
            >>> ctx = AuthorizationContext(principal="agent-1", principal_type="Agent")
            >>> new_ctx = ctx.with_attributes(request_id="req-123")
        """
        new_attrs = {**self.attributes, **kwargs}
        return AuthorizationContext(
            principal=self.principal,
            principal_type=self.principal_type,
            tenant_id=self.tenant_id,
            session_id=self.session_id,
            attributes=new_attrs,
        )


def get_current_context() -> AuthorizationContext | None:
    """Get the current authorization context.

    Returns:
        The current AuthorizationContext, or None if not set.
    """
    return _authorization_context.get()


def set_current_context(ctx: AuthorizationContext | None) -> Token[AuthorizationContext | None]:
    """Set the current authorization context.

    Args:
        ctx: The context to set, or None to clear.

    Returns:
        Token that can be used to reset the context.
    """
    return _authorization_context.set(ctx)


def require_context() -> AuthorizationContext:
    """Get the current context, raising if not set.

    Returns:
        The current AuthorizationContext.

    Raises:
        RuntimeError: If no context is set.
    """
    ctx = get_current_context()
    if ctx is None:
        raise RuntimeError(
            "No authorization context available. "
            "Use 'authorization_context' context manager to set one."
        )
    return ctx


class authorization_context:
    """Context manager for setting authorization context.

    This context manager sets the authorization context for the duration of
    the block. It supports both sync and async usage.

    Example (sync):
        >>> ctx = AuthorizationContext(principal="agent-1", principal_type="Agent")
        >>> with authorization_context(ctx):
        ...     # Context is available here
        ...     pass

    Example (async):
        >>> async with authorization_context(ctx):
        ...     # Context is available here
        ...     pass
    """

    def __init__(self, ctx: AuthorizationContext) -> None:
        """Initialize the context manager.

        Args:
            ctx: The authorization context to set.
        """
        self.ctx = ctx
        self.token: Token[AuthorizationContext | None] | None = None

    def __enter__(self) -> AuthorizationContext:
        """Enter the context manager (sync)."""
        self.token = _authorization_context.set(self.ctx)
        return self.ctx

    def __exit__(self, *args: Any) -> None:
        """Exit the context manager (sync)."""
        if self.token is not None:
            _authorization_context.reset(self.token)

    async def __aenter__(self) -> AuthorizationContext:
        """Enter the context manager (async)."""
        return self.__enter__()

    async def __aexit__(self, *args: Any) -> None:
        """Exit the context manager (async)."""
        self.__exit__(*args)
