"""
Test OpenAI Embeddings Guardrail Translation Handler
"""

from unittest.mock import AsyncMock, MagicMock

import pytest

from litellm.llms.openai.embeddings.guardrail_translation.handler import (
    OpenAIEmbeddingsHandler,
)
from litellm.types.utils import CallTypes


@pytest.mark.asyncio
async def test_embeddings_handler_string_input():
    """Test embeddings handler with single string input"""
    handler = OpenAIEmbeddingsHandler()

    # Mock guardrail
    mock_guardrail = MagicMock()
    mock_guardrail.apply_guardrail = AsyncMock(
        return_value={"texts": ["processed text"]}
    )

    data = {"input": "Hello, world!", "model": "text-embedding-3-small"}

    result = await handler.process_input_messages(
        data=data,
        guardrail_to_apply=mock_guardrail,
    )

    # Verify guardrail was called with correct inputs
    mock_guardrail.apply_guardrail.assert_called_once()
    call_args = mock_guardrail.apply_guardrail.call_args
    assert call_args.kwargs["inputs"]["texts"] == ["Hello, world!"]
    assert call_args.kwargs["inputs"]["model"] == "text-embedding-3-small"

    # Verify result
    assert result["input"] == "processed text"


@pytest.mark.asyncio
async def test_embeddings_handler_list_of_strings_input():
    """Test embeddings handler with list of strings input"""
    handler = OpenAIEmbeddingsHandler()

    # Mock guardrail
    mock_guardrail = MagicMock()
    mock_guardrail.apply_guardrail = AsyncMock(
        return_value={"texts": ["processed text 1", "processed text 2"]}
    )

    data = {
        "input": ["Hello, world!", "How are you?"],
        "model": "text-embedding-3-small",
    }

    result = await handler.process_input_messages(
        data=data,
        guardrail_to_apply=mock_guardrail,
    )

    # Verify guardrail was called with correct inputs
    mock_guardrail.apply_guardrail.assert_called_once()
    call_args = mock_guardrail.apply_guardrail.call_args
    assert call_args.kwargs["inputs"]["texts"] == ["Hello, world!", "How are you?"]

    # Verify result
    assert result["input"] == ["processed text 1", "processed text 2"]


def test_embeddings_guardrail_translation_mappings():
    """Test that embeddings handler is registered for correct call types"""
    from litellm.llms.openai.embeddings.guardrail_translation import (
        guardrail_translation_mappings,
    )

    assert CallTypes.embedding in guardrail_translation_mappings
    assert CallTypes.aembedding in guardrail_translation_mappings
    assert (
        guardrail_translation_mappings[CallTypes.embedding] == OpenAIEmbeddingsHandler
    )
    assert (
        guardrail_translation_mappings[CallTypes.aembedding] == OpenAIEmbeddingsHandler
    )
