"""CrewAI integration for wl-apdp.

This module provides adapters for integrating wl-apdp authorization
into CrewAI agents and tools.

Install with: pip install wl-apdp[crewai]

Example:
    >>> from wl_apdp.contrib.crewai import AuthorizedCrewAITool
    >>> from crewai_tools import SerperDevTool
    >>>
    >>> tool = AuthorizedCrewAITool(
    ...     wrapped_tool=SerperDevTool(),
    ...     resource='Tool::"web_search"'
    ... )
"""

from __future__ import annotations

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

from crewai.tools import BaseTool
from pydantic import Field

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:
    from crewai import Agent, Task


class AuthorizedCrewAITool(BaseTool):
    """Wrapper that adds authorization checks to any CrewAI tool.

    This class wraps a CrewAI BaseTool and checks authorization before
    each execution.

    Example:
        >>> from crewai_tools import SerperDevTool
        >>> from wl_apdp.contrib.crewai import AuthorizedCrewAITool
        >>>
        >>> search_tool = AuthorizedCrewAITool(
        ...     wrapped_tool=SerperDevTool(),
        ...     resource='Tool::"web_search"',
        ...     action="execute"
        ... )
    """

    # Pydantic model fields
    name: str = Field(default="", description="Tool name")
    description: str = Field(default="", description="Tool description")
    wrapped_tool: Any = Field(default=None, description="The wrapped BaseTool")
    resource: str = Field(default="", description="Cedar resource reference")
    auth_action: str = Field(default="execute", description="The action to check")
    base_url: str = Field(default="http://localhost:8081", description="Base URL for wl-apdp API")
    token: str | None = Field(default=None, description="Optional bearer token")
    fail_closed: bool = Field(default=True, description="Deny access when no context is available")

    def __init__(
        self,
        wrapped_tool: BaseTool,
        resource: str | None = None,
        action: str = "execute",
        base_url: str = "http://localhost:8081",
        token: str | None = None,
        fail_closed: bool = True,
        **kwargs: Any,
    ) -> None:
        """Initialize the authorized tool wrapper.

        Args:
            wrapped_tool: The CrewAI tool to wrap.
            resource: Cedar resource reference. If None, uses Tool::"<tool_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.
        """
        # Calculate resource if not provided
        actual_resource = resource or cedar_resource("Tool", wrapped_tool.name)

        # Call parent constructor with required BaseTool fields
        super().__init__(
            name=wrapped_tool.name,
            description=wrapped_tool.description,
            wrapped_tool=wrapped_tool,
            resource=actual_resource,
            auth_action=action,
            base_url=base_url,
            token=token,
            fail_closed=fail_closed,
            **kwargs,
        )

    def _run(self, *args: Any, **kwargs: Any) -> Any:
        """Execute the tool 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.auth_action,
                )
            return self.wrapped_tool._run(*args, **kwargs)

        # Check authorization
        with WlApdpSyncClient(base_url=self.base_url, token=self.token) as client:
            result = client.authorize(
                principal=ctx.to_cedar_principal(),
                action=cedar_action(self.auth_action),
                resource=self.resource,
                context=ctx.to_context_dict(),
            )

        if not result.is_allowed:
            raise AuthorizationDenied(
                f"Agent {ctx.principal} not authorized to use {self.name}",
                response=result,
                principal=ctx.to_cedar_principal(),
                action=cedar_action(self.auth_action),
                resource=self.resource,
            )

        return self.wrapped_tool._run(*args, **kwargs)

    def run(self, *args: Any, **kwargs: Any) -> Any:
        """Public run method that delegates to _run."""
        return self._run(*args, **kwargs)


def create_authorization_guardrail(
    allowed_actions: list[str],
    resource_pattern: str,
    base_url: str = "http://localhost:8081/api",
    token: str | None = None,
) -> Callable[[str], tuple[bool, Any]]:
    """Create a CrewAI task guardrail that checks authorization.

    This function creates a guardrail that can be attached to CrewAI tasks
    to enforce authorization before task completion.

    Args:
        allowed_actions: List of actions that must all be allowed.
        resource_pattern: Cedar resource reference pattern.
        base_url: Base URL for wl-apdp API.
        token: Optional bearer token.

    Returns:
        Guardrail function compatible with CrewAI tasks.

    Example:
        >>> guardrail = create_authorization_guardrail(
        ...     allowed_actions=["read", "summarize"],
        ...     resource_pattern='Document::"*"'
        ... )
        >>> task = Task(
        ...     description="Summarize the document",
        ...     guardrail=guardrail
        ... )
    """

    def guardrail(output: str) -> tuple[bool, Any]:
        ctx = get_current_context()
        if ctx is None:
            return (False, "No authorization context available")

        with WlApdpSyncClient(base_url=base_url, token=token) as client:
            for action in allowed_actions:
                result = client.authorize(
                    principal=ctx.to_cedar_principal(),
                    action=cedar_action(action),
                    resource=resource_pattern,
                    context=ctx.to_context_dict(),
                )
                if not result.is_allowed:
                    return (False, f"Not authorized for action: {action}")

        return (True, output)

    return guardrail


def wrap_agent_tools(
    agent: Agent,
    action: str = "execute",
    base_url: str = "http://localhost:8081/api",
    token: str | None = None,
) -> Agent:
    """Wrap all tools of an agent with authorization checks.

    This function modifies an agent in-place to wrap all its tools
    with authorization checks.

    Args:
        agent: The CrewAI agent to modify.
        action: The action to check for tool execution.
        base_url: Base URL for wl-apdp API.
        token: Optional bearer token.

    Returns:
        The same agent with wrapped tools.

    Example:
        >>> researcher = Agent(
        ...     role="Researcher",
        ...     tools=[SerperDevTool(), WebsiteSearchTool()]
        ... )
        >>> researcher = wrap_agent_tools(researcher)
    """
    if agent.tools:
        agent.tools = [
            AuthorizedCrewAITool(
                wrapped_tool=tool,
                action=action,
                base_url=base_url,
                token=token,
            )
            for tool in agent.tools
        ]
    return agent


class AuthorizedCrew:
    """Wrapper for CrewAI Crew that sets up authorization context.

    This class wraps a CrewAI Crew and ensures that all execution happens
    within an authorization context.

    Example:
        >>> from crewai import Agent, Task, Crew
        >>> from wl_apdp.contrib.crewai import AuthorizedCrew
        >>>
        >>> crew = AuthorizedCrew(
        ...     agents=[researcher, writer],
        ...     tasks=[research_task, write_task],
        ...     principal_id="crew-instance-123",
        ...     tenant_id="tenant-abc"
        ... )
        >>> result = crew.kickoff()
    """

    def __init__(
        self,
        agents: list[Agent],
        tasks: list[Task],
        principal_id: str,
        principal_type: str = "Agent",
        tenant_id: str | None = None,
        attributes: dict[str, Any] | None = None,
        wrap_tools: bool = True,
        base_url: str = "http://localhost:8081/api",
        token: str | None = None,
        **crew_kwargs: Any,
    ) -> None:
        """Initialize the authorized crew.

        Args:
            agents: List of CrewAI agents.
            tasks: List of CrewAI tasks.
            principal_id: Unique identifier for this crew instance.
            principal_type: Type of principal (default: "Agent").
            tenant_id: Optional tenant ID for multi-tenant isolation.
            attributes: Additional attributes for authorization context.
            wrap_tools: If True, automatically wrap all agent tools.
            base_url: Base URL for wl-apdp API.
            token: Optional bearer token.
            **crew_kwargs: Additional arguments passed to Crew constructor.
        """
        from crewai import Crew

        # Optionally wrap all agent tools
        if wrap_tools:
            for agent in agents:
                wrap_agent_tools(agent, base_url=base_url, token=token)

        self.crew = Crew(agents=agents, tasks=tasks, **crew_kwargs)
        self.auth_context = AuthorizationContext(
            principal=principal_id,
            principal_type=principal_type,
            tenant_id=tenant_id,
            attributes=attributes or {},
        )

    def kickoff(self, inputs: dict[str, Any] | None = None) -> Any:
        """Run the crew with authorization context.

        Args:
            inputs: Optional inputs to pass to the crew.

        Returns:
            The crew execution result.
        """
        with authorization_context(self.auth_context):
            return self.crew.kickoff(inputs=inputs)

    async def kickoff_async(self, inputs: dict[str, Any] | None = None) -> Any:
        """Run the crew asynchronously with authorization context.

        Args:
            inputs: Optional inputs to pass to the crew.

        Returns:
            The crew execution result.
        """
        async with authorization_context(self.auth_context):
            return await self.crew.kickoff_async(inputs=inputs)
