"""Authorized executor for agent loops.

This module provides a generic executor that can be integrated into any
agent execution loop to provide authorization checks and audit logging.

Example:
    >>> executor = AuthorizedExecutor(
    ...     principal_id="agent-123",
    ...     principal_type="Agent",
    ...     tenant_id="tenant-abc"
    ... )
    >>>
    >>> # Pre-check before tool execution
    >>> if executor.can_execute(tool_name="web_search"):
    ...     result = my_tool.run(query)
    ...     executor.audit("execute", 'Tool::"web_search"', success=True)
    >>>
    >>> # Or use the execute helper
    >>> result = executor.execute(
    ...     action="execute",
    ...     resource='Tool::"web_search"',
    ...     func=lambda: my_tool.run(query)
    ... )
"""

from __future__ import annotations

import logging
from collections.abc import Callable
from typing import Any, TypeVar

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

logger = logging.getLogger("wl_apdp")

T = TypeVar("T")


class AuthorizedExecutor:
    """Generic authorized execution wrapper for any agent loop.

    This class provides a consistent interface for authorization checks
    and audit logging that can be integrated into any agent framework.

    Attributes:
        context: The authorization context for this executor.

    Example:
        >>> executor = AuthorizedExecutor(
        ...     principal_id="agent-123",
        ...     principal_type="Agent"
        ... )
        >>>
        >>> # Check and execute
        >>> if executor.can_execute("web_search"):
        ...     result = search_tool.run(query)
        ...
        >>> # Or use execute() for automatic checks
        >>> result = executor.execute(
        ...     action="execute",
        ...     resource='Tool::"web_search"',
        ...     func=lambda: search_tool.run(query)
        ... )
    """

    def __init__(
        self,
        principal_id: str,
        principal_type: str = "Agent",
        tenant_id: str | None = None,
        session_id: str | None = None,
        attributes: dict[str, Any] | None = None,
        base_url: str = "http://localhost:8081/api",
        token: str | None = None,
    ) -> None:
        """Initialize the executor.

        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.
            base_url: Base URL for wl-apdp API.
            token: Optional bearer token for authentication.
        """
        self.context = AuthorizationContext(
            principal=principal_id,
            principal_type=principal_type,
            tenant_id=tenant_id,
            session_id=session_id,
            attributes=attributes or {},
        )
        self._base_url = base_url
        self._token = token
        self._client: WlApdpSyncClient | None = None

    def _get_client(self) -> WlApdpSyncClient:
        """Get or create the HTTP client."""
        if self._client is None:
            self._client = WlApdpSyncClient(
                base_url=self._base_url,
                token=self._token,
            )
        return self._client

    def close(self) -> None:
        """Close the HTTP client."""
        if self._client is not None:
            self._client.close()
            self._client = None

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

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

    def authorize(
        self,
        action: str,
        resource: str,
        context: dict[str, Any] | None = None,
    ) -> AuthorizationResponse:
        """Check authorization for an action.

        Args:
            action: The action to check (will be wrapped as Action::"...").
            resource: The Cedar resource reference.
            context: Additional context to merge with executor context.

        Returns:
            AuthorizationResponse with the decision.
        """
        client = self._get_client()

        # Merge contexts
        merged_context = self.context.to_context_dict()
        if context:
            merged_context.update(context)

        return client.authorize(
            principal=self.context.to_cedar_principal(),
            action=cedar_action(action),
            resource=resource,
            context=merged_context,
        )

    def is_allowed(
        self,
        action: str,
        resource: str,
        context: dict[str, Any] | None = None,
    ) -> bool:
        """Check if an action is allowed.

        Args:
            action: The action to check.
            resource: The Cedar resource reference.
            context: Additional context.

        Returns:
            True if allowed, False otherwise.
        """
        result = self.authorize(action, resource, context)
        return result.is_allowed

    def can_execute(
        self,
        tool_name: str,
        context: dict[str, Any] | None = None,
    ) -> bool:
        """Check if a tool can be executed.

        Convenience method that builds the resource reference for a tool.

        Args:
            tool_name: Name of the tool to execute.
            context: Additional context.

        Returns:
            True if tool execution is allowed.
        """
        resource = cedar_resource("Tool", tool_name)
        return self.is_allowed("execute", resource, context)

    def can_read(
        self,
        resource_type: str,
        resource_id: str,
        context: dict[str, Any] | None = None,
    ) -> bool:
        """Check if a resource can be read.

        Args:
            resource_type: Type of resource (e.g., "Document", "File").
            resource_id: ID of the resource.
            context: Additional context.

        Returns:
            True if read is allowed.
        """
        resource = cedar_resource(resource_type, resource_id)
        return self.is_allowed("read", resource, context)

    def can_write(
        self,
        resource_type: str,
        resource_id: str,
        context: dict[str, Any] | None = None,
    ) -> bool:
        """Check if a resource can be written.

        Args:
            resource_type: Type of resource.
            resource_id: ID of the resource.
            context: Additional context.

        Returns:
            True if write is allowed.
        """
        resource = cedar_resource(resource_type, resource_id)
        return self.is_allowed("write", resource, context)

    def execute(
        self,
        action: str,
        resource: str,
        func: Callable[[], T],
        context: dict[str, Any] | None = None,
        audit: bool = True,
    ) -> T:
        """Execute a function with authorization check.

        This method checks authorization before executing the function,
        and optionally logs the result for audit purposes.

        Args:
            action: The action being performed.
            resource: The Cedar resource reference.
            func: The function to execute.
            context: Additional context.
            audit: Whether to log the action for audit.

        Returns:
            The result of the function.

        Raises:
            AuthorizationDenied: If the action is not allowed.
        """
        # Check authorization
        result = self.authorize(action, resource, context)

        if not result.is_allowed:
            if audit:
                self.audit(action, resource, success=False, denied=True)
            raise AuthorizationDenied(
                f"Not authorized to {action} on {resource}",
                response=result,
                principal=self.context.to_cedar_principal(),
                action=cedar_action(action),
                resource=resource,
            )

        # Execute with context set
        try:
            with authorization_context(self.context):
                output = func()
            if audit:
                self.audit(action, resource, success=True)
            return output
        except Exception as e:
            if audit:
                self.audit(action, resource, success=False, error=str(e))
            raise

    def audit(
        self,
        action: str,
        resource: str,
        success: bool = True,
        denied: bool = False,
        error: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Log an action for audit purposes.

        Args:
            action: The action that was performed.
            resource: The resource that was accessed.
            success: Whether the action succeeded.
            denied: Whether the action was denied by authorization.
            error: Error message if the action failed.
            metadata: Additional metadata to log.
        """
        log_data = {
            "audit": True,
            "principal": self.context.to_cedar_principal(),
            "principal_type": str(self.context.principal_type),
            "action": action,
            "resource": resource,
            "success": success,
            "denied": denied,
        }

        if self.context.tenant_id:
            log_data["tenant_id"] = self.context.tenant_id
        if self.context.session_id:
            log_data["session_id"] = self.context.session_id
        if error:
            log_data["error"] = error
        if metadata:
            log_data["metadata"] = metadata

        if denied:
            logger.warning(
                f"AUDIT [DENIED]: {self.context.principal} attempted {action} on {resource}",
                extra=log_data,
            )
        elif success:
            logger.info(
                f"AUDIT: {self.context.principal} performed {action} on {resource}",
                extra=log_data,
            )
        else:
            logger.error(
                f"AUDIT [FAILED]: {self.context.principal} failed {action} on {resource}: {error}",
                extra=log_data,
            )

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

        Args:
            **kwargs: Additional attributes to add.

        Returns:
            New AuthorizedExecutor with merged context.
        """
        new_ctx = self.context.with_attributes(**kwargs)
        return AuthorizedExecutor(
            principal_id=new_ctx.principal,
            principal_type=str(new_ctx.principal_type),
            tenant_id=new_ctx.tenant_id,
            session_id=new_ctx.session_id,
            attributes=new_ctx.attributes,
            base_url=self._base_url,
            token=self._token,
        )


class AsyncAuthorizedExecutor:
    """Async version of AuthorizedExecutor.

    Example:
        >>> async with AsyncAuthorizedExecutor(
        ...     principal_id="agent-123",
        ...     principal_type="Agent"
        ... ) as executor:
        ...     result = await executor.execute(
        ...         action="execute",
        ...         resource='Tool::"web_search"',
        ...         func=search_tool.run_async
        ...     )
    """

    def __init__(
        self,
        principal_id: str,
        principal_type: str = "Agent",
        tenant_id: str | None = None,
        session_id: str | None = None,
        attributes: dict[str, Any] | None = None,
        base_url: str = "http://localhost:8081/api",
        token: str | None = None,
    ) -> None:
        """Initialize the async executor."""
        self.context = AuthorizationContext(
            principal=principal_id,
            principal_type=principal_type,
            tenant_id=tenant_id,
            session_id=session_id,
            attributes=attributes or {},
        )
        self._base_url = base_url
        self._token = token
        self._client: WlApdpClient | None = None

    async def __aenter__(self) -> AsyncAuthorizedExecutor:
        """Enter async context manager."""
        self._client = WlApdpClient(
            base_url=self._base_url,
            token=self._token,
        )
        await self._client.__aenter__()
        return self

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

    async def authorize(
        self,
        action: str,
        resource: str,
        context: dict[str, Any] | None = None,
    ) -> AuthorizationResponse:
        """Check authorization for an action."""
        if self._client is None:
            raise RuntimeError("Client not initialized. Use 'async with' context manager.")

        merged_context = self.context.to_context_dict()
        if context:
            merged_context.update(context)

        return await self._client.authorize(
            principal=self.context.to_cedar_principal(),
            action=cedar_action(action),
            resource=resource,
            context=merged_context,
        )

    async def is_allowed(
        self,
        action: str,
        resource: str,
        context: dict[str, Any] | None = None,
    ) -> bool:
        """Check if an action is allowed."""
        result = await self.authorize(action, resource, context)
        return result.is_allowed

    async def can_execute(
        self,
        tool_name: str,
        context: dict[str, Any] | None = None,
    ) -> bool:
        """Check if a tool can be executed."""
        resource = cedar_resource("Tool", tool_name)
        return await self.is_allowed("execute", resource, context)

    async def execute(
        self,
        action: str,
        resource: str,
        func: Callable[[], T],
        context: dict[str, Any] | None = None,
        audit: bool = True,
    ) -> T:
        """Execute an async function with authorization check."""
        result = await self.authorize(action, resource, context)

        if not result.is_allowed:
            if audit:
                self.audit(action, resource, success=False, denied=True)
            raise AuthorizationDenied(
                f"Not authorized to {action} on {resource}",
                response=result,
                principal=self.context.to_cedar_principal(),
                action=cedar_action(action),
                resource=resource,
            )

        try:
            async with authorization_context(self.context):
                output = await func()  # type: ignore
            if audit:
                self.audit(action, resource, success=True)
            return output
        except Exception as e:
            if audit:
                self.audit(action, resource, success=False, error=str(e))
            raise

    def audit(
        self,
        action: str,
        resource: str,
        success: bool = True,
        denied: bool = False,
        error: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Log an action for audit purposes."""
        log_data = {
            "audit": True,
            "principal": self.context.to_cedar_principal(),
            "action": action,
            "resource": resource,
            "success": success,
            "denied": denied,
        }
        if error:
            log_data["error"] = error
        if metadata:
            log_data["metadata"] = metadata

        level = logging.WARNING if denied else (logging.INFO if success else logging.ERROR)
        logger.log(level, f"AUDIT: {self.context.principal} {action} {resource}", extra=log_data)
