"""ARIAToolInterceptor — ToolCallInterceptor backed by ARIA's AuthZEN PDP.

Integrates with AGT's ToolCallInterceptor protocol to evaluate every
tool call against the ARIA PDP before execution. Posts the AuthZEN
request built by our mapping module directly to the PDP, ensuring
the PDP sees the exact request shape we define.
"""

from __future__ import annotations

import logging
import time
from typing import Any

import httpx
from agent_os.integrations.base import ToolCallRequest, ToolCallResult
from authzen_backend.auth import ClientCredentialsAuth, StaticTokenAuth

from aria_agentkit import compat as _compat  # noqa: F401 — version gate
from aria_agentkit.policy.mapping import map_tool_call

logger = logging.getLogger(__name__)


class ARIAToolInterceptor:
    """ToolCallInterceptor that evaluates tool calls via ARIA's AuthZEN PDP.

    Satisfies the ``agent_os.integrations.base.ToolCallInterceptor`` protocol.
    Posts the AuthZEN request directly to the PDP — our mapping module is the
    single source of truth for the request shape.

    Args:
        pdp_url: AuthZEN PDP evaluation endpoint URL.
        pdp_application: ARIA PDP application identifier.
        token: Static bearer token for PDP authentication.
        client_credentials: ``ClientCredentialsAuth`` instance for
            OAuth 2.0 client-credentials auth.
        http_client: Optional pre-configured httpx.Client.
        timeout: HTTP timeout in seconds.
        fail_closed: If True (default), PDP errors produce denial.
            If False, exceptions propagate to the caller.
    """

    def __init__(
        self,
        pdp_url: str,
        pdp_application: str,
        *,
        token: str | None = None,
        client_credentials: ClientCredentialsAuth | None = None,
        http_client: httpx.Client | None = None,
        timeout: float = 10.0,
        fail_closed: bool = True,
    ) -> None:
        self._pdp_url = pdp_url
        self._pdp_application = pdp_application
        self._fail_closed = fail_closed

        if token is not None and client_credentials is not None:
            raise ValueError("Provide either 'token' or 'client_credentials', not both.")
        self._auth: ClientCredentialsAuth | StaticTokenAuth | None = None
        if client_credentials is not None:
            self._auth = client_credentials
        elif token is not None:
            self._auth = StaticTokenAuth(token=token)

        self._owns_client = http_client is None
        self._client = http_client or httpx.Client(timeout=timeout)

    def intercept(self, request: ToolCallRequest) -> ToolCallResult:
        """Evaluate a tool call against the ARIA PDP.

        Returns a ToolCallResult with the PDP decision. On error, returns
        denial if fail_closed is True (default), otherwise re-raises.
        """
        start = time.monotonic()
        try:
            mapping = map_tool_call(
                tool_name=request.tool_name,
                arguments=request.arguments,
                call_id=request.call_id,
                agent_id=request.agent_id,
                metadata=request.metadata,
                pdp_application=self._pdp_application,
            )

            headers: dict[str, str] = {"Content-Type": "application/json"}
            if self._auth is not None:
                headers["Authorization"] = f"Bearer {self._auth.get_token(self._client)}"

            resp = self._client.post(
                self._pdp_url,
                json=mapping.authzen_request.model_dump(),
                headers=headers,
            )
            resp.raise_for_status()
            body = resp.json()
            elapsed_ms = (time.monotonic() - start) * 1000

            allowed = bool(body.get("decision", False))

            modified_args = None
            reason: str | None = None
            pdp_context = body.get("context", {})
            if isinstance(pdp_context, dict):
                constraints = pdp_context.get("constraints")
                if isinstance(constraints, dict) and constraints:
                    modified_args = {**request.arguments, **constraints}
                admin_reason = pdp_context.get("reason_admin")
                if admin_reason:
                    reason = str(admin_reason)

            audit_entry: dict[str, Any] = {
                "authzen_request": mapping.authzen_request.model_dump(),
                "original_args_hash": mapping.original_args_hash,
                "decision": allowed,
                "reason": reason,
                "evaluation_ms": elapsed_ms,
            }

            return ToolCallResult(
                allowed=allowed,
                reason=reason,
                modified_arguments=modified_args,
                audit_entry=audit_entry,
            )

        except Exception:
            elapsed_ms = (time.monotonic() - start) * 1000
            logger.exception(
                "ARIA PDP evaluation failed for tool=%s agent=%s (%.1fms)",
                request.tool_name,
                request.agent_id,
                elapsed_ms,
            )
            if self._fail_closed:
                return ToolCallResult(
                    allowed=False,
                    reason="Policy evaluation failed — access denied (fail-closed)",
                    audit_entry={
                        "error": "pdp_communication_failure",
                        "evaluation_ms": elapsed_ms,
                    },
                )
            raise

    def close(self) -> None:
        """Release HTTP resources owned by this interceptor."""
        if self._owns_client:
            self._client.close()

    def __repr__(self) -> str:
        return (
            f"ARIAToolInterceptor(pdp_url={self._pdp_url!r}, "
            f"pdp_application={self._pdp_application!r})"
        )

    def __enter__(self) -> ARIAToolInterceptor:
        return self

    def __exit__(self, *args: object) -> None:
        self.close()
