import sys
import os
import io, asyncio
import pytest

sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
    BedrockGuardrail,
    _redact_pii_matches,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching import DualCache
from unittest.mock import MagicMock, AsyncMock, patch


@pytest.mark.asyncio
async def test_bedrock_guardrails_pii_masking():
    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    guardrail = BedrockGuardrail(
        guardrailIdentifier="wf0hkdb5x07f",
        guardrailVersion="DRAFT",
    )

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"},
            {"role": "assistant", "content": "Hello, how can I help you today?"},
            {"role": "user", "content": "I need to cancel my order"},
            {
                "role": "user",
                "content": "ok, my credit card number is 1234-5678-9012-3456",
            },
        ],
    }

    response = await guardrail.async_moderation_hook(
        data=request_data,
        user_api_key_dict=mock_user_api_key_dict,
        call_type="completion",
    )
    print("response after moderation hook", response)

    if response:  # Only assert if response is not None
        assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}"
        assert response["messages"][1]["content"] == "Hello, how can I help you today?"
        assert response["messages"][2]["content"] == "I need to cancel my order"
        assert (
            response["messages"][3]["content"]
            == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}"
        )


@pytest.mark.asyncio
async def test_bedrock_guardrails_pii_masking_content_list():
    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    guardrail = BedrockGuardrail(
        guardrailIdentifier="wf0hkdb5x07f",
        guardrailVersion="DRAFT",
    )

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Hello, my phone number is +1 412 555 1212",
                    },
                    {"type": "text", "text": "what time is it?"},
                ],
            },
            {"role": "assistant", "content": "Hello, how can I help you today?"},
            {"role": "user", "content": "who is the president of the united states?"},
        ],
    }

    response = await guardrail.async_moderation_hook(
        data=request_data,
        user_api_key_dict=mock_user_api_key_dict,
        call_type="completion",
    )
    print(response)

    if response:  # Only assert if response is not None
        # Verify that the list content is properly masked
        assert isinstance(response["messages"][0]["content"], list)
        assert (
            response["messages"][0]["content"][0]["text"]
            == "Hello, my phone number is {PHONE}"
        )
        assert response["messages"][0]["content"][1]["text"] == "what time is it?"
        assert response["messages"][1]["content"] == "Hello, how can I help you today?"
        assert (
            response["messages"][2]["content"]
            == "who is the president of the united states?"
        )


@pytest.mark.asyncio
async def test_bedrock_guardrails_block_messages_api():
    """
    Test that guardrails block messages API requests containing 'coffee' and raise the expected exception.
    """
    from fastapi import HTTPException

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    guardrail = BedrockGuardrail(
        guardrailIdentifier="ff6ujrregl1q",
        guardrailVersion="DRAFT",
    )

    request_data = {
        "model": "claude-sonnet-4-5-20250929",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Hello, my phone number is +1 412 555 1212",
                    },
                    {"type": "text", "text": "what time is it?"},
                ],
            },
            {"role": "user", "content": "tell me about coffee"},
        ],
    }

    with pytest.raises(HTTPException) as exc_info:
        await guardrail.async_pre_call_hook(
            data=request_data,
            user_api_key_dict=mock_user_api_key_dict,
            call_type="anthropic_messages",
            cache=MagicMock(spec=DualCache),
        )

    exception = exc_info.value
    assert exception.status_code == 400
    detail = exception.detail
    assert isinstance(detail, dict)
    assert detail["error"] == "Violated guardrail policy"
    assert (
        detail["bedrock_guardrail_response"]
        == "Sorry, the model cannot answer this question. coffee guardrail applied "
    )


@pytest.mark.asyncio
async def test_bedrock_guardrails_block_responses_api():
    """
    Test that guardrails block responses API requests containing 'coffee' and raise the expected exception.
    """
    from fastapi import HTTPException

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    guardrail = BedrockGuardrail(
        guardrailIdentifier="ff6ujrregl1q",
        guardrailVersion="DRAFT",
    )

    request_data = {
        "model": "gpt-4.1",
        "input": "Tell me a three sentence bedtime story about a unicorn drinking coffee",
        "stream": False,
    }

    with pytest.raises(HTTPException) as exc_info:
        await guardrail.async_pre_call_hook(
            data=request_data,
            user_api_key_dict=mock_user_api_key_dict,
            call_type="responses",
            cache=MagicMock(spec=DualCache),
        )

    exception = exc_info.value
    assert exception.status_code == 400
    detail = exception.detail
    assert isinstance(detail, dict)
    assert detail["error"] == "Violated guardrail policy"
    assert (
        detail["bedrock_guardrail_response"]
        == "Sorry, the model cannot answer this question. coffee guardrail applied "
    )


@pytest.mark.asyncio
async def test_bedrock_guardrails_with_streaming():
    from litellm.proxy.utils import ProxyLogging
    from litellm.types.guardrails import GuardrailEventHooks

    # Create proper mock objects
    mock_user_api_key_cache = MagicMock(spec=DualCache)
    mock_user_api_key_dict = UserAPIKeyAuth()

    with pytest.raises(Exception):  # Assert that this raises an exception
        proxy_logging_obj = ProxyLogging(
            user_api_key_cache=mock_user_api_key_cache,
            premium_user=True,
        )

        guardrail = BedrockGuardrail(
            guardrailIdentifier="ff6ujrregl1q",
            guardrailVersion="DRAFT",
            supported_event_hooks=[GuardrailEventHooks.post_call],
            guardrail_name="bedrock-post-guard",
        )

        litellm.callbacks.append(guardrail)

        request_data = {
            "model": "gpt-5.5",
            "messages": [{"role": "user", "content": "Hi I like coffee"}],
            "stream": True,
            "metadata": {"guardrails": ["bedrock-post-guard"]},
        }

        response = await litellm.acompletion(
            **request_data,
        )

        response = proxy_logging_obj.async_post_call_streaming_iterator_hook(
            user_api_key_dict=mock_user_api_key_dict,
            response=response,
            request_data=request_data,
        )

        async for chunk in response:
            print(chunk)


@pytest.mark.asyncio
async def test_bedrock_guardrails_with_streaming_no_violation():
    from litellm.proxy.utils import ProxyLogging
    from litellm.types.guardrails import GuardrailEventHooks

    # Create proper mock objects
    mock_user_api_key_cache = MagicMock(spec=DualCache)
    mock_user_api_key_dict = UserAPIKeyAuth()

    proxy_logging_obj = ProxyLogging(
        user_api_key_cache=mock_user_api_key_cache,
        premium_user=True,
    )

    guardrail = BedrockGuardrail(
        guardrailIdentifier="ff6ujrregl1q",
        guardrailVersion="DRAFT",
        supported_event_hooks=[GuardrailEventHooks.post_call],
        guardrail_name="bedrock-post-guard",
    )

    litellm.callbacks.append(guardrail)

    request_data = {
        "model": "gpt-5.5",
        "messages": [{"role": "user", "content": "hi"}],
        "stream": True,
        "metadata": {"guardrails": ["bedrock-post-guard"]},
    }

    response = await litellm.acompletion(
        **request_data,
    )

    response = proxy_logging_obj.async_post_call_streaming_iterator_hook(
        user_api_key_dict=mock_user_api_key_dict,
        response=response,
        request_data=request_data,
    )

    async for chunk in response:
        print(chunk)


@pytest.mark.asyncio
async def test_bedrock_guardrails_streaming_request_body_mock():
    """Test that the exact request body sent to Bedrock matches expected format when using streaming"""
    import json
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.caching import DualCache
    from litellm.types.guardrails import GuardrailEventHooks

    # Create mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()
    mock_cache = MagicMock(spec=DualCache)

    # Create the guardrail
    guardrail = BedrockGuardrail(
        guardrailIdentifier="wf0hkdb5x07f",
        guardrailVersion="DRAFT",
        supported_event_hooks=[GuardrailEventHooks.post_call],
        guardrail_name="bedrock-post-guard",
    )

    # Mock the assembled response from streaming
    mock_response = litellm.ModelResponse(
        id="test-id",
        choices=[
            litellm.Choices(
                index=0,
                message=litellm.Message(
                    role="assistant", content="The capital of Spain is Madrid."
                ),
                finish_reason="stop",
            )
        ],
        created=1234567890,
        model="gpt-5.5",
        object="chat.completion",
    )

    # Mock Bedrock API response
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {"action": "NONE", "outputs": []}

    # Patch the async_handler.post method to capture the request body
    with patch.object(guardrail, "async_handler") as mock_async_handler:
        mock_async_handler.post = AsyncMock(return_value=mock_bedrock_response)

        # Test data - simulating request data and assembled response
        request_data = {
            "model": "gpt-5.5",
            "messages": [{"role": "user", "content": "what's the capital of spain?"}],
            "stream": True,
            "metadata": {"guardrails": ["bedrock-post-guard"]},
        }

        # Call the method that should make the Bedrock API request
        await guardrail.make_bedrock_api_request(
            source="OUTPUT", response=mock_response, request_data=request_data
        )

        # Verify the API call was made
        mock_async_handler.post.assert_called_once()

        # Get the request data that was passed
        call_args = mock_async_handler.post.call_args

        # The data should be in the 'data' parameter of the prepared request
        # We need to parse the JSON from the prepared request body
        prepared_request_body = call_args.kwargs.get("data")

        # Parse the JSON body
        if isinstance(prepared_request_body, bytes):
            actual_body = json.loads(prepared_request_body.decode("utf-8"))
        else:
            actual_body = json.loads(prepared_request_body)

        # Expected body based on the convert_to_bedrock_format method behavior
        expected_body = {
            "source": "OUTPUT",
            "content": [{"text": {"text": "The capital of Spain is Madrid."}}],
        }

        print("Actual Bedrock request body:", json.dumps(actual_body, indent=2))
        print("Expected Bedrock request body:", json.dumps(expected_body, indent=2))

        # Assert the request body matches exactly
        assert (
            actual_body == expected_body
        ), f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}"


@pytest.mark.asyncio
async def test_bedrock_guardrail_aws_param_persistence():
    """Test that AWS auth params set on init are used for every request and not popped out."""
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.guardrails import GuardrailEventHooks

    guardrail = BedrockGuardrail(
        guardrailIdentifier="wf0hkdb5x07f",
        guardrailVersion="DRAFT",
        aws_access_key_id="test-access-key",
        aws_secret_access_key="test-secret-key",
        aws_region_name="us-east-1",
        supported_event_hooks=[GuardrailEventHooks.post_call],
        guardrail_name="bedrock-post-guard",
    )

    with patch.object(
        guardrail, "get_credentials", wraps=guardrail.get_credentials
    ) as mock_get_creds:
        for i in range(3):
            request_data = {
                "model": "gpt-5.5",
                "messages": [{"role": "user", "content": f"request {i}"}],
                "stream": False,
                "metadata": {"guardrails": ["bedrock-post-guard"]},
            }
            with patch.object(
                guardrail.async_handler, "post", new_callable=AsyncMock
            ) as mock_post:
                # Configure the mock response properly
                mock_response = AsyncMock()
                mock_response.status_code = 200
                mock_response.json = MagicMock(
                    return_value={"action": "NONE", "outputs": []}
                )
                mock_post.return_value = mock_response
                await guardrail.make_bedrock_api_request(
                    source="INPUT",
                    messages=request_data.get("messages"),
                    request_data=request_data,
                )

        assert mock_get_creds.call_count == 3
        for call in mock_get_creds.call_args_list:
            kwargs = call.kwargs
            print("used the following kwargs to get credentials=", kwargs)
            assert kwargs["aws_access_key_id"] == "test-access-key"
            assert kwargs["aws_secret_access_key"] == "test-secret-key"
            assert kwargs["aws_region_name"] == "us-east-1"


@pytest.mark.asyncio
async def test_bedrock_guardrail_blocked_vs_anonymized_actions():
    """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not"""
    from unittest.mock import MagicMock
    from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
        BedrockGuardrail,
    )
    from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
        BedrockGuardrailResponse,
    )

    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Test 1: ANONYMIZED action should NOT raise exception
    anonymized_response: BedrockGuardrailResponse = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "Hello, my phone number is {PHONE}"}],
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": [
                        {
                            "type": "PHONE",
                            "match": "+1 412 555 1212",
                            "action": "ANONYMIZED",
                        }
                    ]
                }
            }
        ],
    }

    should_raise = guardrail._should_raise_guardrail_blocked_exception(
        anonymized_response
    )
    assert should_raise is False, "ANONYMIZED actions should not raise exceptions"

    # Test 2: BLOCKED action should raise exception
    blocked_response: BedrockGuardrailResponse = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "I can't provide that information."}],
        "assessments": [
            {
                "topicPolicy": {
                    "topics": [
                        {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}
                    ]
                }
            }
        ],
    }

    should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response)
    assert should_raise is True, "BLOCKED actions should raise exceptions"

    # Test 3: Mixed actions - should raise if ANY action is BLOCKED
    mixed_response: BedrockGuardrailResponse = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "I can't provide that information."}],
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": [
                        {
                            "type": "PHONE",
                            "match": "+1 412 555 1212",
                            "action": "ANONYMIZED",
                        }
                    ]
                },
                "topicPolicy": {
                    "topics": [
                        {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"}
                    ]
                },
            }
        ],
    }

    should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response)
    assert (
        should_raise is True
    ), "Mixed actions with any BLOCKED should raise exceptions"

    # Test 4: NONE action should not raise exception
    none_response: BedrockGuardrailResponse = {
        "action": "NONE",
        "outputs": [],
        "assessments": [],
    }

    should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response)
    assert should_raise is False, "NONE actions should not raise exceptions"

    # Test 5: Test other policy types with BLOCKED actions
    content_blocked_response: BedrockGuardrailResponse = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "I can't provide that information."}],
        "assessments": [
            {
                "contentPolicy": {
                    "filters": [
                        {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}
                    ]
                }
            }
        ],
    }

    should_raise = guardrail._should_raise_guardrail_blocked_exception(
        content_blocked_response
    )
    assert (
        should_raise is True
    ), "Content policy BLOCKED actions should raise exceptions"


@pytest.mark.asyncio
async def test_bedrock_guardrail_masking_with_anonymized_response():
    """Test that masking works correctly when guardrail returns ANONYMIZED actions"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.caching import DualCache

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
        mask_request_content=True,
    )

    # Mock the Bedrock API response with ANONYMIZED action
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "Hello, my phone number is {PHONE}"}],
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": [
                        {
                            "type": "PHONE",
                            "match": "+1 412 555 1212",
                            "action": "ANONYMIZED",
                        }
                    ]
                }
            }
        ],
    }

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"},
        ],
    }

    # Patch the async_handler.post method
    with patch.object(
        guardrail.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # This should NOT raise an exception since action is ANONYMIZED
        try:
            response = await guardrail.async_moderation_hook(
                data=request_data,
                user_api_key_dict=mock_user_api_key_dict,
                call_type="completion",
            )
            # Should succeed and return data with masked content
            assert response is not None
            assert (
                response["messages"][0]["content"]
                == "Hello, my phone number is {PHONE}"
            )
        except Exception as e:
            pytest.fail(
                f"Should not raise exception for ANONYMIZED actions, but got: {e}"
            )


@pytest.mark.asyncio
async def test_bedrock_guardrail_uses_masked_output_without_masking_flags():
    """Test that masked output from guardrails is used even when masking flags are not enabled"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    # Create guardrail WITHOUT masking flags enabled
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
        # Note: No mask_request_content=True or mask_response_content=True
    )

    # Mock the Bedrock API response with ANONYMIZED action and masked output
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "Hello, my phone number is {PHONE} and email is {EMAIL}"}],
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": [
                        {
                            "type": "PHONE",
                            "match": "+1 412 555 1212",
                            "action": "ANONYMIZED",
                        },
                        {
                            "type": "EMAIL",
                            "match": "user@example.com",
                            "action": "ANONYMIZED",
                        },
                    ]
                }
            }
        ],
    }

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {
                "role": "user",
                "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com",
            },
        ],
    }

    # Patch the async_handler.post method
    with patch.object(
        guardrail.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # This should use the masked output even without masking flags
        response = await guardrail.async_moderation_hook(
            data=request_data,
            user_api_key_dict=mock_user_api_key_dict,
            call_type="completion",
        )

        # Should use the masked content from guardrail output
        assert response is not None
        assert (
            response["messages"][0]["content"]
            == "Hello, my phone number is {PHONE} and email is {EMAIL}"
        )
        print("✅ Masked output was applied even without masking flags enabled")


@pytest.mark.asyncio
async def test_bedrock_guardrail_response_pii_masking_non_streaming():
    """Test that PII masking is applied to response content in non-streaming scenarios"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    # Create guardrail with response masking enabled
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
    )

    # Mock the Bedrock API response with ANONYMIZED PII
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [
            {
                "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}"
            }
        ],
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": [
                        {
                            "type": "CREDIT_DEBIT_CARD_NUMBER",
                            "match": "1234-5678-9012-3456",
                            "action": "ANONYMIZED",
                        },
                        {
                            "type": "PHONE",
                            "match": "+1 412 555 1212",
                            "action": "ANONYMIZED",
                        },
                    ]
                }
            }
        ],
    }

    # Create a mock response that contains PII
    mock_response = litellm.ModelResponse(
        id="test-id",
        choices=[
            litellm.Choices(
                index=0,
                message=litellm.Message(
                    role="assistant",
                    content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212",
                ),
                finish_reason="stop",
            )
        ],
        created=1234567890,
        model="gpt-5.5",
        object="chat.completion",
    )

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "What's your credit card and phone number?"},
        ],
    }

    # Patch the async_handler.post method
    with patch.object(
        guardrail.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # Call the post-call success hook
        await guardrail.async_post_call_success_hook(
            data=request_data,
            user_api_key_dict=mock_user_api_key_dict,
            response=mock_response,
        )

        # Verify that the response content was masked
        assert (
            mock_response.choices[0].message.content
            == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}"
        )
        print("✓ Non-streaming response PII masking test passed")


@pytest.mark.asyncio
async def test_bedrock_guardrail_response_pii_masking_streaming():
    """Test that PII masking is applied to response content in streaming scenarios"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.utils import ModelResponseStream

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    # Create guardrail with response masking enabled
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
    )

    # Mock the Bedrock API response with ANONYMIZED PII
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "Sure! My email is {EMAIL} and SSN is {US_SSN}"}],
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": [
                        {
                            "type": "EMAIL",
                            "match": "john@example.com",
                            "action": "ANONYMIZED",
                        },
                        {
                            "type": "US_SSN",
                            "match": "123-45-6789",
                            "action": "ANONYMIZED",
                        },
                    ]
                }
            }
        ],
    }

    # Create mock streaming chunks
    async def mock_streaming_response():
        chunks = [
            ModelResponseStream(
                id="test-id",
                choices=[
                    litellm.utils.StreamingChoices(
                        index=0,
                        delta=litellm.utils.Delta(content="Sure! My email is "),
                        finish_reason=None,
                    )
                ],
                created=1234567890,
                model="gpt-5.5",
                object="chat.completion.chunk",
            ),
            ModelResponseStream(
                id="test-id",
                choices=[
                    litellm.utils.StreamingChoices(
                        index=0,
                        delta=litellm.utils.Delta(
                            content="john@example.com and SSN is "
                        ),
                        finish_reason=None,
                    )
                ],
                created=1234567890,
                model="gpt-5.5",
                object="chat.completion.chunk",
            ),
            ModelResponseStream(
                id="test-id",
                choices=[
                    litellm.utils.StreamingChoices(
                        index=0,
                        delta=litellm.utils.Delta(content="123-45-6789"),
                        finish_reason="stop",
                    )
                ],
                created=1234567890,
                model="gpt-5.5",
                object="chat.completion.chunk",
            ),
        ]
        for chunk in chunks:
            yield chunk

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "What's your email and SSN?"},
        ],
        "stream": True,
    }

    # Patch the async_handler.post method
    with patch.object(
        guardrail.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # Call the streaming hook
        masked_stream = guardrail.async_post_call_streaming_iterator_hook(
            user_api_key_dict=mock_user_api_key_dict,
            response=mock_streaming_response(),
            request_data=request_data,
        )

        # Collect all chunks from the masked stream
        masked_chunks = []
        async for chunk in masked_stream:
            masked_chunks.append(chunk)

        # Verify that we got chunks back
        assert len(masked_chunks) > 0

        # Reconstruct the full response from chunks to verify masking
        full_content = ""
        for chunk in masked_chunks:
            if hasattr(chunk, "choices") and chunk.choices:
                if hasattr(chunk.choices[0], "delta") and chunk.choices[0].delta:
                    if (
                        hasattr(chunk.choices[0].delta, "content")
                        and chunk.choices[0].delta.content
                    ):
                        full_content += chunk.choices[0].delta.content

        # Verify that the reconstructed content contains the masked PII
        assert "Sure! My email is {EMAIL} and SSN is {US_SSN}" == full_content
        print("✓ Streaming response PII masking test passed")


@pytest.mark.asyncio
async def test_convert_to_bedrock_format_input_source():
    """Test convert_to_bedrock_format with INPUT source and mock messages"""
    from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
        BedrockGuardrail,
    )
    from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
        BedrockRequest,
    )
    from unittest.mock import patch

    # Create the guardrail instance
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Mock messages
    mock_messages = [
        {"role": "user", "content": "Hello, how are you?"},
        {"role": "assistant", "content": "I'm doing well, thank you!"},
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What's the weather like?"},
                {"type": "text", "text": "Is it sunny today?"},
            ],
        },
    ]

    # Call the method
    result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages)

    # Verify the result structure
    assert isinstance(result, dict)
    assert result.get("source") == "INPUT"
    assert "content" in result
    assert isinstance(result.get("content"), list)

    # Verify content items
    expected_content_items = [
        {"text": {"text": "Hello, how are you?"}},
        {"text": {"text": "I'm doing well, thank you!"}},
        {"text": {"text": "What's the weather like?"}},
        {"text": {"text": "Is it sunny today?"}},
    ]

    assert result.get("content") == expected_content_items
    print("✅ INPUT source test passed - result:", result)


@pytest.mark.asyncio
async def test_convert_to_bedrock_format_output_source():
    """Test convert_to_bedrock_format with OUTPUT source and mock ModelResponse"""
    from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
        BedrockGuardrail,
    )
    from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
        BedrockRequest,
    )
    import litellm
    from unittest.mock import patch

    # Create the guardrail instance
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Mock ModelResponse
    mock_response = litellm.ModelResponse(
        id="test-response-id",
        choices=[
            litellm.Choices(
                index=0,
                message=litellm.Message(
                    role="assistant", content="This is a test response from the model."
                ),
                finish_reason="stop",
            ),
            litellm.Choices(
                index=1,
                message=litellm.Message(
                    role="assistant", content="This is a second choice response."
                ),
                finish_reason="stop",
            ),
        ],
        created=1234567890,
        model="gpt-5.5",
        object="chat.completion",
    )

    # Call the method
    result = guardrail.convert_to_bedrock_format(
        source="OUTPUT", response=mock_response
    )

    # Verify the result structure
    assert isinstance(result, dict)
    assert result.get("source") == "OUTPUT"
    assert "content" in result
    assert isinstance(result.get("content"), list)

    # Verify content items - should contain both choice contents
    expected_content_items = [
        {"text": {"text": "This is a test response from the model."}},
        {"text": {"text": "This is a second choice response."}},
    ]

    assert result.get("content") == expected_content_items
    print("✅ OUTPUT source test passed - result:", result)


@pytest.mark.asyncio
async def test_convert_to_bedrock_format_post_call_streaming_hook():
    """Test async_post_call_streaming_iterator_hook makes OUTPUT bedrock request and applies masking"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.utils import ModelResponseStream
    import litellm

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    # Create guardrail instance
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Mock streaming chunks that contain PII
    async def mock_streaming_response():
        chunks = [
            ModelResponseStream(
                id="test-id",
                choices=[
                    litellm.utils.StreamingChoices(
                        index=0,
                        delta=litellm.utils.Delta(content="My email is "),
                        finish_reason=None,
                    )
                ],
                created=1234567890,
                model="gpt-5.5",
                object="chat.completion.chunk",
            ),
            ModelResponseStream(
                id="test-id",
                choices=[
                    litellm.utils.StreamingChoices(
                        index=0,
                        delta=litellm.utils.Delta(content="john@example.com"),
                        finish_reason="stop",
                    )
                ],
                created=1234567890,
                model="gpt-5.5",
                object="chat.completion.chunk",
            ),
        ]
        for chunk in chunks:
            yield chunk

    # Mock Bedrock API response with PII masking
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "My email is {EMAIL}"}],
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": [
                        {
                            "type": "EMAIL",
                            "match": "john@example.com",
                            "action": "ANONYMIZED",
                        }
                    ]
                }
            }
        ],
    }

    request_data = {
        "model": "gpt-5.5",
        "messages": [{"role": "user", "content": "What's your email?"}],
        "stream": True,
    }

    # Track which bedrock API calls were made
    bedrock_calls = []

    # Mock the make_bedrock_api_request method to track calls
    async def mock_make_bedrock_api_request(
        source,
        messages=None,
        response=None,
        request_data=None,
        logging_event_type=None,
        **kwargs,
    ):
        bedrock_calls.append(
            {
                "source": source,
                "messages": messages,
                "response": response,
                "request_data": request_data,
                "logging_event_type": logging_event_type,
            }
        )
        # Return the mock bedrock response
        from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
            BedrockGuardrailResponse,
        )

        return BedrockGuardrailResponse(**mock_bedrock_response.json())

    # Patch the bedrock API request method
    with patch.object(
        guardrail, "make_bedrock_api_request", side_effect=mock_make_bedrock_api_request
    ):

        # Call the streaming hook
        result_generator = guardrail.async_post_call_streaming_iterator_hook(
            user_api_key_dict=mock_user_api_key_dict,
            response=mock_streaming_response(),
            request_data=request_data,
        )

        # Collect all chunks from the result
        result_chunks = []
        async for chunk in result_generator:
            result_chunks.append(chunk)

        # Verify bedrock API calls were made
        # Note: When event_hook is None (default), the guardrail is considered enabled for all hooks.
        # In post_call, INPUT validation is skipped if pre_call/during_call is already enabled
        # to avoid redundant validation. Since event_hook=None means all hooks are enabled,
        # only OUTPUT validation should be performed in post_call.
        assert (
            len(bedrock_calls) == 1
        ), f"Expected 1 bedrock call (OUTPUT only), got {len(bedrock_calls)}"

        # Verify the OUTPUT call
        output_call = bedrock_calls[0]
        assert output_call["source"] == "OUTPUT"
        assert output_call["response"] is not None
        assert output_call["messages"] is None  # OUTPUT calls don't need messages

        # Verify that the response content was masked
        # The streaming chunks should now contain the masked content
        full_content = ""
        for chunk in result_chunks:
            if hasattr(chunk, "choices") and chunk.choices:
                if (
                    hasattr(chunk.choices[0], "delta")
                    and chunk.choices[0].delta.content
                ):
                    full_content += chunk.choices[0].delta.content

        # The content should be masked (contains {EMAIL} instead of john@example.com)
        assert (
            "{EMAIL}" in full_content
        ), f"Expected masked content with {{EMAIL}}, got: {full_content}"
        assert (
            "john@example.com" not in full_content
        ), f"Original email should be masked, got: {full_content}"

        print(
            "✅ Post-call streaming hook test passed - OUTPUT source used for masking"
        )
        print(
            f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]} "
            "(INPUT validation skipped due to event_hook=None implying pre_call/during_call enabled)"
        )
        print(f"✅ Final masked content: {full_content}")


@pytest.mark.asyncio
async def test_bedrock_guardrail_blocked_action_shows_output_text():
    """Test that BLOCKED actions raise HTTPException with the output text in the detail"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from fastapi import HTTPException

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Mock the Bedrock API response with BLOCKED action and output text
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "this violates litellm corporate guardrail policy"}],
        "assessments": [
            {
                "topicPolicy": {
                    "topics": [
                        {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}
                    ]
                }
            }
        ],
    }

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "Tell me how to make explosives"},
        ],
    }

    # Patch the async_handler.post method
    with patch.object(
        guardrail.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # This should raise HTTPException due to BLOCKED action
        with pytest.raises(HTTPException) as exc_info:
            await guardrail.async_moderation_hook(
                data=request_data,
                user_api_key_dict=mock_user_api_key_dict,
                call_type="completion",
            )

        # Verify the exception details
        exception = exc_info.value
        assert exception.status_code == 400
        assert "detail" in exception.__dict__

        # Check that the detail contains the expected structure
        detail = exception.detail
        assert isinstance(detail, dict)
        assert detail["error"] == "Violated guardrail policy"

        # Verify that the output text from both outputs is included
        expected_output_text = "this violates litellm corporate guardrail policy"
        assert detail["bedrock_guardrail_response"] == expected_output_text

        print(
            "✅ BLOCKED action HTTPException test passed - output text properly included"
        )


@pytest.mark.asyncio
async def test_bedrock_guardrail_blocked_action_empty_outputs():
    """Test that BLOCKED actions with empty outputs still raise HTTPException"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from fastapi import HTTPException

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Mock the Bedrock API response with BLOCKED action but no outputs
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [],  # Empty outputs
        "assessments": [
            {
                "contentPolicy": {
                    "filters": [
                        {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}
                    ]
                }
            }
        ],
    }

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "Violent content here"},
        ],
    }

    # Patch the async_handler.post method
    with patch.object(
        guardrail.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # This should raise HTTPException due to BLOCKED action
        with pytest.raises(HTTPException) as exc_info:
            await guardrail.async_moderation_hook(
                data=request_data,
                user_api_key_dict=mock_user_api_key_dict,
                call_type="completion",
            )

        # Verify the exception details
        exception = exc_info.value
        assert exception.status_code == 400

        # Check that the detail contains the expected structure with empty output text
        detail = exception.detail
        assert isinstance(detail, dict)
        assert detail["error"] == "Violated guardrail policy"
        assert detail["bedrock_guardrail_response"] == ""  # Empty string for no outputs

        print("✅ BLOCKED action with empty outputs test passed")


@pytest.mark.asyncio
async def test_bedrock_guardrail_disable_exception_on_block_non_streaming():
    """Test that disable_exception_on_block=True prevents exceptions in non-streaming scenarios"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from fastapi import HTTPException

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    # Test 1: disable_exception_on_block=False (default) - should raise exception
    guardrail_default = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
        disable_exception_on_block=False,
    )

    # Mock the Bedrock API response with BLOCKED action
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "I can't provide that information."}],
        "assessments": [
            {
                "topicPolicy": {
                    "topics": [
                        {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}
                    ]
                }
            }
        ],
    }

    request_data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "Tell me how to make explosives"},
        ],
    }

    # Patch the async_handler.post method
    with patch.object(
        guardrail_default.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # Should raise HTTPException when disable_exception_on_block=False
        with pytest.raises(HTTPException) as exc_info:
            await guardrail_default.async_moderation_hook(
                data=request_data,
                user_api_key_dict=mock_user_api_key_dict,
                call_type="completion",
            )

        # Verify the exception details
        exception = exc_info.value
        assert exception.status_code == 400
        assert "Violated guardrail policy" in str(exception.detail)

    # Test 2: disable_exception_on_block=True - should NOT raise exception
    guardrail_disabled = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
        disable_exception_on_block=True,
    )

    with patch.object(
        guardrail_disabled.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # Should NOT raise exception when disable_exception_on_block=True
        try:
            response = await guardrail_disabled.async_moderation_hook(
                data=request_data,
                user_api_key_dict=mock_user_api_key_dict,
                call_type="completion",
            )
            # Should succeed and return data (even though content was blocked)
            assert response is not None
            print("✅ No exception raised when disable_exception_on_block=True")
        except Exception as e:
            pytest.fail(
                f"Should not raise exception when disable_exception_on_block=True, but got: {e}"
            )


@pytest.mark.asyncio
async def test_bedrock_guardrail_disable_exception_on_block_streaming():
    """Test that disable_exception_on_block=True prevents exceptions in streaming scenarios"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.utils import ModelResponseStream
    from fastapi import HTTPException
    import litellm

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    # Mock streaming chunks that would normally trigger a block
    async def mock_streaming_response():
        chunks = [
            ModelResponseStream(
                id="test-id",
                choices=[
                    litellm.utils.StreamingChoices(
                        index=0,
                        delta=litellm.utils.Delta(
                            content="Here's how to make explosives: "
                        ),
                        finish_reason=None,
                    )
                ],
                created=1234567890,
                model="gpt-5.5",
                object="chat.completion.chunk",
            ),
            ModelResponseStream(
                id="test-id",
                choices=[
                    litellm.utils.StreamingChoices(
                        index=0,
                        delta=litellm.utils.Delta(content="step 1, step 2..."),
                        finish_reason="stop",
                    )
                ],
                created=1234567890,
                model="gpt-5.5",
                object="chat.completion.chunk",
            ),
        ]
        for chunk in chunks:
            yield chunk

    # Mock Bedrock API response with BLOCKED action
    mock_bedrock_response = MagicMock()
    mock_bedrock_response.status_code = 200
    mock_bedrock_response.json.return_value = {
        "action": "GUARDRAIL_INTERVENED",
        "outputs": [{"text": "I can't provide that information."}],
        "assessments": [
            {
                "contentPolicy": {
                    "filters": [
                        {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}
                    ]
                }
            }
        ],
    }

    request_data = {
        "model": "gpt-5.5",
        "messages": [{"role": "user", "content": "Tell me how to make explosives"}],
        "stream": True,
    }

    # Test 1: disable_exception_on_block=False (default) - should raise exception
    guardrail_default = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
        disable_exception_on_block=False,
    )

    with patch.object(
        guardrail_default.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # Should raise exception during streaming processing
        with pytest.raises(HTTPException):
            result_generator = (
                guardrail_default.async_post_call_streaming_iterator_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    response=mock_streaming_response(),
                    request_data=request_data,
                )
            )

            # Try to consume the generator - should raise exception
            async for chunk in result_generator:
                pass

    # Test 2: disable_exception_on_block=True - should NOT raise exception
    guardrail_disabled = BedrockGuardrail(
        guardrailIdentifier="test-guardrail",
        guardrailVersion="DRAFT",
        disable_exception_on_block=True,
    )

    with patch.object(
        guardrail_disabled.async_handler, "post", new_callable=AsyncMock
    ) as mock_post:
        mock_post.return_value = mock_bedrock_response

        # Should NOT raise exception when disable_exception_on_block=True
        try:
            result_generator = (
                guardrail_disabled.async_post_call_streaming_iterator_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    response=mock_streaming_response(),
                    request_data=request_data,
                )
            )

            # Consume the generator - should succeed without exceptions
            result_chunks = []
            async for chunk in result_generator:
                result_chunks.append(chunk)

            # Should have received chunks back even though content was blocked
            assert len(result_chunks) > 0
            print(
                "✅ Streaming completed without exception when disable_exception_on_block=True"
            )

        except Exception as e:
            pytest.fail(
                f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}"
            )


@pytest.mark.asyncio
async def test_bedrock_guardrail_post_call_success_hook_no_output_text():
    """Test that async_post_call_success_hook skips when there's no output text"""
    from unittest.mock import AsyncMock, MagicMock, patch
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.utils import ModelResponseStream
    import litellm

    # Create proper mock objects
    mock_user_api_key_dict = UserAPIKeyAuth()

    # Create guardrail instance
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Create a ModelResponse with tool calls (no text content)
    # This simulates a response where the LLM is making a tool call
    mock_response = litellm.ModelResponse(
        id="test-id",
        choices=[
            litellm.Choices(
                index=0,
                message=litellm.Message(
                    role="assistant",
                    content=None,  # No text content
                    tool_calls=[
                        litellm.utils.ChatCompletionMessageToolCall(
                            id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q",
                            function=litellm.utils.Function(
                                name="top_song", arguments='{"sign": "WZPZ"}'
                            ),
                            type="function",
                        )
                    ],
                ),
                finish_reason="tool_calls",
            )
        ],
        created=1234567890,
        model="gpt-5.5",
        object="chat.completion",
    )

    data = {
        "model": "gpt-5.5",
        "messages": [
            {"role": "user", "content": "Hello"},
        ],
    }
    mock_user_api_key_dict = UserAPIKeyAuth()

    result = await guardrail.async_post_call_success_hook(
        data=data,
        response=mock_response,
        user_api_key_dict=mock_user_api_key_dict,
    )
    # If no error is raised and result is None, then the test passes
    assert result is None
    print("✅ No output text in response test passed")


@pytest.mark.asyncio
async def test__redact_pii_matches_null_list_fields():
    """Test that explicit null values from Bedrock API are handled correctly.

    The Bedrock API can return explicit JSON null for list fields like
    piiEntities, regexes, customWords, managedWordLists. This would cause
    TypeError: 'NoneType' object is not iterable if not handled.
    """
    # Test 1: null piiEntities and regexes
    response_with_null_pii = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": [
            {
                "sensitiveInformationPolicy": {
                    "piiEntities": None,
                    "regexes": None,
                }
            }
        ],
    }
    redacted = _redact_pii_matches(response_with_null_pii)
    assert redacted is not None
    assert (
        redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] is None
    )
    assert redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] is None

    # Test 2: null customWords and managedWordLists
    response_with_null_words = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": [
            {
                "wordPolicy": {
                    "customWords": None,
                    "managedWordLists": None,
                }
            }
        ],
    }
    redacted = _redact_pii_matches(response_with_null_words)
    assert redacted is not None
    assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None
    assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None

    # Test 3: null assessments at top level
    response_with_null_assessments = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": None,
    }
    redacted = _redact_pii_matches(response_with_null_assessments)
    assert redacted is not None


@pytest.mark.asyncio
async def test__redact_pii_matches_malformed_response():
    """Test _redact_pii_matches with malformed response (should not crash)"""

    # Test with completely malformed response
    malformed_response = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": "not_a_list",
    }
    redacted_response = _redact_pii_matches(malformed_response)
    assert redacted_response == malformed_response

    # Test with missing keys
    missing_keys_response = {
        "action": "GUARDRAIL_INTERVENED",
    }
    redacted_response = _redact_pii_matches(missing_keys_response)
    assert redacted_response == missing_keys_response


@pytest.mark.asyncio
async def test_should_raise_guardrail_blocked_exception_null_fields():
    """Test that _should_raise_guardrail_blocked_exception handles null list fields.

    Validates the or [] null-safety pattern works for all policy fields
    in _should_raise_guardrail_blocked_exception.
    """
    guardrail = BedrockGuardrail(
        guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
    )

    # Test with null assessments
    response_null_assessments = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": None,
    }
    assert (
        guardrail._should_raise_guardrail_blocked_exception(response_null_assessments)
        is False
    )

    # Test with null topics in topicPolicy
    response_null_topics = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": [{"topicPolicy": {"topics": None}}],
    }
    assert (
        guardrail._should_raise_guardrail_blocked_exception(response_null_topics)
        is False
    )

    # Test with null filters in contentPolicy
    response_null_filters = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": [{"contentPolicy": {"filters": None}}],
    }
    assert (
        guardrail._should_raise_guardrail_blocked_exception(response_null_filters)
        is False
    )

    # Test with null customWords and managedWordLists in wordPolicy
    response_null_words = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": [
            {"wordPolicy": {"customWords": None, "managedWordLists": None}}
        ],
    }
    assert (
        guardrail._should_raise_guardrail_blocked_exception(response_null_words)
        is False
    )

    # Test with null piiEntities and regexes in sensitiveInformationPolicy
    response_null_pii = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": [
            {"sensitiveInformationPolicy": {"piiEntities": None, "regexes": None}}
        ],
    }
    assert (
        guardrail._should_raise_guardrail_blocked_exception(response_null_pii) is False
    )

    # Test with null filters in contextualGroundingPolicy
    response_null_grounding = {
        "action": "GUARDRAIL_INTERVENED",
        "assessments": [{"contextualGroundingPolicy": {"filters": None}}],
    }
    assert (
        guardrail._should_raise_guardrail_blocked_exception(response_null_grounding)
        is False
    )
