"""AutoPIL BedrockGuard. See packages/core/autopil/bedrock_guard.py for full docs."""

from __future__ import annotations
from typing import Any, Optional
from .guard import ContextGuard
from .models import SensitivityLevel


class BedrockGuard(ContextGuard):
    """ContextGuard for AWS Bedrock Agents. Tags audit events with source_type='bedrock'."""

    _source_type: str = "bedrock"

    def wrap_invoke_agent(self, client, *, agent_role, user_id, source_id,
                          sensitivity_level=SensitivityLevel.MEDIUM, session_id=None):
        """Wrap a boto3 bedrock-agent-runtime client so every invoke_agent() is policy-evaluated."""
        guard = self
        original_client = client

        def _invoke_impl(input_text: str, **invoke_kwargs) -> Any:
            return original_client.invoke_agent(inputText=input_text, **invoke_kwargs)

        protected = guard.protect(
            agent_role=agent_role, user_id=user_id, source_id=source_id,
            sensitivity_level=sensitivity_level, session_id=session_id,
        )(_invoke_impl)

        class _GuardedBedrockClient:
            def invoke_agent(self, *, inputText: str = "", **kwargs) -> Any:
                return protected(inputText, **kwargs)
            def __getattr__(self, name: str) -> Any:
                return getattr(original_client, name)

        return _GuardedBedrockClient()

    def wrap_invoke_agent_async(self, client, *, agent_role, user_id, source_id,
                                sensitivity_level=SensitivityLevel.MEDIUM, session_id=None):
        """Async version — wraps an aioboto3 bedrock-agent-runtime client."""
        guard = self
        original_client = client

        async def _invoke_impl_async(input_text: str, **invoke_kwargs) -> Any:
            result = original_client.invoke_agent(inputText=input_text, **invoke_kwargs)
            if hasattr(result, "__await__"):
                return await result
            return result

        protected = guard.protect_async(
            agent_role=agent_role, user_id=user_id, source_id=source_id,
            sensitivity_level=sensitivity_level, session_id=session_id,
        )(_invoke_impl_async)

        class _AsyncGuardedBedrockClient:
            async def invoke_agent(self, *, inputText: str = "", **kwargs) -> Any:
                return await protected(inputText, **kwargs)
            def __getattr__(self, name: str) -> Any:
                return getattr(original_client, name)

        return _AsyncGuardedBedrockClient()
