"""
Unit tests for Presidio PII Masking Guardrail
Tests PII detection and masking for different message formats
"""

import asyncio
import os
import sys
from contextlib import asynccontextmanager
from unittest.mock import MagicMock, patch

import pytest

sys.path.insert(0, os.path.abspath("../../../../../.."))

import litellm
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
    _OPTIONAL_PresidioPIIMasking,
)
from litellm.exceptions import GuardrailRaisedException
from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType
from litellm.types.utils import Choices, Message, ModelResponse


def _make_mock_session_iterator(
    json_response, status=200, content_type="application/json", text_response=""
):
    """Create a mock _get_session_iterator that yields a session returning json_response."""

    @asynccontextmanager
    async def mock_iterator():
        class MockResponse:
            def __init__(self):
                self.status = status
                self.content_type = content_type
                self.headers = {"Content-Type": content_type}

            async def text(self):
                if text_response:
                    return text_response
                import json

                try:
                    return json.dumps(json_response)
                except Exception:
                    return str(json_response)

            async def json(self):
                return json_response

            async def __aenter__(self):
                return self

            async def __aexit__(self, *args):
                pass

        class MockSession:
            def post(self, *args, **kwargs):
                self.last_kwargs = kwargs
                return MockResponse()

            async def __aenter__(self):
                return self

            async def __aexit__(self, *args):
                pass

        yield MockSession()

    return mock_iterator


@pytest.fixture
def presidio_guardrail():
    """Create a Presidio guardrail instance for testing"""
    return _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=False,
        pii_entities_config={
            PiiEntityType.CREDIT_CARD: PiiAction.MASK,
            PiiEntityType.EMAIL_ADDRESS: PiiAction.MASK,
            PiiEntityType.PHONE_NUMBER: PiiAction.MASK,
        },
    )


@pytest.fixture
def mock_user_api_key():
    """Create a mock user API key auth object"""
    return UserAPIKeyAuth(
        api_key="test_key",
        user_id="test_user",
    )


@pytest.fixture
def mock_cache():
    """Create a mock cache object"""
    return MagicMock(spec=DualCache)


@pytest.mark.asyncio
async def test_multimodal_message_format_completion_call_type(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test Presidio PII masking with multimodal message format (content as list)
    for completion call type.

    Tests the message format:
    {
        "role": "user",
        "content": [
            {
                "type": "text",
                "text": "My credit card number is 4111-1111-1111-1111..."
            }
        ]
    }
    """
    # Prepare test data with multimodal message format
    test_data = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "My credit card number is 4111-1111-1111-1111, my email is test@example.com, and my phone is 555-123-4567",
                    }
                ],
            }
        ],
        "model": "gpt-4",
    }

    # Mock the check_pii method to return redacted text
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        # Simulate PII detection and masking
        redacted_text = text
        redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]")
        redacted_text = redacted_text.replace("test@example.com", "[EMAIL]")
        redacted_text = redacted_text.replace("555-123-4567", "[PHONE]")
        return redacted_text

    presidio_guardrail.check_pii = mock_check_pii

    # Call the async_pre_call_hook with call_type="completion"
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    # Verify that PII was masked in the text field
    assert result is not None
    assert "messages" in result
    assert len(result["messages"]) == 1

    message = result["messages"][0]
    assert "content" in message
    assert isinstance(message["content"], list)
    assert len(message["content"]) == 1

    content_item = message["content"][0]
    assert content_item["type"] == "text"
    assert "[CREDIT_CARD]" in content_item["text"]
    assert "[EMAIL]" in content_item["text"]
    assert "[PHONE]" in content_item["text"]

    # Verify original PII is not present
    assert "4111-1111-1111-1111" not in content_item["text"]
    assert "test@example.com" not in content_item["text"]
    assert "555-123-4567" not in content_item["text"]

    print("✓ Multimodal message format test for completion call type passed")


@pytest.mark.asyncio
async def test_multimodal_message_format_anthropic_messages_call_type(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test Presidio PII masking with multimodal message format (content as list)
    for anthropic_messages call type.

    Tests the same message format but with anthropic_messages call type.
    """
    # Prepare test data with multimodal message format
    test_data = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "My credit card number is 4111-1111-1111-1111, my email is test@example.com, and my phone is 555-123-4567",
                    }
                ],
            }
        ],
        "model": "claude-3-opus-20240229",
    }

    # Mock the check_pii method to return redacted text
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        # Simulate PII detection and masking
        redacted_text = text
        redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]")
        redacted_text = redacted_text.replace("test@example.com", "[EMAIL]")
        redacted_text = redacted_text.replace("555-123-4567", "[PHONE]")
        return redacted_text

    presidio_guardrail.check_pii = mock_check_pii

    # Call the async_pre_call_hook with call_type="anthropic_messages"
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="anthropic_messages",
    )

    # Verify that PII was masked in the text field
    assert result is not None
    assert "messages" in result
    assert len(result["messages"]) == 1

    message = result["messages"][0]
    assert "content" in message
    assert isinstance(message["content"], list)
    assert len(message["content"]) == 1

    content_item = message["content"][0]
    assert content_item["type"] == "text"
    assert "[CREDIT_CARD]" in content_item["text"]
    assert "[EMAIL]" in content_item["text"]
    assert "[PHONE]" in content_item["text"]

    # Verify original PII is not present
    assert "4111-1111-1111-1111" not in content_item["text"]
    assert "test@example.com" not in content_item["text"]
    assert "555-123-4567" not in content_item["text"]

    print("✓ Multimodal message format test for anthropic_messages call type passed")


@pytest.mark.asyncio
async def test_multimodal_message_multiple_content_items(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test Presidio PII masking with multiple content items in the content list.
    """
    # Prepare test data with multiple content items
    test_data = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "My credit card is 4111-1111-1111-1111",
                    },
                    {
                        "type": "text",
                        "text": "My email is test@example.com",
                    },
                ],
            }
        ],
        "model": "gpt-4",
    }

    # Mock the check_pii method
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        redacted_text = text
        redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]")
        redacted_text = redacted_text.replace("test@example.com", "[EMAIL]")
        return redacted_text

    presidio_guardrail.check_pii = mock_check_pii

    # Call the async_pre_call_hook
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    # Verify both content items were processed
    assert result is not None
    message = result["messages"][0]
    content_items = message["content"]

    assert len(content_items) == 2
    assert "[CREDIT_CARD]" in content_items[0]["text"]
    assert "[EMAIL]" in content_items[1]["text"]

    print("✓ Multiple content items test passed")


@pytest.mark.asyncio
async def test_mixed_string_and_list_content(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test Presidio PII masking with mixed string and list content formats.
    """
    # Prepare test data with mixed content formats
    test_data = {
        "messages": [
            {
                "role": "user",
                "content": "My credit card is 4111-1111-1111-1111",
            },
            {
                "role": "assistant",
                "content": "I can help you with that.",
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "My email is test@example.com",
                    }
                ],
            },
        ],
        "model": "gpt-4",
    }

    # Mock the check_pii method
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        redacted_text = text
        redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]")
        redacted_text = redacted_text.replace("test@example.com", "[EMAIL]")
        return redacted_text

    presidio_guardrail.check_pii = mock_check_pii

    # Call the async_pre_call_hook
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    # Verify all messages were processed correctly
    assert result is not None
    messages = result["messages"]

    # First message (string content)
    assert isinstance(messages[0]["content"], str)
    assert "[CREDIT_CARD]" in messages[0]["content"]

    # Second message (string content, no PII)
    assert isinstance(messages[1]["content"], str)
    assert messages[1]["content"] == "I can help you with that."

    # Third message (list content)
    assert isinstance(messages[2]["content"], list)
    assert "[EMAIL]" in messages[2]["content"][0]["text"]

    print("✓ Mixed string and list content test passed")


@pytest.mark.asyncio
async def test_content_list_without_text_field(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test Presidio PII masking gracefully handles content items without text field
    (e.g., image content items).
    """
    # Prepare test data with image content (no text field)
    test_data = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {"url": "https://example.com/image.jpg"},
                    },
                    {
                        "type": "text",
                        "text": "What's in this image? My email is test@example.com",
                    },
                ],
            }
        ],
        "model": "gpt-4-vision",
    }

    # Mock the check_pii method
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        redacted_text = text.replace("test@example.com", "[EMAIL]")
        return redacted_text

    presidio_guardrail.check_pii = mock_check_pii

    # Call the async_pre_call_hook
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    # Verify that image content is preserved and text content is processed
    assert result is not None
    content_items = result["messages"][0]["content"]

    assert len(content_items) == 2
    # Image content should remain unchanged
    assert content_items[0]["type"] == "image_url"
    assert content_items[0]["image_url"]["url"] == "https://example.com/image.jpg"

    # Text content should be redacted
    assert content_items[1]["type"] == "text"
    assert "[EMAIL]" in content_items[1]["text"]

    print("✓ Content list without text field test passed")


@pytest.mark.asyncio
async def test_empty_messages(presidio_guardrail, mock_user_api_key, mock_cache):
    """
    Test that Presidio handles empty messages gracefully.
    """
    test_data = {
        "messages": [],
        "model": "gpt-4",
    }

    # Call the async_pre_call_hook
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    # Should return data unchanged
    assert result == test_data
    print("✓ Empty messages test passed")


@pytest.mark.asyncio
async def test_no_messages_field(presidio_guardrail, mock_user_api_key, mock_cache):
    """
    Test that Presidio handles missing messages field gracefully.
    """
    test_data = {
        "model": "gpt-4",
        "prompt": "This is a completion request",
    }

    # Call the async_pre_call_hook
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    # Should return data unchanged
    assert result == test_data
    print("✓ No messages field test passed")


@pytest.mark.asyncio
async def test_logging_hook_multimodal_message_format(presidio_guardrail):
    """
    Test Presidio async_logging_hook with multimodal message format for completion call type.
    This hook is used to mask PII before logging to external services.
    """
    # Prepare kwargs with multimodal message format
    test_kwargs = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "My credit card number is 4111-1111-1111-1111, my email is test@example.com",
                    }
                ],
            }
        ],
        "model": "gpt-4",
    }

    # Mock result
    mock_result = {"choices": [{"message": {"content": "Response"}}]}

    # Mock the check_pii method
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        redacted_text = text
        redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]")
        redacted_text = redacted_text.replace("test@example.com", "[EMAIL]")
        return redacted_text

    presidio_guardrail.check_pii = mock_check_pii

    # Call the async_logging_hook
    result_kwargs, result_response = await presidio_guardrail.async_logging_hook(
        kwargs=test_kwargs,
        result=mock_result,
        call_type="completion",
    )

    # Verify that PII was masked in the kwargs
    assert result_kwargs is not None
    assert "messages" in result_kwargs
    message = result_kwargs["messages"][0]
    content_item = message["content"][0]

    assert "[CREDIT_CARD]" in content_item["text"]
    assert "[EMAIL]" in content_item["text"]
    assert "4111-1111-1111-1111" not in content_item["text"]
    assert "test@example.com" not in content_item["text"]

    print("✓ Logging hook multimodal message format test passed")


@pytest.mark.asyncio
async def test_logging_hook_multiple_content_items(presidio_guardrail):
    """
    Test Presidio async_logging_hook with multiple content items in a single message.
    """
    test_kwargs = {
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "My credit card is 4111-1111-1111-1111",
                    },
                    {
                        "type": "text",
                        "text": "My email is test@example.com",
                    },
                ],
            }
        ],
        "model": "gpt-4",
    }

    mock_result = {"choices": [{"message": {"content": "Response"}}]}

    # Mock the check_pii method
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        redacted_text = text
        redacted_text = redacted_text.replace("4111-1111-1111-1111", "[CREDIT_CARD]")
        redacted_text = redacted_text.replace("test@example.com", "[EMAIL]")
        return redacted_text

    presidio_guardrail.check_pii = mock_check_pii

    # Call the async_logging_hook
    result_kwargs, result_response = await presidio_guardrail.async_logging_hook(
        kwargs=test_kwargs,
        result=mock_result,
        call_type="completion",
    )

    # Verify both content items were processed
    message = result_kwargs["messages"][0]
    content_items = message["content"]

    assert len(content_items) == 2
    assert "[CREDIT_CARD]" in content_items[0]["text"]
    assert "[EMAIL]" in content_items[1]["text"]

    print("✓ Logging hook multiple content items test passed")


@pytest.mark.asyncio
async def test_presidio_sets_guardrail_information_in_request_data():
    """
    Test that Presidio populates guardrail information into request_data metadata.

    This validates that add_standard_logging_guardrail_information_to_request_data
    correctly sets the guardrail information that will be used for logging.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        guardrail_name="test_presidio",
        output_parse_pii=True,
        mock_testing=True,
    )

    request_data = {
        "messages": [{"role": "user", "content": "Test"}],
        "model": "gpt-4o",
        "metadata": {},
    }

    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        assert request_data is not None

        presidio.add_standard_logging_guardrail_information_to_request_data(
            guardrail_provider="presidio",
            guardrail_json_response=[],
            request_data=request_data,
            guardrail_status="success",
            start_time=1234567890.0,
            end_time=1234567891.0,
            duration=1.0,
            masked_entity_count={"EMAIL_ADDRESS": 1, "PERSON": 1},
        )

        return text

    with patch.object(presidio, "check_pii", mock_check_pii):
        await presidio.apply_guardrail(
            inputs={"texts": ["Test message"]},
            request_data=request_data,
            input_type="request",
        )

    assert "metadata" in request_data
    assert "standard_logging_guardrail_information" in request_data["metadata"]

    guardrail_info_list = request_data["metadata"][
        "standard_logging_guardrail_information"
    ]
    assert isinstance(guardrail_info_list, list)
    assert len(guardrail_info_list) > 0

    guardrail_info = guardrail_info_list[0]
    assert "masked_entity_count" in guardrail_info
    assert guardrail_info["masked_entity_count"]["EMAIL_ADDRESS"] == 1
    assert guardrail_info["masked_entity_count"]["PERSON"] == 1

    print("✓ Presidio sets guardrail_information in request_data")


@pytest.mark.asyncio
async def test_request_data_flows_to_apply_guardrail():
    """
    Test that request_data is correctly passed to apply_guardrail method.

    This validates the fix where guardrail translation handler passes data
    as request_data to apply_guardrail so guardrails can store metadata for logging.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        guardrail_name="test_presidio",
        output_parse_pii=True,
        mock_testing=True,
    )

    request_data = {
        "messages": [{"role": "user", "content": "Test message"}],
        "model": "gpt-4o",
        "metadata": {},
    }

    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        assert request_data is not None, "request_data should be passed to check_pii"
        assert "metadata" in request_data, "request_data should have metadata"

        request_data.setdefault("metadata", {})
        request_data["metadata"]["test_flag"] = "passed_correctly"

        return text

    with patch.object(presidio, "check_pii", mock_check_pii):
        await presidio.apply_guardrail(
            inputs={"texts": ["Test message"]},
            request_data=request_data,
            input_type="request",
        )

        assert "metadata" in request_data
        assert request_data["metadata"].get("test_flag") == "passed_correctly"

    print("✓ request_data correctly passed to apply_guardrail")


@pytest.mark.asyncio
async def test_output_masking_apply_to_output_only(mock_user_api_key):
    """
    Ensure output masking runs when apply_to_output is enabled.
    """

    presidio = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        apply_to_output=True,
        pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.MASK},
    )

    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        return text.replace("4111-1111-1111-1111", "[CREDIT_CARD]")

    presidio.check_pii = mock_check_pii

    response = ModelResponse(
        id="1",
        object="chat.completion",
        created=0,
        model="gpt-test",
        choices=[
            Choices(
                message=Message(
                    role="assistant",
                    content="Card is 4111-1111-1111-1111",
                ),
                index=0,
                finish_reason="stop",
            )
        ],
    )

    result = await presidio.async_post_call_success_hook(
        data={},
        user_api_key_dict=mock_user_api_key,
        response=response,
    )

    assert "[CREDIT_CARD]" in result.choices[0].message.content
    assert "4111-1111-1111-1111" not in result.choices[0].message.content


@pytest.mark.asyncio
async def test_presidio_filter_scope_initializer(monkeypatch):
    """
    Ensure initializer respects presidio_filter_scope for input/output/both.
    """

    created = []

    class DummyGuardrail:
        def __init__(self, apply_to_output: bool = False, event_hook=None, **kwargs):
            self.apply_to_output = apply_to_output
            self.event_hook = event_hook
            created.append(self)

        def update_in_memory_litellm_params(self, litellm_params):
            pass

    class DummyManager:
        def __init__(self):
            self.added = []

        def add_litellm_callback(self, cb):
            self.added.append(cb)

    mgr = DummyManager()
    monkeypatch.setattr(litellm, "logging_callback_manager", mgr, raising=False)
    import litellm.proxy.guardrails.guardrail_hooks.presidio as presidio_mod
    import litellm.proxy.guardrails.guardrail_initializers as gi

    monkeypatch.setattr(
        presidio_mod, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False
    )
    monkeypatch.setattr(
        gi, "_OPTIONAL_PresidioPIIMasking", DummyGuardrail, raising=False
    )

    # input-only
    created.clear()
    from litellm.proxy.guardrails.guardrail_initializers import initialize_presidio

    params_input = LitellmParams(
        guardrail="presidio", mode="pre_call", presidio_filter_scope="input"
    )
    guardrail_dict = {"guardrail_name": "g1"}
    cb = initialize_presidio(params_input, guardrail_dict)
    assert cb is created[0]
    assert created[0].apply_to_output is False

    # output-only
    created.clear()
    params_output = LitellmParams(
        guardrail="presidio", mode="pre_call", presidio_filter_scope="output"
    )
    cb = initialize_presidio(params_output, guardrail_dict)
    assert len(created) == 1
    assert created[0].apply_to_output is True

    # both -> expect two callbacks (input + output)
    created.clear()
    params_both = LitellmParams(
        guardrail="presidio", mode="pre_call", presidio_filter_scope="both"
    )
    cb = initialize_presidio(params_both, guardrail_dict)
    assert len(created) == 2
    assert any(not c.apply_to_output for c in created)
    assert any(c.apply_to_output for c in created)


@pytest.mark.asyncio
async def test_empty_content_handling(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test that Presidio handles empty content gracefully.

    This is common in tool/function calling where assistant messages have
    empty content but include tool_calls.

    Bug fix: Previously crashed with:
    TypeError: argument after ** must be a mapping, not str
    """
    test_data = {
        "messages": [
            {"role": "user", "content": "What is 2+2?"},
            {
                "role": "assistant",
                "content": "",  # Empty content - common in tool calls
                "tool_calls": [
                    {
                        "id": "call_123",
                        "type": "function",
                        "function": {
                            "name": "calculator",
                            "arguments": '{"a":2,"b":2}',
                        },
                    }
                ],
            },
            {"role": "tool", "tool_call_id": "call_123", "content": "4"},
        ],
        "model": "gpt-4",
    }

    # Mock check_pii to simulate PII processing without needing Presidio API
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        # Empty text returns as-is (this is what our fix ensures)
        return text

    presidio_guardrail.check_pii = mock_check_pii

    # This should not raise an exception
    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    assert result is not None
    assert "messages" in result
    # Verify messages are preserved
    assert len(result["messages"]) == 3

    print("✓ Empty content handling test passed")


@pytest.mark.asyncio
async def test_whitespace_only_content(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test that Presidio handles whitespace-only content gracefully.

    Whitespace-only content should be treated the same as empty content.
    """
    test_data = {
        "messages": [
            {"role": "user", "content": "   "},  # Whitespace only
            {"role": "assistant", "content": "\n\t  "},  # Tabs and newlines
            {"role": "user", "content": "Real question here"},
        ],
        "model": "gpt-4",
    }

    # Mock check_pii to simulate PII processing
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        return text

    presidio_guardrail.check_pii = mock_check_pii

    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    assert result is not None
    assert len(result["messages"]) == 3

    print("✓ Whitespace-only content test passed")


@pytest.mark.asyncio
async def test_analyze_text_with_empty_string():
    """
    Test analyze_text method directly with empty string.

    Should return empty list without making API call to Presidio.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test:5002/",
        presidio_anonymizer_api_base="http://test:5001/",
        output_parse_pii=False,
    )

    # Test with empty string - should return immediately without API call
    result = await presidio.analyze_text(
        text="",
        presidio_config=None,
        request_data={},
    )
    assert result == [], "Empty text should return empty list"

    # Test with whitespace only - should return immediately
    result = await presidio.analyze_text(
        text="   \n\t   ",
        presidio_config=None,
        request_data={},
    )
    assert result == [], "Whitespace-only text should return empty list"

    print("✓ analyze_text empty string test passed")


@pytest.mark.asyncio
async def test_analyze_text_error_dict_handling():
    """
    Test that analyze_text handles error dict responses from Presidio API.

    When Presidio returns {'error': 'No text provided'}, should handle gracefully
    instead of crashing with TypeError.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://mock-presidio:5002/",
        presidio_anonymizer_api_base="http://mock-presidio:5001/",
        output_parse_pii=False,
    )

    with patch.object(
        presidio,
        "_get_session_iterator",
        _make_mock_session_iterator({"error": "No text provided"}),
    ):
        result = await presidio.analyze_text(
            text="some text",
            presidio_config=None,
            request_data={},
        )
    assert result == [], "Error dict should be handled gracefully"

    print("✓ analyze_text error dict handling test passed")


@pytest.mark.asyncio
async def test_analyze_text_string_response_handling():
    """
    Test that analyze_text handles string responses from Presidio API.

    When Presidio returns a string (e.g. error message from websearch/hosted models),
    should handle gracefully instead of crashing with TypeError about mapping vs str.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://mock-presidio:5002/",
        presidio_anonymizer_api_base="http://mock-presidio:5001/",
        output_parse_pii=False,
    )

    with patch.object(
        presidio,
        "_get_session_iterator",
        _make_mock_session_iterator("Internal Server Error"),
    ):
        result = await presidio.analyze_text(
            text="some text",
            presidio_config=None,
            request_data={},
        )
    assert result == [], "String response should be handled gracefully"


@pytest.mark.asyncio
async def test_analyze_text_invalid_response_raises_when_block_configured():
    """
    When pii_entities_config has BLOCK and Presidio returns invalid response,
    should raise GuardrailRaisedException (fail-closed) rather than silently allowing content.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://mock-presidio:5002/",
        presidio_anonymizer_api_base="http://mock-presidio:5001/",
        output_parse_pii=False,
        pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.BLOCK},
    )

    with patch.object(
        presidio,
        "_get_session_iterator",
        _make_mock_session_iterator("Internal Server Error"),
    ):
        with pytest.raises(GuardrailRaisedException) as exc_info:
            await presidio.analyze_text(
                text="some text",
                presidio_config=None,
                request_data={},
            )
    assert "BLOCK" in str(exc_info.value) or "Presidio" in str(exc_info.value)


@pytest.mark.asyncio
async def test_analyze_text_invalid_response_raises_when_mask_configured():
    """
    When pii_entities_config has MASK and Presidio returns invalid response,
    should raise GuardrailRaisedException (fail-closed) because PII masking is expected.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://mock-presidio:5002/",
        presidio_anonymizer_api_base="http://mock-presidio:5001/",
        output_parse_pii=False,
        pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.MASK},
    )

    with patch.object(
        presidio,
        "_get_session_iterator",
        _make_mock_session_iterator("Internal Server Error"),
    ):
        with pytest.raises(GuardrailRaisedException) as exc_info:
            await presidio.analyze_text(
                text="some text",
                presidio_config=None,
                request_data={},
            )
    assert "PII protection is configured" in str(exc_info.value)


@pytest.mark.asyncio
async def test_analyze_text_list_with_non_dict_items():
    """
    Test that analyze_text skips non-dict items in the result list.

    When Presidio returns a list containing strings (malformed response),
    should skip invalid items and return parsed valid ones.
    """
    presidio = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://mock-presidio:5002/",
        presidio_anonymizer_api_base="http://mock-presidio:5001/",
        output_parse_pii=False,
    )

    json_response = [
        {"entity_type": "PERSON", "start": 0, "end": 5, "score": 0.9},
        "invalid_string_item",
        {"entity_type": "EMAIL", "start": 10, "end": 25, "score": 0.85},
    ]
    with patch.object(
        presidio, "_get_session_iterator", _make_mock_session_iterator(json_response)
    ):
        result = await presidio.analyze_text(
            text="some text",
            presidio_config=None,
            request_data={},
        )
    assert len(result) == 2, "Should parse 2 valid dict items and skip the string"
    assert result[0].get("entity_type") == "PERSON"
    assert result[1].get("entity_type") == "EMAIL"


@pytest.mark.asyncio
async def test_tool_calling_complete_scenario(
    presidio_guardrail, mock_user_api_key, mock_cache
):
    """
    Test complete tool calling scenario with PII in user message.

    This tests the real-world scenario where:
    1. User provides a query with PII
    2. Assistant responds with empty content + tool_calls
    3. Tool provides response
    4. Assistant provides final answer
    """
    test_data = {
        "messages": [
            {
                "role": "user",
                "content": "My email is john.doe@example.com. Can you look up my account?",
            },
            {
                "role": "assistant",
                "content": "",  # Empty - tool call
                "tool_calls": [
                    {
                        "id": "call_abc",
                        "type": "function",
                        "function": {"name": "lookup_account", "arguments": "{}"},
                    }
                ],
            },
            {"role": "tool", "tool_call_id": "call_abc", "content": "Account found"},
            {"role": "assistant", "content": "I found your account information."},
        ],
        "model": "gpt-4",
    }

    # Mock check_pii to simulate PII masking
    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        if "john.doe@example.com" in text:
            return text.replace("john.doe@example.com", "[EMAIL]")
        return text

    presidio_guardrail.check_pii = mock_check_pii

    result = await presidio_guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    assert result is not None
    # Verify PII was masked in user message
    assert "[EMAIL]" in result["messages"][0]["content"]
    assert "john.doe@example.com" not in result["messages"][0]["content"]
    # Verify other messages preserved
    assert len(result["messages"]) == 4

    print("✓ Tool calling complete scenario test passed")


def test_filter_drops_low_score_detection():
    """
    Detections below the configured score threshold should be removed.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.8},
    )
    analyze_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4}
    ]

    filtered = guardrail.filter_analyze_results_by_score(analyze_results)
    assert filtered == []


def test_filter_preserves_high_score_detection():
    """
    Detections meeting the score threshold should be preserved.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.8},
    )
    analyze_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.9, "start": 0, "end": 4}
    ]

    filtered = guardrail.filter_analyze_results_by_score(analyze_results)
    assert len(filtered) == 1
    assert filtered[0]["entity_type"] == PiiEntityType.CREDIT_CARD


def test_no_thresholds_returns_all():
    """
    With no thresholds configured, all detections are kept.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True)
    analyze_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.1, "start": 0, "end": 4},
        {
            "entity_type": PiiEntityType.EMAIL_ADDRESS,
            "score": 0.2,
            "start": 5,
            "end": 9,
        },
    ]

    filtered = guardrail.filter_analyze_results_by_score(analyze_results)
    assert len(filtered) == 2


def test_entity_specific_threshold_only_applies_to_that_entity():
    """
    Entity-specific thresholds do not affect other entity types.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.8},
    )
    analyze_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4},
        {
            "entity_type": PiiEntityType.EMAIL_ADDRESS,
            "score": 0.1,
            "start": 5,
            "end": 9,
        },
    ]

    filtered = guardrail.filter_analyze_results_by_score(analyze_results)
    # CREDIT_CARD is filtered, EMAIL_ADDRESS is kept because no threshold
    assert len(filtered) == 1
    assert filtered[0]["entity_type"] == PiiEntityType.EMAIL_ADDRESS


def test_filter_uses_default_all_threshold():
    """
    Default ALL threshold applies to any entity without a specific override.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_score_thresholds={"ALL": 0.75},
    )
    analyze_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4},
        {
            "entity_type": PiiEntityType.EMAIL_ADDRESS,
            "score": 0.8,
            "start": 5,
            "end": 9,
        },
    ]

    filtered = guardrail.filter_analyze_results_by_score(analyze_results)
    assert len(filtered) == 1
    assert filtered[0]["entity_type"] == PiiEntityType.EMAIL_ADDRESS


def test_entity_specific_overrides_default_threshold():
    """
    Entity-specific threshold should override the ALL default.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_score_thresholds={
            "ALL": 0.8,
            PiiEntityType.CREDIT_CARD: 0.6,
        },
    )
    analyze_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.65, "start": 0, "end": 4},
        {
            "entity_type": PiiEntityType.EMAIL_ADDRESS,
            "score": 0.75,
            "start": 5,
            "end": 9,
        },
    ]

    filtered = guardrail.filter_analyze_results_by_score(analyze_results)
    # CREDIT_CARD passes due to override, EMAIL_ADDRESS dropped by ALL threshold
    assert len(filtered) == 1
    assert filtered[0]["entity_type"] == PiiEntityType.CREDIT_CARD


@pytest.mark.asyncio
async def test_anonymize_skips_when_no_detections_after_filter():
    """
    When all detections are filtered out, anonymize_text should return the original text.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.8},
    )
    masked_entity_count = {}
    text = "4111"

    filtered = guardrail.filter_analyze_results_by_score(
        [{"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4}]
    )

    result = await guardrail.anonymize_text(
        text=text,
        analyze_results=filtered,
        output_parse_pii=False,
        masked_entity_count=masked_entity_count,
    )

    assert result == text
    assert masked_entity_count == {}


def test_blocking_respects_threshold_filter():
    """
    Entities filtered out by score should not trigger blocking, but high-score detections should.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        pii_entities_config={PiiEntityType.CREDIT_CARD: PiiAction.BLOCK},
        presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.9},
    )

    low_score_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.7, "start": 0, "end": 4}
    ]
    filtered = guardrail.filter_analyze_results_by_score(low_score_results)
    guardrail.raise_exception_if_blocked_entities_detected(filtered)

    high_score_results = [
        {"entity_type": PiiEntityType.CREDIT_CARD, "score": 0.95, "start": 0, "end": 4}
    ]
    filtered_high = guardrail.filter_analyze_results_by_score(high_score_results)
    with pytest.raises(Exception):
        guardrail.raise_exception_if_blocked_entities_detected(filtered_high)


def test_update_in_memory_applies_score_thresholds():
    """
    update_in_memory_litellm_params should refresh score thresholds.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(mock_testing=True)
    assert guardrail.presidio_score_thresholds == {}

    params = LitellmParams(
        guardrail="presidio",
        mode="pre_call",
        presidio_score_thresholds={PiiEntityType.CREDIT_CARD: 0.85},
    )
    guardrail.update_in_memory_litellm_params(params)

    assert guardrail.presidio_score_thresholds == {PiiEntityType.CREDIT_CARD: 0.85}


@pytest.mark.asyncio
async def test_get_session_iterator_thread_safety(presidio_guardrail):
    """
    Test that _get_session_iterator yields:
    1. The shared session when in the main thread.
    2. A loop-bound cached session when in a background thread (reused per loop for efficiency).
    """
    import threading

    import aiohttp

    # 1. Main Thread Case
    # We are in the "main thread" relative to the guardrail initialization
    async with presidio_guardrail._get_session_iterator() as session:
        assert isinstance(session, aiohttp.ClientSession)
        assert session is presidio_guardrail._http_session
        shared_session_id = id(session)

    # 2. Background Thread Case
    # Define a helper function to run in a thread
    def thread_target(loop, result_future):
        async def run_in_loop():
            # This runs in the thread's loop
            async with presidio_guardrail._get_session_iterator() as session:
                return session, id(session)

        try:
            # Create a new loop for this thread to run async code
            new_loop = asyncio.new_event_loop()
            asyncio.set_event_loop(new_loop)
            session_obj, session_id = new_loop.run_until_complete(run_in_loop())
            result_future.set_result((session_obj, session_id))
            new_loop.close()
        except Exception as e:
            result_future.set_exception(e)

    # Run the background thread test
    bg_future = asyncio.Future()
    t = threading.Thread(
        target=thread_target, args=(asyncio.get_running_loop(), bg_future)
    )
    t.start()
    t.join()

    bg_session, bg_session_id = await bg_future

    # Assertions
    # The background session should be DIFFERENT from the shared session
    assert bg_session_id != shared_session_id
    # The shared session should still be open (not closed by the background thread)
    assert not presidio_guardrail._http_session.closed
    # The background session should be cached in _loop_sessions and remain open for reuse
    # (Changed behavior: no longer closes immediately, cached per loop for efficiency)
    assert not bg_session.closed, "Background session should remain open for reuse"

    print("✓ Session iterator thread safety test passed")


from litellm.types.utils import ModelResponseStream


@pytest.mark.asyncio
async def test_streaming_with_bytes_chunks_does_not_crash(mock_user_api_key):
    """
    Regression test: async_post_call_streaming_iterator_hook should
    gracefully handle raw bytes in the stream instead of crashing with
    'bytes' object has no attribute 'id'.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        apply_to_output=True,
        mock_redacted_text={"text": "redacted"},
    )

    async def mock_stream():
        yield b'data: {"id":"chatcmpl-1"}\n\n'  # raw bytes
        yield ModelResponseStream(
            id="chatcmpl-1",
            choices=[],
            created=1,
            model="gpt-4",
            object="chat.completion.chunk",
            system_fingerprint=None,
        )  # proper chunk

    chunks = []
    async for chunk in guardrail.async_post_call_streaming_iterator_hook(
        user_api_key_dict=mock_user_api_key,
        response=mock_stream(),
        request_data={},
    ):
        chunks.append(chunk)

    # Should not crash, should produce at least one valid chunk
    assert len(chunks) >= 1


def test_entity_deny_list_filters_detections():
    """
    Verify presidio_entities_deny_list removes matching entity types.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_entities_deny_list=["US_DRIVER_LICENSE"],
    )

    results = [
        {"entity_type": "US_DRIVER_LICENSE", "start": 0, "end": 2, "score": 0.6},
        {"entity_type": "CREDIT_CARD", "start": 10, "end": 26, "score": 0.95},
    ]

    filtered = guardrail.filter_analyze_results_by_score(results)

    assert len(filtered) == 1
    assert filtered[0]["entity_type"] == "CREDIT_CARD"


def test_deny_list_and_score_threshold_combined():
    """
    Verify deny list + score threshold work together correctly.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        presidio_entities_deny_list=["US_DRIVER_LICENSE"],
        presidio_score_thresholds={"ALL": 0.8},
    )

    results = [
        {"entity_type": "US_DRIVER_LICENSE", "start": 0, "end": 2, "score": 0.95},
        {"entity_type": "CREDIT_CARD", "start": 10, "end": 26, "score": 0.6},
        {"entity_type": "EMAIL_ADDRESS", "start": 30, "end": 50, "score": 0.9},
    ]

    filtered = guardrail.filter_analyze_results_by_score(results)

    # US_DRIVER_LICENSE excluded by deny list (even though score > 0.8)
    # CREDIT_CARD excluded by score threshold (0.6 < 0.8)
    # EMAIL_ADDRESS passes both filters
    assert len(filtered) == 1
    assert filtered[0]["entity_type"] == "EMAIL_ADDRESS"


@pytest.mark.asyncio
async def test_analyze_text_non_json_content_type_fail_closed():
    """
    Test that analyze_text raises GuardrailRaisedException when Presidio health
    endpoint returns text/html and fail-closed is enabled.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        pii_entities_config={"PERSON": PiiAction.BLOCK},
        mock_testing=False,
    )

    mock_iterator = _make_mock_session_iterator(
        json_response=None,
        status=200,
        content_type="text/html; charset=utf-8",
        text_response="Presidio Analyzer service is up.",
    )

    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        with pytest.raises(GuardrailRaisedException) as exc_info:
            await guardrail.analyze_text(
                text="Hello world",
                presidio_config=None,
                request_data={},
            )
        assert "expected application/json Content-Type" in str(exc_info.value)
        assert "text/html" in str(exc_info.value)


@pytest.mark.asyncio
async def test_analyze_text_non_json_content_type_fail_open():
    """
    Test that analyze_text returns empty list when Presidio returns text/html
    and fail-closed is NOT enabled.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        mock_testing=False,
    )

    mock_iterator = _make_mock_session_iterator(
        json_response=None,
        status=200,
        content_type="text/html; charset=utf-8",
        text_response="Presidio Analyzer service is up.",
    )

    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        results = await guardrail.analyze_text(
            text="Hello world",
            presidio_config=None,
            request_data={},
        )
        assert results == []


@pytest.mark.asyncio
async def test_analyze_text_http_error_status():
    """
    Test that analyze_text handles 5xx HTTP errors properly.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        pii_entities_config={"PERSON": PiiAction.BLOCK},
        mock_testing=False,
    )

    mock_iterator = _make_mock_session_iterator(
        json_response=None,
        status=500,
        content_type="text/plain",
        text_response="Internal Server Error",
    )

    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        with pytest.raises(GuardrailRaisedException) as exc_info:
            await guardrail.analyze_text(
                text="Hello world",
                presidio_config=None,
                request_data={},
            )
        assert "HTTP 500" in str(exc_info.value)


@pytest.mark.asyncio
async def test_anonymize_text_non_json_content_type():
    """
    Test that anonymize_text raises Exception for non-JSON responses.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        mock_testing=False,
    )

    mock_iterator = _make_mock_session_iterator(
        json_response=None,
        status=200,
        content_type="text/html",
        text_response="Presidio Anonymizer service is up.",
    )

    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        with pytest.raises(
            Exception, match="Presidio anonymizer returned non-JSON Content-Type"
        ):
            await guardrail.anonymize_text(
                text="Hello world",
                analyze_results=[{"start": 0, "end": 5, "entity_type": "PERSON"}],
                output_parse_pii=False,
                masked_entity_count={},
            )


@pytest.mark.asyncio
async def test_anonymize_text_http_error_status():
    """
    Test that anonymize_text raises Exception on HTTP error.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        mock_testing=False,
    )

    mock_iterator = _make_mock_session_iterator(
        json_response=None,
        status=502,
        content_type="text/plain",
        text_response="Bad Gateway",
    )

    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        with pytest.raises(Exception, match="Presidio anonymizer returned HTTP 502"):
            await guardrail.anonymize_text(
                text="Hello world",
                analyze_results=[{"start": 0, "end": 5, "entity_type": "PERSON"}],
                output_parse_pii=False,
                masked_entity_count={},
            )


@pytest.mark.asyncio
async def test_pii_tokens_stored_in_metadata_not_top_level(presidio_guardrail):
    """
    Regression test: pii_tokens must be stored in data['metadata']['pii_tokens'],
    NOT in data['pii_tokens']. Storing at the top level leaks the field to LLM
    providers like Anthropic, which reject unknown fields with
    'pii_tokens: Extra inputs are not permitted'.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
        pii_entities_config={
            PiiEntityType.PERSON: PiiAction.MASK,
            PiiEntityType.PHONE_NUMBER: PiiAction.MASK,
        },
    )

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
    mock_cache = DualCache()

    test_data = {
        "messages": [
            {"role": "user", "content": "My name is John and my phone is 555-123-4567"}
        ],
        "model": "claude-haiku-4-5-20251001",
        "metadata": {},
    }

    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        # Simulate PII masking with token storage (mimics real anonymize_text behavior)
        if request_data is not None and output_parse_pii:
            if "metadata" not in request_data:
                request_data["metadata"] = {}
            if "pii_tokens" not in request_data["metadata"]:
                request_data["metadata"]["pii_tokens"] = {}
            pii_tokens = request_data["metadata"]["pii_tokens"]
            seq = len(pii_tokens) + 1
            token = f"<PERSON_{seq}>"
            pii_tokens[token] = "John"
            text = text.replace("John", token)
        return text

    guardrail.check_pii = mock_check_pii

    result = await guardrail.async_pre_call_hook(
        user_api_key_dict=mock_user_api_key,
        cache=mock_cache,
        data=test_data,
        call_type="completion",
    )

    # pii_tokens must NOT be at the top level of data (would leak to providers)
    assert "pii_tokens" not in result, (
        "pii_tokens must not be a top-level key in request data — "
        "it would leak to LLM providers and cause 'Extra inputs are not permitted' errors"
    )

    # pii_tokens must be inside metadata (safe from provider leakage)
    assert "metadata" in result
    assert "pii_tokens" in result["metadata"]
    assert len(result["metadata"]["pii_tokens"]) > 0


@pytest.mark.asyncio
async def test_pii_tokens_in_metadata_used_for_unmasking():
    """
    Regression test: _process_response_for_pii must read pii_tokens from
    data['metadata']['pii_tokens'] and correctly unmask the response.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
    )

    token_key = "<PERSON_1>"
    request_data = {
        "model": "claude-haiku-4-5-20251001",
        "metadata": {"pii_tokens": {token_key: "John"}},
    }

    response = ModelResponse(
        choices=[
            Choices(
                message=Message(
                    role="assistant",
                    content=f"Hello {token_key}, how can I help you?",
                ),
                index=0,
                finish_reason="stop",
            )
        ]
    )

    await guardrail._process_response_for_pii(
        response=response,
        request_data=request_data,
        mode="unmask",
    )

    assert response.choices[0].message.content == "Hello John, how can I help you?"


@pytest.mark.parametrize(
    "initial_hook",
    ["pre_call", "during_call", "pre_mcp_call"],
)
def test_event_hook_auto_expansion_for_all_string_hooks(initial_hook):
    """
    Regression test: when output_parse_pii is True, the guardrail must add
    'post_call' to event_hook regardless of the initial string hook value,
    not just when it's 'pre_call'.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
        event_hook=initial_hook,
    )
    assert isinstance(guardrail.event_hook, list)
    assert initial_hook in guardrail.event_hook
    assert "post_call" in guardrail.event_hook


def test_event_hook_no_expansion_when_already_post_call():
    """post_call alone should stay as-is — no expansion needed."""
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
        event_hook="post_call",
    )
    # Should remain a string "post_call", not expanded to a list
    assert guardrail.event_hook == "post_call"


@pytest.mark.asyncio
async def test_metadata_none_does_not_crash():
    """
    Regression test: if metadata is explicitly None in request_data,
    the guardrail must not crash with TypeError on the write or read path.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
    )

    token_key = "<PERSON_1>"
    # metadata explicitly None — must not crash
    request_data = {
        "model": "gpt-3.5-turbo",
        "metadata": None,
    }

    response = ModelResponse(
        choices=[
            Choices(
                message=Message(
                    role="assistant",
                    content=f"Hello {token_key}, how can I help you?",
                ),
                index=0,
                finish_reason="stop",
            )
        ]
    )

    # Should not raise TypeError
    await guardrail._process_response_for_pii(
        response=response,
        request_data=request_data,
        mode="unmask",
    )

    # No pii_tokens to unmask, so content stays as-is
    assert (
        response.choices[0].message.content == f"Hello {token_key}, how can I help you?"
    )


# ---------------------------------------------------------------------------
# Tests for sequential-numbered token unmasking in _unmask_pii_text
# ---------------------------------------------------------------------------


def test_unmask_exact_match_with_sequential_tokens():
    """
    Normal unmasking: LLM echoes numbered tokens verbatim → original PII restored.
    """
    from litellm.proxy.guardrails.guardrail_hooks.presidio import (
        _OPTIONAL_PresidioPIIMasking,
    )

    pii_tokens = {
        "<PERSON_1>": "John Smith",
        "<PHONE_NUMBER_1>": "555-123-4567",
    }
    text = "Hello <PERSON_1>, your number is <PHONE_NUMBER_1>."
    result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens)
    assert result == "Hello John Smith, your number is 555-123-4567."


def test_unmask_multiple_same_entity_type():
    """
    Two phone numbers get distinct numbered tokens and unmask correctly.
    """
    from litellm.proxy.guardrails.guardrail_hooks.presidio import (
        _OPTIONAL_PresidioPIIMasking,
    )

    pii_tokens = {
        "<PHONE_NUMBER_1>": "555-111-0000",
        "<PHONE_NUMBER_2>": "555-222-0000",
    }
    text = "Call <PHONE_NUMBER_1> or <PHONE_NUMBER_2>."
    result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens)
    assert result == "Call 555-111-0000 or 555-222-0000."


def test_unmask_graceful_degradation():
    """
    If the LLM doesn't echo the token back, the numbered label stays
    in the output — clean and readable, not garbage hex.
    """
    from litellm.proxy.guardrails.guardrail_hooks.presidio import (
        _OPTIONAL_PresidioPIIMasking,
    )

    pii_tokens = {
        "<PERSON_1>": "John",
    }
    # LLM paraphrased instead of echoing the token
    text = "I see you provided a name."
    result = _OPTIONAL_PresidioPIIMasking._unmask_pii_text(text, pii_tokens)
    # No change — no garbage, just clean text
    assert result == text


# ---------------------------------------------------------------------------
# Fix 1: Position bug — reverse sort + original text coordinates
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_anonymize_text_multiple_items_position_correctness():
    """
    Regression test: when multiple PII items exist, coordinates reference the
    ORIGINAL text. Processing in reverse order prevents coordinate drift.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        mock_testing=False,
    )

    # "Call John at 555-123-4567"
    #   "John" at [5:9], "555-123-4567" at [13:25]
    anonymizer_response = {
        "text": "Call <PERSON> at <PHONE_NUMBER>",
        "items": [
            {
                "start": 5,
                "end": 9,
                "entity_type": "PERSON",
                "text": "<PERSON>",
                "operator": "replace",
            },
            {
                "start": 13,
                "end": 25,
                "entity_type": "PHONE_NUMBER",
                "text": "<PHONE_NUMBER>",
                "operator": "replace",
            },
        ],
    }

    mock_iterator = _make_mock_session_iterator(anonymizer_response)

    request_data = {"metadata": {}}
    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        result = await guardrail.anonymize_text(
            text="Call John at 555-123-4567",
            analyze_results=[
                {"start": 5, "end": 9, "entity_type": "PERSON", "score": 0.9},
                {"start": 13, "end": 25, "entity_type": "PHONE_NUMBER", "score": 0.95},
            ],
            output_parse_pii=True,
            masked_entity_count={},
            request_data=request_data,
        )

    pii_tokens = request_data["metadata"]["pii_tokens"]

    # Verify tokens captured the correct ORIGINAL text values
    person_token = [k for k in pii_tokens if "PERSON" in k][0]
    phone_token = [k for k in pii_tokens if "PHONE" in k][0]
    assert pii_tokens[person_token] == "John"
    assert pii_tokens[phone_token] == "555-123-4567"

    # Verify both PII values are masked in the result
    assert "John" not in result
    assert "555-123-4567" not in result


# ---------------------------------------------------------------------------
# Fix 2: Anthropic native dict response handling
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_anthropic_native_response_unmasking():
    """
    Anthropic native dict responses (type='message') should be unmasked
    when output_parse_pii is enabled.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
    )

    request_data = {
        "model": "claude-3-haiku",
        "metadata": {
            "pii_tokens": {
                "<PERSON_1>": "John Smith",
                "<PHONE_NUMBER_1>": "555-123-4567",
            }
        },
    }

    anthropic_response = {
        "type": "message",
        "id": "msg_123",
        "model": "claude-3-haiku",
        "role": "assistant",
        "content": [
            {
                "type": "text",
                "text": "Hello <PERSON_1>, your number is <PHONE_NUMBER_1>.",
            }
        ],
        "stop_reason": "end_turn",
        "usage": {"input_tokens": 10, "output_tokens": 20},
    }

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")

    result = await guardrail.async_post_call_success_hook(
        data=request_data,
        user_api_key_dict=mock_user_api_key,
        response=anthropic_response,
    )

    assert result["content"][0]["text"] == (
        "Hello John Smith, your number is 555-123-4567."
    )


@pytest.mark.asyncio
async def test_anthropic_native_response_masking():
    """
    Anthropic native dict responses should be masked when
    apply_to_output is enabled.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        apply_to_output=True,
    )

    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        return text.replace("John Smith", "[PERSON]").replace("555-123-4567", "[PHONE]")

    guardrail.check_pii = mock_check_pii

    anthropic_response = {
        "type": "message",
        "id": "msg_123",
        "model": "claude-3-haiku",
        "role": "assistant",
        "content": [{"type": "text", "text": "Hello John Smith, call 555-123-4567."}],
        "stop_reason": "end_turn",
        "usage": {"input_tokens": 10, "output_tokens": 20},
    }

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")

    result = await guardrail.async_post_call_success_hook(
        data={},
        user_api_key_dict=mock_user_api_key,
        response=anthropic_response,
    )

    assert "[PERSON]" in result["content"][0]["text"]
    assert "[PHONE]" in result["content"][0]["text"]
    assert "John Smith" not in result["content"][0]["text"]


@pytest.mark.asyncio
async def test_anthropic_native_response_non_text_blocks_untouched():
    """
    Non-text blocks (tool_use, thinking) in Anthropic responses
    should be left untouched during unmasking.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
    )

    request_data = {
        "model": "claude-3-haiku",
        "metadata": {"pii_tokens": {"<PERSON_1>": "John"}},
    }

    anthropic_response = {
        "type": "message",
        "id": "msg_123",
        "content": [
            {"type": "text", "text": "Hello <PERSON_1>"},
            {
                "type": "tool_use",
                "id": "call_1",
                "name": "search",
                "input": {"q": "test"},
            },
        ],
        "role": "assistant",
        "stop_reason": "end_turn",
        "usage": {"input_tokens": 10, "output_tokens": 20},
    }

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")

    result = await guardrail.async_post_call_success_hook(
        data=request_data,
        user_api_key_dict=mock_user_api_key,
        response=anthropic_response,
    )

    assert result["content"][0]["text"] == "Hello John"
    assert result["content"][1]["type"] == "tool_use"
    assert result["content"][1]["name"] == "search"


# ---------------------------------------------------------------------------
# Fix 3: Anthropic native SSE streaming — bytes passthrough
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_streaming_bytes_chunks_are_yielded_not_discarded():
    """
    Regression test: bytes chunks (Anthropic native SSE) should be yielded
    through the streaming hook, not silently discarded.
    """

    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        apply_to_output=True,
    )

    byte_chunk = b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n'

    async def mock_stream():
        yield byte_chunk

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
    chunks = []
    async for chunk in guardrail.async_post_call_streaming_iterator_hook(
        user_api_key_dict=mock_user_api_key,
        response=mock_stream(),
        request_data={},
    ):
        chunks.append(chunk)

    assert any(
        isinstance(c, bytes) for c in chunks
    ), "bytes chunks must not be discarded"
    assert byte_chunk in chunks


@pytest.mark.asyncio
async def test_streaming_unmask_path_bytes_passthrough():
    """
    Bytes chunks in the unmasking path should also pass through.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
    )

    byte_chunk = b'data: {"type":"content_block_delta"}\n\n'
    request_data = {
        "metadata": {"pii_tokens": {"<PERSON_1>": "John"}},
    }

    async def mock_stream():
        yield byte_chunk

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
    chunks = []
    async for chunk in guardrail.async_post_call_streaming_iterator_hook(
        user_api_key_dict=mock_user_api_key,
        response=mock_stream(),
        request_data=request_data,
    ):
        chunks.append(chunk)

    assert len(chunks) == 1
    assert chunks[0] == byte_chunk


@pytest.mark.asyncio
async def test_apply_to_output_streaming_unknown_events_passthrough():
    """
    Regression test: /v1/responses-style event objects (neither bytes nor
    ModelResponseStream) must be preserved in order and not dropped.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        apply_to_output=True,
    )

    class FakeResponsesEvent:
        def __init__(self, event_type: str):
            self.type = event_type

    events = [
        FakeResponsesEvent("response.created"),
        FakeResponsesEvent("response.output_text.delta"),
        FakeResponsesEvent("response.completed"),
    ]

    async def mock_stream():
        for event in events:
            yield event

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
    received = []
    async for chunk in guardrail.async_post_call_streaming_iterator_hook(
        user_api_key_dict=mock_user_api_key,
        response=mock_stream(),
        request_data={},
    ):
        received.append(chunk)

    # Preserve exact objects and ordering so clients receive full event lifecycle.
    assert received == events
    assert [e.type for e in received] == [
        "response.created",
        "response.output_text.delta",
        "response.completed",
    ]


@pytest.mark.asyncio
async def test_apply_to_output_streaming_mixed_chunks_flushes_and_warns():
    """
    Regression test for mixed stream shape:
    a buffered ModelResponseStream chunk followed by unknown responses-style
    events should be preserved, and masking skip should be visible via warnings.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        apply_to_output=True,
    )

    class FakeResponsesEvent:
        def __init__(self, event_type: str):
            self.type = event_type

    model_chunk = ModelResponseStream(
        id="chatcmpl-mixed-1",
        choices=[],
        created=1,
        model="gpt-4",
        object="chat.completion.chunk",
        system_fingerprint=None,
    )
    response_completed = FakeResponsesEvent("response.completed")

    async def mock_stream():
        yield model_chunk
        yield response_completed

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")
    received = []
    with patch(
        "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger"
    ) as mock_logger:
        async for chunk in guardrail.async_post_call_streaming_iterator_hook(
            user_api_key_dict=mock_user_api_key,
            response=mock_stream(),
            request_data={},
        ):
            received.append(chunk)

        # Preserve original ordering across mixed stream types.
        assert received == [model_chunk, response_completed]

        # Two warnings are expected:
        # 1) mixed stream detected + unmasked flush
        # 2) passthrough mode skipped output masking
        assert mock_logger.warning.call_count == 2
        warning_messages = [call.args[0] for call in mock_logger.warning.call_args_list]
        assert any("mixed stream detected" in msg for msg in warning_messages)
        assert any("unknown event objects" in msg for msg in warning_messages)


# ---------------------------------------------------------------------------
# Fix 4: apply_guardrail unmask path for input_type="response"
# ---------------------------------------------------------------------------


@pytest.mark.asyncio
async def test_apply_guardrail_unmask_on_response():
    """
    When input_type is 'response' and pii_tokens exist, apply_guardrail
    should unmask text instead of masking it.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        guardrail_name="test_presidio",
        output_parse_pii=True,
        mock_testing=True,
    )

    request_data = {
        "model": "gpt-4o",
        "metadata": {
            "pii_tokens": {
                "<PERSON_1>": "John Smith",
                "<PHONE_NUMBER_1>": "555-123-4567",
            }
        },
    }

    inputs = {
        "texts": [
            "Hello <PERSON_1>, your number is <PHONE_NUMBER_1>.",
        ]
    }

    result = await guardrail.apply_guardrail(
        inputs=inputs,
        request_data=request_data,
        input_type="response",
    )

    assert result["texts"][0] == "Hello John Smith, your number is 555-123-4567."


@pytest.mark.asyncio
async def test_apply_guardrail_masks_on_request():
    """
    When input_type is 'request', apply_guardrail should mask as before.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        guardrail_name="test_presidio",
        output_parse_pii=True,
        mock_testing=True,
    )

    async def mock_check_pii(text, output_parse_pii, presidio_config, request_data):
        return text.replace("John Smith", "<PERSON>")

    guardrail.check_pii = mock_check_pii

    result = await guardrail.apply_guardrail(
        inputs={"texts": ["Hello John Smith"]},
        request_data={"model": "gpt-4o", "metadata": {}},
        input_type="request",
    )

    assert "<PERSON>" in result["texts"][0]
    assert "John Smith" not in result["texts"][0]


@pytest.mark.asyncio
async def test_apply_to_output_streaming_bytes_only_logs_warning():
    """
    Regression test: when apply_to_output=True and the stream contains only
    bytes chunks (Anthropic native SSE), output masking is skipped.
    A warning must be logged so operators are aware.
    """
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        apply_to_output=True,
    )

    byte_chunks = [
        b'data: {"type":"content_block_delta","delta":{"text":"Hello"}}\n\n',
        b'data: {"type":"content_block_delta","delta":{"text":" world"}}\n\n',
    ]

    async def mock_stream():
        for b in byte_chunks:
            yield b

    mock_user_api_key = UserAPIKeyAuth(api_key="test-key")

    collected = []
    with patch(
        "litellm.proxy.guardrails.guardrail_hooks.presidio.verbose_proxy_logger"
    ) as mock_logger:
        async for chunk in guardrail.async_post_call_streaming_iterator_hook(
            user_api_key_dict=mock_user_api_key,
            response=mock_stream(),
            request_data={},
        ):
            collected.append(chunk)

        # All bytes should be yielded through
        assert len(collected) == len(byte_chunks)
        for original, received in zip(byte_chunks, collected):
            assert original == received

        # Warning must be logged about skipped masking
        mock_logger.warning.assert_called_once()
        warning_msg = mock_logger.warning.call_args[0][0]
        assert "Output PII masking was skipped" in warning_msg


@pytest.mark.asyncio
async def test_anonymize_text_uses_correct_positions_no_parse_pii():
    """
    Regression test for anonymizer offset bug (fixes #24160).

    The Presidio anonymizer returns items with start/end positions that
    reference the *anonymized output* text, not the original input text.
    When output_parse_pii is False, anonymize_text must return
    redacted_text["text"] directly instead of manually splicing the
    original text using those positions, which produces garbled output
    with remnants of original PII data.
    """
    original_text = (
        "My name is John Smith, my email is john@example.com, phone 555-867-5309"
    )
    # Positions as returned by the analyzer (reference original text)
    analyze_results = [
        {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35},
        {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11},
        {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59},
    ]
    # Anonymizer response — positions reference the *anonymized* text
    anonymizer_response = {
        "text": "My name is <PERSON>, my email is <EMAIL_ADDRESS>, phone <PHONE_NUMBER>",
        "items": [
            {
                "start": 56,
                "end": 70,
                "entity_type": "PHONE_NUMBER",
                "text": "<PHONE_NUMBER>",
                "operator": "replace",
            },
            {
                "start": 33,
                "end": 48,
                "entity_type": "EMAIL_ADDRESS",
                "text": "<EMAIL_ADDRESS>",
                "operator": "replace",
            },
            {
                "start": 11,
                "end": 19,
                "entity_type": "PERSON",
                "text": "<PERSON>",
                "operator": "replace",
            },
        ],
    }

    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        mock_testing=False,
    )

    mock_iterator = _make_mock_session_iterator(
        json_response=anonymizer_response,
    )

    masked_entity_count = {}
    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        result = await guardrail.anonymize_text(
            text=original_text,
            analyze_results=analyze_results,
            output_parse_pii=False,
            masked_entity_count=masked_entity_count,
        )

    expected = "My name is <PERSON>, my email is <EMAIL_ADDRESS>, phone <PHONE_NUMBER>"
    assert (
        result == expected
    ), f"anonymize_text produced garbled output with PII remnants.\nExpected: {expected!r}\nGot:      {result!r}"
    assert masked_entity_count == {
        "PERSON": 1,
        "EMAIL_ADDRESS": 1,
        "PHONE_NUMBER": 1,
    }


@pytest.mark.asyncio
async def test_anonymize_text_uses_correct_positions_with_parse_pii():
    """
    Regression test for anonymizer offset bug with output_parse_pii=True
    (fixes #24160).

    When output_parse_pii is True, anonymize_text must use positions from
    analyze_results (which reference the original text) to build numbered
    tokens and the pii_tokens mapping, not positions from anonymizer items
    (which reference the anonymized output text).
    """
    original_text = (
        "My name is John Smith, my email is john@example.com, phone 555-867-5309"
    )
    analyze_results = [
        {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35},
        {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11},
        {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59},
    ]
    anonymizer_response = {
        "text": "My name is <PERSON>, my email is <EMAIL_ADDRESS>, phone <PHONE_NUMBER>",
        "items": [
            {
                "start": 56,
                "end": 70,
                "entity_type": "PHONE_NUMBER",
                "text": "<PHONE_NUMBER>",
                "operator": "replace",
            },
            {
                "start": 33,
                "end": 48,
                "entity_type": "EMAIL_ADDRESS",
                "text": "<EMAIL_ADDRESS>",
                "operator": "replace",
            },
            {
                "start": 11,
                "end": 19,
                "entity_type": "PERSON",
                "text": "<PERSON>",
                "operator": "replace",
            },
        ],
    }

    guardrail = _OPTIONAL_PresidioPIIMasking(
        presidio_analyzer_api_base="http://test-analyzer/",
        presidio_anonymizer_api_base="http://test-anonymizer/",
        mock_testing=False,
        output_parse_pii=True,
    )

    mock_iterator = _make_mock_session_iterator(
        json_response=anonymizer_response,
    )

    masked_entity_count = {}
    request_data = {"metadata": {}}
    with patch.object(guardrail, "_get_session_iterator", mock_iterator):
        result = await guardrail.anonymize_text(
            text=original_text,
            analyze_results=analyze_results,
            output_parse_pii=True,
            masked_entity_count=masked_entity_count,
            request_data=request_data,
        )

    # Result must not contain any remnants of original PII
    assert "John" not in result
    assert "john@example.com" not in result
    assert "555-867-5309" not in result

    # pii_tokens must map numbered tokens back to correct original values
    pii_tokens = request_data["metadata"]["pii_tokens"]
    token_values = set(pii_tokens.values())
    assert "John Smith" in token_values
    assert "john@example.com" in token_values
    assert "555-867-5309" in token_values

    # Tokens must be numbered in left-to-right order of appearance:
    # PERSON (pos 11) → _1, EMAIL_ADDRESS (pos 35) → _2, PHONE_NUMBER (pos 59) → _3
    assert pii_tokens.get("<PERSON_1>") == "John Smith"
    assert pii_tokens.get("<EMAIL_ADDRESS_2>") == "john@example.com"
    assert pii_tokens.get("<PHONE_NUMBER_3>") == "555-867-5309"


def test_unmask_sse_bytes_chunk_replaces_text_delta():
    import json

    pii_tokens = {"<PERSON_1>": "Bobby"}
    event = {
        "type": "content_block_delta",
        "index": 0,
        "delta": {"type": "text_delta", "text": "Hello <PERSON_1>, how are you?"},
    }
    chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8")

    result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens)

    decoded = result.decode("utf-8")
    parsed = json.loads(decoded.split("data: ", 1)[1].strip())
    assert parsed["delta"]["text"] == "Hello Bobby, how are you?"


def test_unmask_sse_bytes_chunk_ignores_non_text_delta():
    import json

    pii_tokens = {"<PERSON_1>": "Bobby"}

    # message_start event — no delta
    event = {"type": "message_start", "message": {"id": "msg_01", "role": "assistant"}}
    chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8")
    result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens)
    assert result == chunk

    # input_json_delta — should not be touched
    event2 = {
        "type": "content_block_delta",
        "index": 1,
        "delta": {"type": "input_json_delta", "partial_json": '{"name": "<PERSON_1>"}'},
    }
    chunk2 = ("data: " + json.dumps(event2) + "\n\n").encode("utf-8")
    result2 = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk2, pii_tokens)
    assert result2 == chunk2


def test_unmask_sse_bytes_chunk_handles_malformed_json():
    chunk = b"data: {not valid json}\n\n"
    result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(
        chunk, {"<PERSON_1>": "Bobby"}
    )
    assert result == chunk


def test_unmask_sse_bytes_chunk_handles_unicode_decode_error():
    chunk = b"\xff\xfe invalid utf-8"
    result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(
        chunk, {"<PERSON_1>": "Bobby"}
    )
    assert result == chunk


def test_unmask_sse_bytes_chunk_non_ascii_pii_not_escaped():
    import json

    pii_tokens = {"<PERSON_1>": "José"}
    event = {
        "type": "content_block_delta",
        "index": 0,
        "delta": {"type": "text_delta", "text": "Hello <PERSON_1>!"},
    }
    chunk = ("data: " + json.dumps(event) + "\n\n").encode("utf-8")

    result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(chunk, pii_tokens)

    decoded = result.decode("utf-8")
    assert "Jos\\u" not in decoded
    parsed = json.loads(decoded.split("data: ", 1)[1].strip())
    assert parsed["delta"]["text"] == "Hello José!"


def test_unmask_sse_bytes_chunk_handles_crlf_line_endings():
    import json

    pii_tokens = {"<PERSON_1>": "Bobby"}
    event = {
        "type": "content_block_delta",
        "index": 0,
        "delta": {"type": "text_delta", "text": "Hi <PERSON_1>!"},
    }
    crlf_chunk = ("data: " + json.dumps(event) + "\r\ndata: [DONE]\r\n").encode("utf-8")

    result = _OPTIONAL_PresidioPIIMasking._unmask_sse_bytes_chunk(
        crlf_chunk, pii_tokens
    )

    decoded = result.decode("utf-8")
    parsed = json.loads(decoded.split("data: ", 1)[1].split("\n")[0].strip())
    assert parsed["delta"]["text"] == "Hi Bobby!"
    assert "data: [DONE]" in decoded


@pytest.mark.asyncio
async def test_stream_pii_unmasking_unmaskes_bytes_chunks(mock_user_api_key):
    import json

    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
    )

    pii_tokens = {"<PERSON_1>": "Bobby"}
    request_data = {"metadata": {"pii_tokens": pii_tokens}}

    def _make_sse_chunk(text: str) -> bytes:
        event = {
            "type": "content_block_delta",
            "index": 0,
            "delta": {"type": "text_delta", "text": text},
        }
        return ("data: " + json.dumps(event) + "\n\n").encode("utf-8")

    async def mock_stream():
        yield _make_sse_chunk("Hello <PERSON_1>!")
        yield _make_sse_chunk(" How can I help?")

    chunks = []
    async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data):
        chunks.append(chunk)

    assert len(chunks) == 2
    first = chunks[0].decode("utf-8")
    first_event = json.loads(first.split("data: ", 1)[1].strip())
    assert first_event["delta"]["text"] == "Hello Bobby!"

    second = chunks[1].decode("utf-8")
    second_event = json.loads(second.split("data: ", 1)[1].strip())
    assert second_event["delta"]["text"] == " How can I help?"


@pytest.mark.asyncio
async def test_stream_pii_unmasking_passthrough_when_no_tokens(mock_user_api_key):
    guardrail = _OPTIONAL_PresidioPIIMasking(
        mock_testing=True,
        output_parse_pii=True,
    )

    raw_chunk = b"data: {}\n\n"
    request_data: dict = {"metadata": {}}

    async def mock_stream():
        yield raw_chunk

    chunks = []
    async for chunk in guardrail._stream_pii_unmasking(mock_stream(), request_data):
        chunks.append(chunk)

    assert chunks == [raw_chunk]
