"""AutoGen integration for wl-apdp.

This module provides adapters for integrating wl-apdp authorization
into Microsoft AutoGen agents and functions.

Install with: pip install wl-apdp[autogen]

Example:
    >>> from wl_apdp.contrib.autogen import AuthorizedFunction
    >>>
    >>> def search(query: str) -> str:
    ...     return f"Results for: {query}"
    >>>
    >>> authorized_search = AuthorizedFunction(
    ...     func=search,
    ...     resource='Tool::"search"'
    ... )
"""

from __future__ import annotations

import functools
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, TypeVar

from ..client import WlApdpSyncClient
from ..context import AuthorizationContext, authorization_context, get_current_context
from ..exceptions import AuthorizationDenied
from ..models import cedar_action, cedar_resource

if TYPE_CHECKING:
    pass

T = TypeVar("T")


class AuthorizedFunction:
    """Wrapper for functions with authorization checks.

    This class wraps any callable function to add authorization checks
    before execution. It's compatible with AutoGen's function registration.

    Example:
        >>> def calculate(expression: str) -> float:
        ...     return eval(expression)
        >>>
        >>> authorized_calc = AuthorizedFunction(
        ...     func=calculate,
        ...     resource='Tool::"calculator"'
        ... )
        >>>
        >>> # Register with AutoGen
        >>> assistant.register_for_llm()(authorized_calc)
    """

    def __init__(
        self,
        func: Callable[..., T],
        resource: str | None = None,
        action: str = "execute",
        name: str | None = None,
        description: str | None = None,
        base_url: str = "http://localhost:8081/api",
        token: str | None = None,
        fail_closed: bool = True,
    ) -> None:
        """Initialize the authorized function.

        Args:
            func: The function to wrap.
            resource: Cedar resource reference. If None, uses Tool::"<func_name>".
            action: The action to check (default: "execute").
            name: Override name for the function.
            description: Override description for the function.
            base_url: Base URL for wl-apdp API.
            token: Optional bearer token.
            fail_closed: If True, deny access when no context is available.
        """
        self.func = func
        self._name = name or func.__name__
        self._description = description or func.__doc__ or ""
        self.resource = resource or cedar_resource("Tool", self._name)
        self.action = action
        self._base_url = base_url
        self._token = token
        self.fail_closed = fail_closed

        # Copy function metadata
        functools.update_wrapper(self, func)

    @property
    def name(self) -> str:
        """Return the function name."""
        return self._name

    @property
    def description(self) -> str:
        """Return the function description."""
        return self._description

    def __call__(self, *args: Any, **kwargs: Any) -> T:
        """Execute the function with authorization check."""
        ctx = get_current_context()

        if ctx is None:
            if self.fail_closed:
                raise AuthorizationDenied(
                    "No authorization context available",
                    resource=self.resource,
                    action=self.action,
                )
            return self.func(*args, **kwargs)

        with WlApdpSyncClient(base_url=self._base_url, token=self._token) as client:
            result = client.authorize(
                principal=ctx.to_cedar_principal(),
                action=cedar_action(self.action),
                resource=self.resource,
                context=ctx.to_context_dict(),
            )

        if not result.is_allowed:
            raise AuthorizationDenied(
                f"Not authorized to execute {self._name}",
                response=result,
                principal=ctx.to_cedar_principal(),
                action=cedar_action(self.action),
                resource=self.resource,
            )

        return self.func(*args, **kwargs)


def authorized_function(
    resource: str | None = None,
    action: str = "execute",
    base_url: str = "http://localhost:8081/api",
    token: str | None = None,
    fail_closed: bool = True,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
    """Decorator to add authorization to any function.

    This decorator wraps a function with authorization checks compatible
    with AutoGen's function calling.

    Args:
        resource: Cedar resource reference. If None, derived from function name.
        action: The action to check (default: "execute").
        base_url: Base URL for wl-apdp API.
        token: Optional bearer token.
        fail_closed: If True, deny access when no context is available.

    Returns:
        Decorated function with authorization checks.

    Example:
        >>> @authorized_function(resource='Tool::"weather"')
        ... def get_weather(city: str) -> str:
        ...     '''Get weather for a city.'''
        ...     return f"Weather in {city}: Sunny"
    """

    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        wrapper = AuthorizedFunction(
            func=func,
            resource=resource,
            action=action,
            base_url=base_url,
            token=token,
            fail_closed=fail_closed,
        )
        return wrapper  # type: ignore

    return decorator


def wrap_autogen_functions(
    functions: list[Callable[..., Any]],
    action: str = "execute",
    base_url: str = "http://localhost:8081/api",
    token: str | None = None,
) -> list[AuthorizedFunction]:
    """Wrap a list of functions with authorization checks.

    Args:
        functions: List of functions to wrap.
        action: The action to check for function execution.
        base_url: Base URL for wl-apdp API.
        token: Optional bearer token.

    Returns:
        List of wrapped functions with authorization checks.

    Example:
        >>> functions = [search, calculate, get_weather]
        >>> authorized_funcs = wrap_autogen_functions(functions)
    """
    return [
        AuthorizedFunction(
            func=func,
            action=action,
            base_url=base_url,
            token=token,
        )
        for func in functions
    ]


class AutoGenAuthContext:
    """Context manager for running AutoGen conversations with authorization.

    This class sets up authorization context for AutoGen agent conversations.

    Example:
        >>> with AutoGenAuthContext(
        ...     principal_id="user-123",
        ...     principal_type="User",
        ...     tenant_id="org-abc"
        ... ) as ctx:
        ...     user_proxy.initiate_chat(assistant, message="Hello")
    """

    def __init__(
        self,
        principal_id: str,
        principal_type: str = "User",
        tenant_id: str | None = None,
        session_id: str | None = None,
        attributes: dict[str, Any] | None = None,
    ) -> None:
        """Initialize the auth context.

        Args:
            principal_id: Unique identifier for the principal.
            principal_type: Type of principal (User, Agent, Service, Bot).
            tenant_id: Optional tenant ID for multi-tenant isolation.
            session_id: Optional session ID for tracking.
            attributes: Additional attributes for policy evaluation.
        """
        self.context = AuthorizationContext(
            principal=principal_id,
            principal_type=principal_type,
            tenant_id=tenant_id,
            session_id=session_id,
            attributes=attributes or {},
        )
        self._ctx_manager = authorization_context(self.context)

    def __enter__(self) -> AuthorizationContext:
        """Enter context manager."""
        return self._ctx_manager.__enter__()

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

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

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


def create_authorized_agent_config(
    principal_id: str,
    principal_type: str = "Agent",
    tenant_id: str | None = None,
    allowed_tools: list[str] | None = None,
    base_url: str = "http://localhost:8081/api",
    token: str | None = None,
) -> dict[str, Any]:
    """Create an agent configuration with authorization metadata.

    This function creates a configuration dict that can be used when
    creating AutoGen agents with authorization information embedded.

    Args:
        principal_id: Unique identifier for this agent.
        principal_type: Type of principal (default: "Agent").
        tenant_id: Optional tenant ID.
        allowed_tools: List of tool names this agent is allowed to use.
        base_url: Base URL for wl-apdp API.
        token: Optional bearer token.

    Returns:
        Configuration dict with authorization metadata.

    Example:
        >>> config = create_authorized_agent_config(
        ...     principal_id="assistant-1",
        ...     tenant_id="org-abc",
        ...     allowed_tools=["search", "calculate"]
        ... )
        >>> assistant = AssistantAgent("assistant", **config)
    """
    return {
        "_wl_apdp": {
            "principal_id": principal_id,
            "principal_type": principal_type,
            "tenant_id": tenant_id,
            "allowed_tools": allowed_tools or [],
            "base_url": base_url,
            "token": token,
        }
    }
