import copy
import os
from typing import cast
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest

import litellm
from litellm import ModelResponse
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.noma import (
    NomaGuardrail,
    NomaV2Guardrail,
    initialize_guardrail,
)
import litellm.proxy.guardrails.guardrail_hooks.noma.noma as noma_legacy_module
from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message


@pytest.fixture
def noma_guardrail():
    """Create a NomaGuardrail instance for testing"""
    return NomaGuardrail(
        api_key="test-api-key",
        api_base="https://api.test.noma.security/",
        application_id="test-app",
        monitor_mode=False,
        block_failures=True,
        guardrail_name="test-noma-guardrail",
        event_hook="pre_call",
        default_on=True,
    )


@pytest.fixture
def mock_user_api_key_dict():
    """Create a mock UserAPIKeyAuth object"""
    return UserAPIKeyAuth(
        user_id="test-user-id",
        user_email="test@example.com",
        key_name="test-key",
        key_alias=None,
        team_id=None,
        team_alias=None,
        user_role=None,
        api_key="test-api-key",
        permissions={},
        models=[],
        spend=0.0,
        max_budget=None,
        soft_budget=None,
        tpm_limit=None,
        rpm_limit=None,
        metadata={},
        max_parallel_requests=None,
        allowed_cache_controls=[],
        model_spend={},
        model_max_budget={},
    )


@pytest.fixture
def mock_request_data():
    """Create mock request data"""
    return {
        "messages": [
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": "Hello, how are you?"},
        ],
        "litellm_call_id": "test-call-id",
        "metadata": {"requester_ip_address": "192.168.1.1"},
    }


class TestNomaGuardrailConfiguration:
    """Test configuration and initialization of Noma guardrail"""

    def test_legacy_guardrail_emits_deprecation_warning(self, monkeypatch):
        monkeypatch.setattr(
            noma_legacy_module, "_LEGACY_NOMA_DEPRECATION_WARNED", False
        )
        with pytest.warns(DeprecationWarning, match="deprecated"):
            NomaGuardrail(api_key="test-api-key")

    def test_init_with_config(self):
        """Test initializing Noma guardrail via init_guardrails_v2"""
        with patch.dict(
            os.environ,
            {
                "NOMA_API_KEY": "test-api-key",
                "NOMA_API_BASE": "https://api.test.noma.security/",
            },
        ):
            init_guardrails_v2(
                all_guardrails=[
                    {
                        "guardrail_name": "noma-pre-guard",
                        "litellm_params": {
                            "guardrail": "noma",
                            "mode": "pre_call",
                            "application_id": "test-app",
                            "monitor_mode": False,
                            "block_failures": True,
                        },
                    }
                ],
                config_file_path="",
            )

    def test_init_with_env_vars(self):
        """Test initialization with environment variables"""
        with patch.dict(
            os.environ,
            {
                "NOMA_API_KEY": "env-api-key",
                "NOMA_API_BASE": "https://env.api.noma.security/",
                "NOMA_APPLICATION_ID": "env-app-id",
                "NOMA_MONITOR_MODE": "true",
                "NOMA_BLOCK_FAILURES": "false",
            },
        ):
            guardrail = NomaGuardrail()
            assert guardrail.api_key == "env-api-key"
            assert guardrail.api_base == "https://env.api.noma.security/"
            assert guardrail.application_id == "env-app-id"
            assert guardrail.monitor_mode is True
            assert guardrail.block_failures is False

    def test_init_with_params_override_env(self):
        """Test that constructor params override environment variables"""
        with patch.dict(
            os.environ,
            {
                "NOMA_API_KEY": "env-api-key",
                "NOMA_MONITOR_MODE": "true",
            },
        ):
            guardrail = NomaGuardrail(
                api_key="param-api-key",
                monitor_mode=False,
            )
            assert guardrail.api_key == "param-api-key"
            assert guardrail.monitor_mode is False

    def test_initialize_guardrail_function(self):
        """Test the initialize_guardrail function"""
        from litellm.types.guardrails import Guardrail, LitellmParams

        litellm_params = LitellmParams(
            guardrail="noma",
            mode="pre_call",
            api_key="test-key",
            api_base="https://test.api/",
            application_id="test-app",
            monitor_mode=True,
            block_failures=False,
        )

        guardrail = Guardrail(
            guardrail_name="test-guardrail",
            litellm_params=litellm_params,
        )

        with patch("litellm.logging_callback_manager.add_litellm_callback") as mock_add:
            result = initialize_guardrail(litellm_params, guardrail)

            assert isinstance(result, NomaGuardrail)
            assert result.api_key == "test-key"
            assert result.api_base == "https://test.api/"
            assert result.application_id == "test-app"
            assert result.monitor_mode is True
            assert result.block_failures is False
            mock_add.assert_called_once_with(result)

    def test_initialize_guardrail_use_v2_routes_to_noma_v2(self):
        """Test migration routing: guardrail=noma + use_v2=True initializes NomaV2Guardrail."""
        from litellm.types.guardrails import Guardrail, LitellmParams

        litellm_params = LitellmParams(
            guardrail="noma",
            mode="pre_call",
            use_v2=True,
            api_key="test-key",
            api_base="https://test.api/",
            application_id="test-app",
        )

        guardrail = Guardrail(
            guardrail_name="test-guardrail",
            litellm_params=litellm_params,
        )

        with patch("litellm.logging_callback_manager.add_litellm_callback") as mock_add:
            result = initialize_guardrail(litellm_params, guardrail)

            assert isinstance(result, NomaV2Guardrail)
            assert result.api_key == "test-key"
            assert result.api_base == "https://test.api"
            assert result.application_id == "test-app"
            mock_add.assert_called_once_with(result)


class TestNomaApplicationIdResolution:
    """Tests for determining which applicationId is sent to Noma."""

    @staticmethod
    def _clone_user_auth(user_auth: UserAPIKeyAuth) -> UserAPIKeyAuth:
        return UserAPIKeyAuth(**user_auth.model_dump())

    @staticmethod
    def _clone_request_data(request_data: dict) -> dict:
        return copy.deepcopy(request_data)

    @staticmethod
    async def _get_application_id(
        guardrail: NomaGuardrail,
        user_auth: UserAPIKeyAuth,
        request_data: dict,
        extra_data: dict,
    ) -> str:
        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": False,
            "scanResult": [],
        }
        mock_response.raise_for_status = MagicMock()
        mock_post = AsyncMock(return_value=mock_response)

        with patch.object(guardrail.async_handler, "post", mock_post):
            await guardrail._call_noma_api(
                payload={"input": []},
                llm_request_id=None,
                request_data=request_data,
                user_auth=user_auth,
                extra_data=extra_data,
            )

        sent_payload = mock_post.call_args.kwargs["json"]
        return sent_payload["x-noma-context"]["applicationId"]

    @pytest.mark.asyncio
    async def test_application_id_prefers_extra_body(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        request_data = self._clone_request_data(mock_request_data)
        request_data.setdefault("metadata", {}).setdefault("headers", {})[
            "x-noma-application-id"
        ] = "header-app"
        user_auth = self._clone_user_auth(mock_user_api_key_dict)
        user_auth.key_alias = "alias-app"

        application_id = await self._get_application_id(
            guardrail=noma_guardrail,
            user_auth=user_auth,
            request_data=request_data,
            extra_data={"application_id": "dynamic-app"},
        )

        assert application_id == "dynamic-app"

    @pytest.mark.asyncio
    async def test_application_id_prefers_header_over_alias(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        request_data = self._clone_request_data(mock_request_data)
        request_data.setdefault("metadata", {}).setdefault("headers", {})[
            "x-noma-application-id"
        ] = "header-app"
        user_auth = self._clone_user_auth(mock_user_api_key_dict)
        user_auth.key_alias = "alias-app"
        original_app_id = noma_guardrail.application_id
        noma_guardrail.application_id = None

        try:
            application_id = await self._get_application_id(
                guardrail=noma_guardrail,
                user_auth=user_auth,
                request_data=request_data,
                extra_data={},
            )
        finally:
            noma_guardrail.application_id = original_app_id

        assert application_id == "header-app"

    @pytest.mark.asyncio
    async def test_application_id_uses_guardrail_config_before_alias(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        request_data = self._clone_request_data(mock_request_data)
        user_auth = self._clone_user_auth(mock_user_api_key_dict)
        user_auth.key_alias = "alias-app"
        original_app_id = noma_guardrail.application_id
        noma_guardrail.application_id = "config-app"

        try:
            application_id = await self._get_application_id(
                guardrail=noma_guardrail,
                user_auth=user_auth,
                request_data=request_data,
                extra_data={},
            )
        finally:
            noma_guardrail.application_id = original_app_id

        assert application_id == "config-app"

    @pytest.mark.asyncio
    async def test_application_id_falls_back_to_alias(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        request_data = self._clone_request_data(mock_request_data)
        user_auth = self._clone_user_auth(mock_user_api_key_dict)
        user_auth.key_alias = "alias-app"
        original_app_id = noma_guardrail.application_id
        noma_guardrail.application_id = None

        try:
            application_id = await self._get_application_id(
                guardrail=noma_guardrail,
                user_auth=user_auth,
                request_data=request_data,
                extra_data={},
            )
        finally:
            noma_guardrail.application_id = original_app_id

        assert application_id == "alias-app"

    @pytest.mark.asyncio
    async def test_application_id_defaults_to_litellm(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        request_data = self._clone_request_data(mock_request_data)
        user_auth = self._clone_user_auth(mock_user_api_key_dict)
        user_auth.key_alias = None
        original_app_id = noma_guardrail.application_id
        noma_guardrail.application_id = None

        try:
            application_id = await self._get_application_id(
                guardrail=noma_guardrail,
                user_auth=user_auth,
                request_data=request_data,
                extra_data={},
            )
        finally:
            noma_guardrail.application_id = original_app_id

        assert application_id == "litellm"


class TestNomaBlockedMessage:
    """Test the NomaBlockedMessage exception class"""

    def test_blocked_message_basic(self):
        """Test basic blocked message creation"""
        response = {
            "aggregatedScanResult": True,
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "harmfulContent": {
                            "result": True,
                            "probability": 0.9,
                            "status": "SUCCESS",
                        },
                        "code": {
                            "result": False,
                            "probability": 0.1,
                            "status": "SUCCESS",
                        },
                    },
                }
            ],
        }

        exception = NomaBlockedMessage(response)
        assert exception.status_code == 400
        assert exception.detail["error"] == "Request blocked by Noma guardrail"

    def test_blocked_message_with_data_detection(self):
        """Test blocked message with data detection"""
        response = {
            "aggregatedScanResult": True,
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "sensitiveData": {
                            "PII": {
                                "result": True,
                                "probability": 0.8,
                                "status": "SUCCESS",
                            },
                            "PCI": {
                                "result": False,
                                "probability": 0,
                                "status": "SUCCESS",
                            },
                        },
                    },
                }
            ],
        }

        exception = NomaBlockedMessage(response)
        assert exception.detail["error"] == "Request blocked by Noma guardrail"

    def test_blocked_message_with_topics(self):
        """Test blocked message with topic guardrails"""
        response = {
            "aggregatedScanResult": True,
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "customLlm": {
                            "topic1": {
                                "result": True,
                                "probability": 0.95,
                                "status": "SUCCESS",
                            },
                            "topic2": {
                                "result": False,
                                "probability": 0.2,
                                "status": "SUCCESS",
                            },
                        },
                    },
                }
            ],
        }

        exception = NomaBlockedMessage(response)
        assert exception.detail["error"] == "Request blocked by Noma guardrail"


class TestNomaGuardrailHooks:
    """Test the guardrail hook methods"""

    @pytest.mark.asyncio
    async def test_pre_call_hook_allowed(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test pre-call hook when content is allowed"""
        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [{"role": "user", "type": "message", "results": {}}],
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ) as mock_post:
            result = await noma_guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=mock_request_data,
                call_type="completion",
            )

            assert result == mock_request_data
            mock_post.assert_called_once()

            # Verify API call details
            call_args = mock_post.call_args
            # Verify the URL endpoint
            assert call_args.args[0].endswith("/ai-dr/v2/prompt/scan")
            # Verify headers and JSON payload
            if "headers" in call_args.kwargs:
                headers = call_args.kwargs["headers"]
                assert "Authorization" in headers
                assert headers["Authorization"] == "Bearer test-api-key"
            # Verify application ID is in JSON payload (not headers)
            if "json" in call_args.kwargs:
                json_payload = call_args.kwargs["json"]
                assert "x-noma-context" in json_payload
                assert json_payload["x-noma-context"]["applicationId"] == "test-app"

    @pytest.mark.asyncio
    async def test_pre_call_hook_with_system_prompt(
        self, noma_guardrail, mock_user_api_key_dict
    ):
        """Test pre-call hook includes system prompt in Noma API request"""
        request_data = {
            "messages": [
                {"role": "system", "content": "You are a helpful assistant"},
                {"role": "user", "content": "Hello, how are you?"},
            ],
            "litellm_call_id": "test-call-id",
            "metadata": {"requester_ip_address": "192.168.1.1"},
        }

        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [
                {"role": "system", "type": "message", "results": {}},
                {"role": "user", "type": "message", "results": {}},
            ],
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ) as mock_post:
            result = await noma_guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=request_data,
                call_type="completion",
            )

            assert result == request_data
            mock_post.assert_called_once()

            # Verify the payload includes both system and user messages
            call_args = mock_post.call_args
            json_payload = call_args.kwargs["json"]
            assert "input" in json_payload
            messages = json_payload["input"]

            # Should have 2 messages: system and user
            assert len(messages) == 2

            # First message should be system
            assert messages[0]["type"] == "message"
            assert messages[0]["role"] == "system"
            assert messages[0]["content"][0]["type"] == "input_text"
            assert messages[0]["content"][0]["text"] == "You are a helpful assistant"

            # Second message should be user
            assert messages[1]["type"] == "message"
            assert messages[1]["role"] == "user"
            assert messages[1]["content"][0]["type"] == "input_text"
            assert messages[1]["content"][0]["text"] == "Hello, how are you?"

    @pytest.mark.asyncio
    async def test_pre_call_hook_with_multiple_system_prompts(
        self, noma_guardrail, mock_user_api_key_dict
    ):
        """Test pre-call hook combines multiple system prompts into single message"""
        request_data = {
            "messages": [
                {"role": "system", "content": "You are a helpful assistant"},
                {"role": "system", "content": "You should be polite and respectful"},
                {"role": "user", "content": "Hello, how are you?"},
            ],
            "litellm_call_id": "test-call-id",
        }

        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": False,
            "scanResult": [
                {"role": "system", "type": "message", "results": {}},
                {"role": "user", "type": "message", "results": {}},
            ],
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ) as mock_post:
            result = await noma_guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=request_data,
                call_type="completion",
            )

            assert result == request_data
            mock_post.assert_called_once()

            # Verify the payload combines system prompts into single message
            call_args = mock_post.call_args
            json_payload = call_args.kwargs["json"]
            messages = json_payload["input"]

            # Should have 2 messages: 1 combined system and 1 user
            assert len(messages) == 2

            # First message should be system with combined content
            assert messages[0]["role"] == "system"
            assert messages[0]["content"][0]["type"] == "input_text"
            assert (
                messages[0]["content"][0]["text"]
                == "You are a helpful assistant You should be polite and respectful"
            )

            # Second message should be user
            assert messages[1]["role"] == "user"
            assert messages[1]["content"][0]["type"] == "input_text"
            assert messages[1]["content"][0]["text"] == "Hello, how are you?"

    @pytest.mark.asyncio
    async def test_pre_call_hook_blocked(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test pre-call hook when content is blocked"""
        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "harmfulContent": {
                            "result": True,
                            "probability": 0.9,
                            "status": "SUCCESS",
                        }
                    },
                }
            ],
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ):
            with pytest.raises(NomaBlockedMessage) as exc_info:
                await noma_guardrail.async_pre_call_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    cache=MagicMock(),
                    data=mock_request_data,
                    call_type="completion",
                )

            assert exc_info.value.status_code == 400

    @pytest.mark.asyncio
    async def test_pre_call_hook_monitor_mode(
        self, mock_user_api_key_dict, mock_request_data
    ):
        """Test pre-call hook in monitor mode (logs but doesn't block)"""
        guardrail = NomaGuardrail(
            api_key="test-key",
            monitor_mode=True,
            guardrail_name="test-guardrail",
            event_hook="pre_call",
            default_on=True,
        )

        with patch.object(
            guardrail, "_create_background_noma_check"
        ) as mock_create_background:
            # Should return immediately without waiting for API call
            result = await guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=mock_request_data,
                call_type="completion",
            )

            assert result == mock_request_data
            # Verify background task was created
            mock_create_background.assert_called_once()

    @pytest.mark.asyncio
    async def test_pre_call_hook_monitor_mode_background_task_failure(
        self, mock_user_api_key_dict, mock_request_data
    ):
        """Test pre-call hook in monitor mode when background task creation fails"""
        guardrail = NomaGuardrail(
            api_key="test-key",
            monitor_mode=True,
            guardrail_name="test-guardrail",
            event_hook="pre_call",
            default_on=True,
        )

        with patch.object(
            guardrail,
            "_create_background_noma_check",
            side_effect=Exception("Task creation failed"),
        ):
            # Should still return successfully even if background task creation fails
            result = await guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=mock_request_data,
                call_type="completion",
            )

            assert result == mock_request_data

    @pytest.mark.asyncio
    async def test_post_call_success_hook(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test post-call success hook"""
        # Create a mock ModelResponse
        response = ModelResponse(
            id="test-response-id",
            choices=[
                Choices(
                    finish_reason="stop",
                    index=0,
                    message=Message(
                        content="I'm doing well, thank you!", role="assistant"
                    ),
                )
            ],
            created=1234567890,
            model="gpt-3.5-turbo",
            object="chat.completion",
            system_fingerprint=None,
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )

        mock_api_response = MagicMock()
        mock_api_response.json.return_value = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [{"role": "assistant", "type": "message", "results": {}}],
        }
        mock_api_response.raise_for_status = MagicMock()

        # Update guardrail to use post_call event hook
        noma_guardrail.event_hook = "post_call"

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_api_response
        ) as mock_post:
            result = await noma_guardrail.async_post_call_success_hook(
                data=mock_request_data,
                user_api_key_dict=mock_user_api_key_dict,
                response=response,
            )

            assert result == response
            mock_post.assert_called_once()

    @pytest.mark.asyncio
    async def test_moderation_hook(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test moderation hook (during_call)"""
        # Update guardrail to use during_call event hook
        noma_guardrail.event_hook = "during_call"

        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [{"role": "user", "type": "message", "results": {}}],
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ):
            result = await noma_guardrail.async_moderation_hook(
                data=mock_request_data,
                user_api_key_dict=mock_user_api_key_dict,
                call_type="completion",
            )

            assert result == mock_request_data

    @pytest.mark.asyncio
    async def test_api_failure_handling(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        with patch.object(
            noma_guardrail.async_handler,
            "post",
            side_effect=httpx.HTTPStatusError(
                "API Error", request=MagicMock(), response=MagicMock(status_code=500)
            ),
        ):
            with pytest.raises(httpx.HTTPStatusError):
                await noma_guardrail.async_pre_call_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    cache=MagicMock(),
                    data=mock_request_data,
                    call_type="completion",
                )

    @pytest.mark.asyncio
    async def test_api_failure_no_block(
        self, mock_user_api_key_dict, mock_request_data
    ):
        guardrail = NomaGuardrail(
            api_key="test-key",
            block_failures=False,
            guardrail_name="test-guardrail",
            event_hook="pre_call",
            default_on=True,
        )

        with patch.object(
            guardrail.async_handler,
            "post",
            side_effect=httpx.HTTPStatusError(
                "API Error", request=MagicMock(), response=MagicMock(status_code=500)
            ),
        ):
            result = await guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=mock_request_data,
                call_type="completion",
            )

            assert result == mock_request_data


class TestBackgroundProcessing:
    """Test the new background processing functionality"""

    @pytest.fixture
    def monitor_mode_guardrail(self):
        """Create a guardrail with monitor mode enabled"""
        return NomaGuardrail(
            api_key="test-api-key",
            api_base="https://api.test.noma.security/",
            application_id="test-app",
            monitor_mode=True,  # Enable monitor mode
            block_failures=True,
            guardrail_name="test-noma-guardrail",
            event_hook="pre_call",
            default_on=True,
        )

    @pytest.mark.asyncio
    async def test_process_user_message_check_monitor_mode(
        self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test shared helper method in monitor mode"""
        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "harmfulContent": {"result": True, "status": "SUCCESS"}
                    },
                }
            ],
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            monitor_mode_guardrail.async_handler, "post", return_value=mock_response
        ) as mock_post:
            with patch.object(
                monitor_mode_guardrail, "_handle_verdict_background"
            ) as mock_handle_verdict:
                result = await monitor_mode_guardrail._process_user_message_check(
                    mock_request_data, mock_user_api_key_dict
                )

                assert result is not None
                mock_post.assert_called_once()
                mock_handle_verdict.assert_called_once()

    @pytest.mark.asyncio
    async def test_process_user_message_check_non_monitor_mode(
        self, noma_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test shared helper method in non-monitor mode"""
        mock_response = MagicMock()
        mock_response.json.return_value = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [{"role": "user", "type": "message", "results": {}}],
        }
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ) as mock_post:
            with patch.object(noma_guardrail, "_check_verdict") as mock_check_verdict:
                result = await noma_guardrail._process_user_message_check(
                    mock_request_data, mock_user_api_key_dict
                )

                assert result is not None
                mock_post.assert_called_once()
                mock_check_verdict.assert_called_once()

    @pytest.mark.asyncio
    async def test_process_llm_response_check_monitor_mode(
        self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test LLM response processing in monitor mode"""
        from litellm.types.utils import Choices, Message

        response = ModelResponse(
            id="test-response-id",
            choices=[
                Choices(
                    finish_reason="stop",
                    index=0,
                    message=Message(
                        content="I'm doing well, thank you!", role="assistant"
                    ),
                )
            ],
            created=1234567890,
            model="gpt-3.5-turbo",
            object="chat.completion",
            system_fingerprint=None,
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )

        mock_api_response = MagicMock()
        mock_api_response.json.return_value = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [{"role": "assistant", "type": "message", "results": {}}],
        }
        mock_api_response.raise_for_status = MagicMock()

        with patch.object(
            monitor_mode_guardrail.async_handler, "post", return_value=mock_api_response
        ) as mock_post:
            with patch.object(
                monitor_mode_guardrail, "_handle_verdict_background"
            ) as mock_handle_verdict:
                result = await monitor_mode_guardrail._process_llm_response_check(
                    mock_request_data, response, mock_user_api_key_dict
                )

                assert result == "I'm doing well, thank you!"
                mock_post.assert_called_once()
                mock_handle_verdict.assert_called_once()

    @pytest.mark.asyncio
    async def test_check_user_message_background(
        self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test background user message check method"""
        with patch.object(
            monitor_mode_guardrail, "_process_user_message_check"
        ) as mock_process:
            await monitor_mode_guardrail._check_user_message_background(
                mock_request_data, mock_user_api_key_dict
            )

            mock_process.assert_called_once_with(
                mock_request_data, mock_user_api_key_dict
            )

    @pytest.mark.asyncio
    async def test_check_user_message_background_exception_handling(
        self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test background user message check handles exceptions gracefully"""
        with patch.object(
            monitor_mode_guardrail,
            "_process_user_message_check",
            side_effect=Exception("API failed"),
        ):
            # Should not raise exception, just log error
            await monitor_mode_guardrail._check_user_message_background(
                mock_request_data, mock_user_api_key_dict
            )

    @pytest.mark.asyncio
    async def test_check_llm_response_background(
        self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test background LLM response check method"""
        from litellm.types.utils import Choices, Message

        response = ModelResponse(
            id="test-response-id",
            choices=[
                Choices(
                    finish_reason="stop",
                    index=0,
                    message=Message(content="Test response", role="assistant"),
                )
            ],
            created=1234567890,
            model="gpt-3.5-turbo",
            object="chat.completion",
            system_fingerprint=None,
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )

        with patch.object(
            monitor_mode_guardrail, "_process_llm_response_check"
        ) as mock_process:
            await monitor_mode_guardrail._check_llm_response_background(
                mock_request_data, response, mock_user_api_key_dict
            )

            mock_process.assert_called_once_with(
                mock_request_data, response, mock_user_api_key_dict
            )

    @pytest.mark.asyncio
    async def test_handle_verdict_background_blocked(self, monitor_mode_guardrail):
        """Test background verdict handling for blocked content"""
        response_json = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "harmfulContent": {"result": True, "status": "SUCCESS"}
                    },
                }
            ],
        }

        with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning:
            await monitor_mode_guardrail._handle_verdict_background(
                "user", "test message", response_json
            )

            mock_warning.assert_called_once()
            assert "blocked user message" in mock_warning.call_args[0][0]

    @pytest.mark.asyncio
    async def test_handle_verdict_background_allowed(self, monitor_mode_guardrail):
        """Test background verdict handling for allowed content"""
        response_json = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [{"role": "assistant", "type": "message", "results": {}}],
        }

        with patch("litellm._logging.verbose_proxy_logger.info") as mock_info:
            await monitor_mode_guardrail._handle_verdict_background(
                "assistant", "test response", response_json
            )

            mock_info.assert_called_once()
            assert "allowed assistant message" in mock_info.call_args[0][0]

    @pytest.mark.asyncio
    async def test_create_background_noma_check(self, monitor_mode_guardrail):
        """Test background task creation"""

        async def dummy_coroutine():
            return "completed"

        with patch("asyncio.create_task") as mock_create_task:
            monitor_mode_guardrail._create_background_noma_check(dummy_coroutine())
            mock_create_task.assert_called_once()

    @pytest.mark.asyncio
    async def test_create_background_noma_check_exception(self, monitor_mode_guardrail):
        """Test background task creation with exception handling"""

        async def dummy_coroutine():
            return "completed"

        with patch(
            "asyncio.create_task", side_effect=Exception("Task creation failed")
        ):
            # Should not raise exception, just log error
            monitor_mode_guardrail._create_background_noma_check(dummy_coroutine())

    @pytest.mark.asyncio
    async def test_moderation_hook_monitor_mode(
        self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test moderation hook in monitor mode"""
        # Update event hook to during_call
        monitor_mode_guardrail.event_hook = "during_call"

        with patch.object(
            monitor_mode_guardrail, "_create_background_noma_check"
        ) as mock_create_background:
            result = await monitor_mode_guardrail.async_moderation_hook(
                data=mock_request_data,
                user_api_key_dict=mock_user_api_key_dict,
                call_type="completion",
            )

            assert result == mock_request_data
            mock_create_background.assert_called_once()

    @pytest.mark.asyncio
    async def test_post_call_success_hook_monitor_mode(
        self, monitor_mode_guardrail, mock_user_api_key_dict, mock_request_data
    ):
        """Test post-call success hook in monitor mode"""
        from litellm.types.utils import Choices, Message

        # Update event hook to post_call
        monitor_mode_guardrail.event_hook = "post_call"

        response = ModelResponse(
            id="test-response-id",
            choices=[
                Choices(
                    finish_reason="stop",
                    index=0,
                    message=Message(content="Test response", role="assistant"),
                )
            ],
            created=1234567890,
            model="gpt-3.5-turbo",
            object="chat.completion",
            system_fingerprint=None,
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )

        with patch.object(
            monitor_mode_guardrail, "_create_background_noma_check"
        ) as mock_create_background:
            result = await monitor_mode_guardrail.async_post_call_success_hook(
                data=mock_request_data,
                user_api_key_dict=mock_user_api_key_dict,
                response=response,
            )

            assert result == response
            mock_create_background.assert_called_once()


class TestNomaImageProcessing:
    """Test image processing functionality for multimodal content"""

    @pytest.fixture
    def noma_guardrail(self):
        """Create a NomaGuardrail instance for testing"""
        return NomaGuardrail(
            api_key="test-api-key",
            api_base="https://api.test.noma.security/",
            application_id="test-app",
            monitor_mode=False,
            block_failures=True,
            guardrail_name="test-noma-guardrail",
            event_hook="pre_call",
            default_on=True,
        )

    @pytest.fixture
    def mock_user_api_key_dict(self):
        """Create a mock UserAPIKeyAuth object"""
        return UserAPIKeyAuth(
            user_id="test-user-id",
            user_email="test@example.com",
            key_name="test-key",
            api_key="test-api-key",
            permissions={},
            models=[],
            spend=0.0,
            metadata={},
        )

    def test_extract_user_message_with_image_url(self):
        """User message with only image_url becomes a single input_image content item."""
        from litellm.completion_extras.litellm_responses_transformation.transformation import (
            LiteLLMResponsesTransformationHandler,
        )

        handler = LiteLLMResponsesTransformationHandler()
        messages: list[AllMessageValues] = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {"url": "https://example.com/image.jpg"},
                    }
                ],
            }
        ]

        input_items, _ = handler.convert_chat_completion_messages_to_responses_api(
            messages
        )

        assert len(input_items) == 1
        message = input_items[0]["content"]
        assert message is not None
        assert len(message) == 1
        assert message[0]["type"] == "input_image"
        assert message[0]["image_url"] == "https://example.com/image.jpg"

    def test_extract_user_message_with_mixed_content(self):
        """User message with text + image becomes input_text then input_image in content list."""
        from litellm.completion_extras.litellm_responses_transformation.transformation import (
            LiteLLMResponsesTransformationHandler,
        )

        handler = LiteLLMResponsesTransformationHandler()
        messages: list[AllMessageValues] = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "What's in this image?",
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": "https://example.com/image.jpg"},
                    },
                ],
            }
        ]

        input_items, _ = handler.convert_chat_completion_messages_to_responses_api(
            messages
        )

        # Match the original assertions: `message` is the content list
        assert len(input_items) == 1
        message = input_items[0]["content"]
        assert message is not None
        assert len(message) == 2
        # First item should be text
        assert message[0]["type"] == "input_text"
        assert message[0]["text"] == "What's in this image?"
        # Second item should be image
        assert message[1]["type"] == "input_image"
        assert message[1]["image_url"] == "https://example.com/image.jpg"

    def test_extract_user_message_with_multiple_images(self):
        """User message with multiple images becomes multiple input_image items."""
        from litellm.completion_extras.litellm_responses_transformation.transformation import (
            LiteLLMResponsesTransformationHandler,
        )

        handler = LiteLLMResponsesTransformationHandler()

        messages: list[AllMessageValues] = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": "Compare these images",
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": "https://example.com/image1.jpg"},
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": "https://example.com/image2.jpg"},
                    },
                ],
            }
        ]

        input_items, _ = handler.convert_chat_completion_messages_to_responses_api(
            messages
        )

        # Match the original assertions
        assert len(input_items) == 1
        message = input_items[0]["content"]
        assert message is not None
        assert len(message) == 3
        assert message[0]["type"] == "input_text"
        assert message[1]["type"] == "input_image"
        assert message[1]["image_url"] == "https://example.com/image1.jpg"
        assert message[2]["type"] == "input_image"
        assert message[2]["image_url"] == "https://example.com/image2.jpg"

    @pytest.mark.asyncio
    async def test_pre_call_hook_with_image_content(
        self, noma_guardrail, mock_user_api_key_dict
    ):
        """Test pre-call hook with image content"""
        request_data = {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {"url": "https://example.com/test-image.jpg"},
                        }
                    ],
                }
            ],
            "litellm_call_id": "test-call-id",
            "metadata": {"requester_ip_address": "192.168.1.1"},
        }

        # Mock Noma API response for image content
        noma_response = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "harmfulContent": {
                            "result": False,
                            "probability": 0.1,
                            "status": "SUCCESS",
                        }
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ) as mock_post:
            result = await noma_guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=request_data,
                call_type="completion",
            )

            assert result == request_data
            mock_post.assert_called_once()

            # Verify the API call payload includes image
            call_args = mock_post.call_args
            payload = call_args[1]["json"]
            assert "input" in payload
            assert len(payload["input"]) > 0
            assert payload["input"][0]["role"] == "user"
            assert "content" in payload["input"][0]

    @pytest.mark.asyncio
    async def test_pre_call_hook_with_mixed_content(
        self, noma_guardrail, mock_user_api_key_dict
    ):
        """Test pre-call hook with mixed text and image content"""
        request_data = {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Analyze this image for harmful content",
                        },
                        {
                            "type": "image_url",
                            "image_url": {"url": "https://example.com/test-image.jpg"},
                        },
                    ],
                }
            ],
            "litellm_call_id": "test-call-id",
        }

        # Mock Noma API response
        noma_response = {
            "aggregatedScanResult": False,
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "harmfulContent": {
                            "result": False,
                            "probability": 0.05,
                            "status": "SUCCESS",
                        }
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ) as mock_post:
            result = await noma_guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=request_data,
                call_type="completion",
            )

            assert result == request_data
            mock_post.assert_called_once()

    @pytest.mark.asyncio
    async def test_image_content_blocked(self, noma_guardrail, mock_user_api_key_dict):
        """Test that image content can be blocked by Noma"""
        request_data = {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": "https://example.com/inappropriate-image.jpg"
                            },
                        }
                    ],
                }
            ],
            "litellm_call_id": "test-call-id",
        }

        # Mock Noma API response indicating harmful content in image
        noma_response = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "harmfulContent": {
                            "result": True,
                            "probability": 0.95,
                            "status": "SUCCESS",
                        }
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        from litellm.proxy.guardrails.guardrail_hooks.noma.noma import (
            NomaBlockedMessage,
        )

        with patch.object(
            noma_guardrail.async_handler, "post", return_value=mock_response
        ):
            with pytest.raises(NomaBlockedMessage) as exc_info:
                await noma_guardrail.async_pre_call_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    cache=MagicMock(),
                    data=request_data,
                    call_type="completion",
                )

            assert exc_info.value.status_code == 400

    @pytest.mark.asyncio
    async def test_image_with_base64_data(self, noma_guardrail, mock_user_api_key_dict):
        """Test extracting image with base64 data URL"""
        from litellm.completion_extras.litellm_responses_transformation.transformation import (
            LiteLLMResponsesTransformationHandler,
        )

        data = {
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": "data:image/jpeg;base64,/9j/4AAQSkZJRg..."
                            },
                        }
                    ],
                }
            ]
        }

        handler = LiteLLMResponsesTransformationHandler()
        messages = cast(list[AllMessageValues], data["messages"])

        input_items, _ = handler.convert_chat_completion_messages_to_responses_api(
            messages
        )

        assert len(input_items) == 1
        message = input_items[0]["content"]
        assert message is not None
        assert len(message) == 1
        assert message[0]["type"] == "input_image"
        assert message[0]["image_url"].startswith("data:image/jpeg;base64,")


class TestIntegration:
    @pytest.mark.asyncio
    async def test_full_guardrail_flow(self):
        """Test full guardrail flow with multiple hooks"""
        with patch.dict(
            os.environ,
            {
                "NOMA_API_KEY": "test-api-key",
                "NOMA_API_BASE": "https://api.test.noma.security/",
            },
        ):
            init_guardrails_v2(
                all_guardrails=[
                    {
                        "guardrail_name": "noma-pre-guard",
                        "litellm_params": {
                            "guardrail": "noma",
                            "mode": "pre_call",
                            "application_id": "test-app",
                        },
                    },
                    {
                        "guardrail_name": "noma-post-guard",
                        "litellm_params": {
                            "guardrail": "noma",
                            "mode": "post_call",
                            "application_id": "test-app",
                        },
                    },
                ],
                config_file_path="",
            )

            custom_loggers = (
                litellm.logging_callback_manager.get_custom_loggers_for_type(
                    callback_type=litellm.integrations.custom_guardrail.CustomGuardrail
                )
            )
            assert len(custom_loggers) >= 2


class TestNomaAnonymizationConfiguration:
    """Test anonymize_input configuration parameter"""

    def test_init_with_anonymize_input_env_var(self):
        """Test initialization with NOMA_ANONYMIZE_INPUT environment variable"""
        with patch.dict(
            os.environ,
            {
                "NOMA_ANONYMIZE_INPUT": "true",
            },
        ):
            guardrail = NomaGuardrail()
            assert guardrail.anonymize_input is True

        with patch.dict(
            os.environ,
            {
                "NOMA_ANONYMIZE_INPUT": "false",
            },
        ):
            guardrail = NomaGuardrail()
            assert guardrail.anonymize_input is False

    def test_init_with_anonymize_input_default(self):
        """Test default value for anonymize_input"""
        guardrail = NomaGuardrail()
        assert guardrail.anonymize_input is False

    def test_init_with_anonymize_input_param_override_env(self):
        """Test that constructor param overrides environment variable"""
        with patch.dict(
            os.environ,
            {
                "NOMA_ANONYMIZE_INPUT": "true",
            },
        ):
            guardrail = NomaGuardrail(anonymize_input=False)
            assert guardrail.anonymize_input is False

    def test_initialize_guardrail_with_anonymize_input(self):
        """Test the initialize_guardrail function with anonymize_input"""
        from litellm.types.guardrails import Guardrail, LitellmParams

        litellm_params = LitellmParams(
            guardrail="noma",
            mode="pre_call",
            api_key="test-key",
            anonymize_input=True,
        )

        guardrail = Guardrail(
            guardrail_name="test-guardrail",
            litellm_params=litellm_params,
        )

        with patch("litellm.logging_callback_manager.add_litellm_callback"):
            result = initialize_guardrail(litellm_params, guardrail)
            assert result.anonymize_input is True


class TestNomaAnonymizationLogic:
    """Test the anonymization logic helper methods"""

    @pytest.fixture
    def anonymize_guardrail(self):
        """Create a guardrail with anonymize_input enabled"""
        return NomaGuardrail(
            api_key="test-api-key",
            anonymize_input=True,
            monitor_mode=False,
            block_failures=True,
        )

    def test_is_result_true(self, anonymize_guardrail):
        """Test _is_result_true helper method"""
        assert anonymize_guardrail._is_result_true({"result": True}) is True
        assert anonymize_guardrail._is_result_true({"result": False}) is False
        assert anonymize_guardrail._is_result_true({"other": True}) is False
        assert anonymize_guardrail._is_result_true(None) is False
        assert anonymize_guardrail._is_result_true({}) is False
        assert anonymize_guardrail._is_result_true("not a dict") is False

    def test_should_only_data_detector_failed_true(self, anonymize_guardrail):
        """Test _should_only_sensitive_data_failed when only sensitive data detector triggered"""
        classification = {
            "sensitiveData": {
                "PII": {"result": True, "status": "SUCCESS"},
                "PCI": {"result": True, "status": "SUCCESS"},
                "secrets": {"result": False, "probability": 0, "status": "SUCCESS"},
            },
            "harmfulContent": {"result": False, "status": "SUCCESS"},
            "maliciousIntent": {"result": False, "status": "SUCCESS"},
            "code": {"result": False, "status": "SUCCESS"},
        }

        result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
        assert result is True

    def test_should_only_data_detector_failed_false_other_detectors(
        self, anonymize_guardrail
    ):
        """Test _should_only_sensitive_data_failed when other detectors also triggered"""
        classification = {
            "sensitiveData": {
                "PII": {"result": True, "status": "SUCCESS"},
            },
            "harmfulContent": {
                "result": True,
                "status": "SUCCESS",
            },  # This should cause False
            "maliciousIntent": {"result": False, "status": "SUCCESS"},
        }

        result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
        assert result is False

    def test_should_only_data_detector_failed_false_no_data_detected(
        self, anonymize_guardrail
    ):
        """Test _should_only_sensitive_data_failed when no sensitive data detected"""
        classification = {
            "sensitiveData": {
                "PII": {"result": False, "status": "SUCCESS"},
                "PCI": {"result": False, "status": "SUCCESS"},
            },
            "harmfulContent": {"result": False, "status": "SUCCESS"},
            "maliciousIntent": {"result": False, "status": "SUCCESS"},
        }

        result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
        assert result is False

    def test_should_only_data_detector_failed_with_nested_detectors(
        self, anonymize_guardrail
    ):
        """Test _should_only_sensitive_data_failed with nested detectors like topicDetector"""
        classification = {
            "sensitiveData": {
                "PII": {"result": True, "status": "SUCCESS"},
            },
            "customLlm": {
                "topic1": {
                    "result": True,
                    "status": "SUCCESS",
                },  # This should cause False
            },
            "harmfulContent": {"result": False, "status": "SUCCESS"},
        }

        result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
        assert result is False

    def test_extract_anonymized_content_user(self, anonymize_guardrail):
        """Test _extract_anonymized_content for user messages"""
        response_json = {
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "anonymizedContent": {
                            "anonymized": "My email is ******* and phone is *******"
                        }
                    },
                }
            ]
        }

        result = anonymize_guardrail._extract_anonymized_content(response_json, "user")
        assert result == "My email is ******* and phone is *******"

    def test_extract_anonymized_content_assistant(self, anonymize_guardrail):
        """Test _extract_anonymized_content for assistant messages"""
        response_json = {
            "scanResult": [
                {
                    "role": "assistant",
                    "type": "message",
                    "results": {
                        "anonymizedContent": {
                            "anonymized": "I can't help with that request."
                        }
                    },
                }
            ]
        }

        result = anonymize_guardrail._extract_anonymized_content(
            response_json, "assistant"
        )
        assert result == "I can't help with that request."

    def test_extract_anonymized_content_missing(self, anonymize_guardrail):
        """Test _extract_anonymized_content when anonymized content is missing"""
        response_json = {
            "scanResult": [{"role": "user", "type": "message", "results": {}}]
        }

        result = anonymize_guardrail._extract_anonymized_content(response_json, "user")
        assert result == ""

    def test_should_anonymize_verdict_true(self, anonymize_guardrail):
        """Test _should_anonymize when aggregatedScanResult is False (safe)"""
        response_json = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [{"role": "user", "type": "message", "results": {}}],
        }

        result = anonymize_guardrail._should_anonymize(response_json, "user")
        assert result is True

    def test_should_anonymize_verdict_false_only_sensitive(self, anonymize_guardrail):
        """Test _should_anonymize when aggregatedScanResult is True but only sensitive data detector triggered"""
        response_json = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "sensitiveData": {"PCI": {"result": True, "status": "SUCCESS"}},
                        "harmfulContent": {"result": False, "status": "SUCCESS"},
                    },
                }
            ],
        }

        result = anonymize_guardrail._should_anonymize(response_json, "user")
        assert result is True

    def test_should_anonymize_verdict_false_other_detectors(self, anonymize_guardrail):
        """Test _should_anonymize when aggregatedScanResult is True and other detectors triggered"""
        response_json = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "sensitiveData": {"PCI": {"result": True, "status": "SUCCESS"}},
                        "harmfulContent": {"result": True, "status": "SUCCESS"},
                    },
                }
            ],
        }

        result = anonymize_guardrail._should_anonymize(response_json, "user")
        assert result is False

    def test_should_anonymize_monitor_mode(self):
        """Test _should_anonymize in monitor mode (should never anonymize)"""
        guardrail = NomaGuardrail(
            anonymize_input=True,
            monitor_mode=True,
        )

        response_json = {
            "aggregatedScanResult": False,
            "scanResult": [{"role": "user", "type": "message", "results": {}}],
        }
        result = guardrail._should_anonymize(response_json, "user")
        assert result is False

    def test_should_anonymize_disabled(self):
        """Test _should_anonymize when anonymize_input is disabled"""
        guardrail = NomaGuardrail(
            anonymize_input=False,
            monitor_mode=False,
        )

        response_json = {
            "aggregatedScanResult": False,
            "scanResult": [{"role": "user", "type": "message", "results": {}}],
        }
        result = guardrail._should_anonymize(response_json, "user")
        assert result is False

    def test_replace_user_message_content(self, anonymize_guardrail):
        """Test _replace_user_message_content"""
        request_data = {
            "messages": [
                {"role": "system", "content": "System prompt"},
                {"role": "user", "content": "My email is test@example.com"},
                {"role": "assistant", "content": "I can help you"},
                {"role": "user", "content": "My phone is 123-456-7890"},
            ]
        }

        anonymize_guardrail._replace_user_message_content(
            request_data, "My phone is *******"
        )

        # Should replace the last user message
        assert request_data["messages"][-1]["content"] == "My phone is *******"
        assert (
            request_data["messages"][1]["content"] == "My email is test@example.com"
        )  # Unchanged

    def test_replace_llm_response_content(self, anonymize_guardrail):
        """Test _replace_llm_response_content"""
        response = ModelResponse(
            id="test-id",
            choices=[
                Choices(
                    finish_reason="stop",
                    index=0,
                    message=Message(
                        content="Your email is test@example.com", role="assistant"
                    ),
                )
            ],
            created=1234567890,
            model="gpt-3.5-turbo",
            object="chat.completion",
            system_fingerprint=None,
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )

        anonymize_guardrail._replace_llm_response_content(
            response, "Your email is *******"
        )

        assert response.choices[0].message.content == "Your email is *******"


class TestNomaAnonymizationFlow:
    """Test full anonymization flow with real Noma response objects"""

    @pytest.fixture
    def anonymize_guardrail(self):
        """Create a guardrail with anonymize_input enabled"""
        return NomaGuardrail(
            api_key="test-api-key",
            api_base="https://api.test.noma.security/",
            application_id="test-app",
            anonymize_input=True,
            monitor_mode=False,
            block_failures=True,
            guardrail_name="test-noma-guardrail",
            event_hook="pre_call",
            default_on=True,
        )

    @pytest.fixture
    def mock_user_api_key_dict(self):
        """Create a mock UserAPIKeyAuth object"""
        return UserAPIKeyAuth(
            user_id="test-user-id",
            user_email="test@example.com",
            key_name="test-key",
            api_key="test-api-key",
            permissions={},
            models=[],
            spend=0.0,
            metadata={},
        )

    @pytest.mark.asyncio
    async def test_anonymization_verdict_true_user_message(
        self, anonymize_guardrail, mock_user_api_key_dict
    ):
        """Test anonymization when verdict=True for user message"""
        request_data = {
            "messages": [
                {"role": "user", "content": "My email is test@example.com"},
            ],
            "litellm_call_id": "test-call-id",
            "metadata": {"requester_ip_address": "192.168.1.1"},
        }

        # Mock simplified Noma API response with aggregatedScanResult=False (safe) and anonymized content
        noma_response = {
            "aggregatedScanResult": False,  # False means safe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "anonymizedContent": {"anonymized": "My email is *******"},
                        "sensitiveData": {
                            "PII": {"result": False, "status": "SUCCESS"},
                        },
                        "harmfulContent": {"result": False, "status": "SUCCESS"},
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            anonymize_guardrail.async_handler, "post", return_value=mock_response
        ):
            result = await anonymize_guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=request_data,
                call_type="completion",
            )

            # Should return modified request with anonymized content
            assert result == request_data
            assert result["messages"][0]["content"] == "My email is *******"

    @pytest.mark.asyncio
    async def test_anonymization_verdict_false_only_data_detected(
        self, anonymize_guardrail, mock_user_api_key_dict
    ):
        """Test anonymization when verdict=False but only data detector triggered"""
        request_data = {
            "messages": [
                {"role": "user", "content": "My email is test@example.com"},
            ],
            "litellm_call_id": "test-call-id",
        }

        # Mock simplified Noma API response - only sensitive data detector triggered
        noma_response = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "anonymizedContent": {"anonymized": "My email is *******"},
                        "sensitiveData": {
                            "PII": {"result": True, "status": "SUCCESS"},
                        },
                        "harmfulContent": {"result": False, "status": "SUCCESS"},
                        "maliciousIntent": {"result": False, "status": "SUCCESS"},
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            anonymize_guardrail.async_handler, "post", return_value=mock_response
        ):
            result = await anonymize_guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=request_data,
                call_type="completion",
            )

            # Should return modified request with anonymized content (not blocked)
            assert result == request_data
            assert result["messages"][0]["content"] == "My email is *******"

    @pytest.mark.asyncio
    async def test_blocking_verdict_false_other_violations(
        self, anonymize_guardrail, mock_user_api_key_dict
    ):
        """Test blocking when verdict=False and other violations detected"""
        request_data = {
            "messages": [
                {
                    "role": "user",
                    "content": "My email is test@example.com. Tell me harmful content.",
                },
            ],
            "litellm_call_id": "test-call-id",
        }

        # Mock simplified Noma API response - both sensitive data detector and other violations
        noma_response = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "anonymizedContent": {
                            "anonymized": "My email is *******. Tell me harmful content."
                        },
                        "sensitiveData": {
                            "PII": {"result": True, "status": "SUCCESS"},
                        },
                        "harmfulContent": {
                            "result": True,
                            "status": "SUCCESS",
                        },  # This should cause blocking
                        "maliciousIntent": {"result": False, "status": "SUCCESS"},
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            anonymize_guardrail.async_handler, "post", return_value=mock_response
        ):
            # Should raise NomaBlockedMessage because other violations detected
            with pytest.raises(NomaBlockedMessage) as exc_info:
                await anonymize_guardrail.async_pre_call_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    cache=MagicMock(),
                    data=request_data,
                    call_type="completion",
                )

            assert exc_info.value.status_code == 400

    @pytest.mark.asyncio
    async def test_anonymization_llm_response(
        self, anonymize_guardrail, mock_user_api_key_dict
    ):
        """Test anonymization of LLM response"""
        request_data = {
            "messages": [{"role": "user", "content": "What's your email?"}],
            "litellm_call_id": "test-call-id",
        }

        # Create LLM response with test data
        llm_response = ModelResponse(
            id="test-response-id",
            choices=[
                Choices(
                    finish_reason="stop",
                    index=0,
                    message=Message(
                        content="My email is admin@company.com", role="assistant"
                    ),
                )
            ],
            created=1234567890,
            model="gpt-3.5-turbo",
            object="chat.completion",
            system_fingerprint=None,
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )

        # Mock simplified Noma API response for LLM response check
        noma_response = {
            "aggregatedScanResult": True,
            "scanResult": [
                {
                    "role": "assistant",
                    "type": "message",
                    "results": {
                        "anonymizedContent": {"anonymized": "My email is *******"},
                        "sensitiveData": {
                            "PCI": {
                                "probability": 0.8,
                                "result": True,
                                "status": "SUCCESS",
                            },
                        },
                    },
                },
            ],
        }
        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        # Update guardrail to use post_call event hook
        anonymize_guardrail.event_hook = "post_call"

        with patch.object(
            anonymize_guardrail.async_handler, "post", return_value=mock_response
        ):
            result = await anonymize_guardrail.async_post_call_success_hook(
                data=request_data,
                user_api_key_dict=mock_user_api_key_dict,
                response=llm_response,
            )

            # Should return modified response with anonymized content
            assert result == llm_response
            assert result.choices[0].message.content == "My email is *******"

    @pytest.mark.asyncio
    async def test_no_anonymization_when_disabled(self, mock_user_api_key_dict):
        """Test that no anonymization occurs when anonymize_input=False"""
        guardrail = NomaGuardrail(
            api_key="test-api-key",
            anonymize_input=False,  # Disabled
            monitor_mode=False,
            block_failures=True,
        )

        request_data = {
            "messages": [
                {"role": "user", "content": "My email is test@example.com"},
            ],
        }

        noma_response = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "anonymizedContent": {"anonymized": "My email is *******"},
                        "sensitiveData": {
                            "PII": {"result": True, "status": "SUCCESS"},
                        },
                        "harmfulContent": {"result": False, "status": "SUCCESS"},
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        with patch.object(guardrail.async_handler, "post", return_value=mock_response):
            # Should raise NomaBlockedMessage because anonymization is disabled
            with pytest.raises(NomaBlockedMessage):
                await guardrail.async_pre_call_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    cache=MagicMock(),
                    data=request_data,
                    call_type="completion",
                )

    @pytest.mark.asyncio
    async def test_no_anonymization_in_monitor_mode(self, mock_user_api_key_dict):
        """Test that no anonymization occurs in monitor mode"""
        guardrail = NomaGuardrail(
            api_key="test-api-key",
            anonymize_input=True,
            monitor_mode=True,  # Monitor mode
            block_failures=True,
        )

        request_data = {
            "messages": [
                {"role": "user", "content": "My email is test@example.com"},
            ],
        }

        with patch.object(
            guardrail, "_create_background_noma_check"
        ) as mock_create_background:
            result = await guardrail.async_pre_call_hook(
                user_api_key_dict=mock_user_api_key_dict,
                cache=MagicMock(),
                data=request_data,
                call_type="completion",
            )

            # Should return original data unchanged
            assert result == request_data
            assert (
                request_data["messages"][0]["content"] == "My email is test@example.com"
            )
            mock_create_background.assert_called_once()

    @pytest.mark.asyncio
    async def test_anonymization_no_anonymized_content_available(
        self, anonymize_guardrail, mock_user_api_key_dict
    ):
        """Test behavior when anonymized content is not available"""
        request_data = {
            "messages": [
                {"role": "user", "content": "My email is test@example.com"},
            ],
        }

        noma_response = {
            "aggregatedScanResult": True,  # True means unsafe
            "scanResult": [
                {
                    "role": "user",
                    "type": "message",
                    "results": {
                        "sensitiveData": {
                            "PII": {"result": True, "status": "SUCCESS"},
                        },
                        "harmfulContent": {"result": False, "status": "SUCCESS"},
                    },
                }
            ],
        }

        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        with patch.object(
            anonymize_guardrail.async_handler, "post", return_value=mock_response
        ):
            # Should raise NomaBlockedMessage because no anonymized content available
            with pytest.raises(NomaBlockedMessage):
                await anonymize_guardrail.async_pre_call_hook(
                    user_api_key_dict=mock_user_api_key_dict,
                    cache=MagicMock(),
                    data=request_data,
                    call_type="completion",
                )

    @pytest.mark.asyncio
    async def test_anonymization_llm_response_no_anonymized_content_available(
        self, anonymize_guardrail, mock_user_api_key_dict
    ):
        """Test behavior when LLM response has no anonymized content available"""
        request_data = {
            "messages": [{"role": "user", "content": "What's your email?"}],
            "litellm_call_id": "test-call-id",
        }

        # Create LLM response with test data
        llm_response = ModelResponse(
            id="test-response-id",
            choices=[
                Choices(
                    finish_reason="stop",
                    index=0,
                    message=Message(
                        content="My email is admin@company.com", role="assistant"
                    ),
                )
            ],
            created=1234567890,
            model="gpt-3.5-turbo",
            object="chat.completion",
            system_fingerprint=None,
            usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
        )

        # Mock Noma API response with no anonymized content available
        noma_response = {
            "aggregatedScanResult": True,
            "scanResult": [
                {
                    "role": "assistant",
                    "type": "message",
                    "results": {
                        "sensitiveData": {
                            "PCI": {
                                "probability": 0.8,
                                "result": True,
                                "status": "SUCCESS",
                            },
                        },
                    },
                },
            ],
        }
        mock_response = MagicMock()
        mock_response.json.return_value = noma_response
        mock_response.raise_for_status = MagicMock()

        # Update guardrail to use post_call event hook
        anonymize_guardrail.event_hook = "post_call"

        with patch.object(
            anonymize_guardrail.async_handler, "post", return_value=mock_response
        ):
            # Should raise NomaBlockedMessage because no anonymized content available for LLM response
            with pytest.raises(NomaBlockedMessage):
                await anonymize_guardrail.async_post_call_success_hook(
                    data=request_data,
                    user_api_key_dict=mock_user_api_key_dict,
                    response=llm_response,
                )
