from unittest.mock import Mock

from verse_sdk.integrations.anthropic.utils import (
    format_messages_for_tracing,
    get_completion_choices,
    get_finish_reason,
    get_request_parameters,
    get_stream_text,
    get_usage_from_response,
)


class TestGetFinishReason:
    def test_response_with_stop_reason(self):
        mock_response = Mock()
        mock_response.stop_reason = "end_turn"
        assert get_finish_reason(mock_response) == "end_turn"

    def test_response_without_stop_reason(self):
        mock_response = Mock()
        del mock_response.stop_reason
        assert get_finish_reason(mock_response) is None

    def test_none_response(self):
        assert get_finish_reason(None) is None

    def test_response_with_max_tokens_stop_reason(self):
        mock_response = Mock()
        mock_response.stop_reason = "max_tokens"
        assert get_finish_reason(mock_response) == "max_tokens"


class TestGetCompletionChoices:
    def test_response_with_text_content(self):
        mock_content_block = Mock()
        mock_content_block.text = "Hello, world!"
        mock_content_block.type = "text"

        mock_response = Mock()
        mock_response.content = [mock_content_block]
        mock_response.role = "assistant"

        choices = get_completion_choices(mock_response)
        assert len(choices) == 1
        assert choices[0]["message"]["content"] == "Hello, world!"
        assert choices[0]["message"]["role"] == "assistant"
        assert choices[0]["text"] == "Hello, world!"

    def test_response_with_tool_use(self):
        mock_tool_block = Mock()
        mock_tool_block.type = "tool_use"
        mock_tool_block.id = "tool_123"
        mock_tool_block.name = "get_weather"
        mock_tool_block.input = {"location": "San Francisco"}

        mock_response = Mock()
        mock_response.content = [mock_tool_block]
        mock_response.role = "assistant"

        choices = get_completion_choices(mock_response)
        assert len(choices) == 1
        assert choices[0]["message"]["tool_calls"] is not None
        assert len(choices[0]["message"]["tool_calls"]) == 1
        assert choices[0]["message"]["tool_calls"][0]["id"] == "tool_123"
        assert choices[0]["message"]["tool_calls"][0]["function"]["name"] == "get_weather"

        import json
        arguments = choices[0]["message"]["tool_calls"][0]["function"]["arguments"]
        assert isinstance(arguments, str)
        parsed_args = json.loads(arguments)
        assert parsed_args == {"location": "San Francisco"}

    def test_response_without_content(self):
        mock_response = Mock()
        mock_response.content = []
        assert get_completion_choices(mock_response) == []

    def test_none_response(self):
        assert get_completion_choices(None) == []

    def test_response_with_multiple_content_blocks(self):
        mock_text_block = Mock()
        mock_text_block.text = "Here's the weather:"
        mock_text_block.type = "text"

        mock_tool_block = Mock()
        mock_tool_block.type = "tool_use"
        mock_tool_block.id = "tool_456"
        mock_tool_block.name = "get_weather"
        mock_tool_block.input = {"location": "NYC"}

        mock_response = Mock()
        mock_response.content = [mock_text_block, mock_tool_block]
        mock_response.role = "assistant"

        choices = get_completion_choices(mock_response)
        assert len(choices) == 2


class TestGetUsageFromResponse:
    def test_response_with_usage(self):
        mock_usage = Mock()
        mock_usage.input_tokens = 100
        mock_usage.output_tokens = 50

        mock_response = Mock()
        mock_response.usage = mock_usage

        usage = get_usage_from_response(mock_response)
        assert usage["input_tokens"] == 100
        assert usage["output_tokens"] == 50
        assert usage["prompt_tokens"] == 100
        assert usage["completion_tokens"] == 50
        assert usage["total_tokens"] == 150

    def test_response_with_partial_usage(self):
        mock_usage = Mock()
        mock_usage.input_tokens = 100
        del mock_usage.output_tokens

        mock_response = Mock()
        mock_response.usage = mock_usage

        usage = get_usage_from_response(mock_response)
        assert usage["input_tokens"] == 100
        assert usage["output_tokens"] is None
        assert usage["total_tokens"] is None

    def test_response_without_usage(self):
        mock_response = Mock()
        del mock_response.usage
        assert get_usage_from_response(mock_response) == {}

    def test_none_response(self):
        assert get_usage_from_response(None) == {}


class TestGetRequestParameters:
    def test_basic_parameters(self):
        kwargs = {
            "max_tokens": 1024,
            "temperature": 0.7,
            "top_p": 0.9,
            "stream": True,
        }
        params = get_request_parameters(kwargs)
        assert params["max_tokens"] == 1024
        assert params["temperature"] == 0.7
        assert params["top_p"] == 0.9
        assert params["stream"] is True

    def test_missing_parameters(self):
        kwargs = {"max_tokens": 512}
        params = get_request_parameters(kwargs)
        assert params["max_tokens"] == 512
        assert params["temperature"] is None
        assert params["top_p"] is None
        assert params["stream"] is False

    def test_empty_kwargs(self):
        params = get_request_parameters({})
        assert params["max_tokens"] is None
        assert params["temperature"] is None
        assert params["top_p"] is None
        assert params["stream"] is False


class TestFormatMessagesForTracing:
    def test_simple_messages(self):
        messages = [
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi there!"},
        ]
        formatted = format_messages_for_tracing(messages)
        assert len(formatted) == 2
        assert formatted[0]["role"] == "user"
        assert formatted[0]["content"] == "Hello"
        assert formatted[1]["role"] == "assistant"
        assert formatted[1]["content"] == "Hi there!"

    def test_messages_with_content_blocks(self):
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "What's the weather?"},
                ],
            }
        ]
        formatted = format_messages_for_tracing(messages)
        assert len(formatted) == 1
        assert formatted[0]["content"] == "What's the weather?"

    def test_messages_with_tool_results(self):
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "tool_result", "content": "Sunny, 72°F"},
                ],
            }
        ]
        formatted = format_messages_for_tracing(messages)
        assert len(formatted) == 1
        assert "[Tool Result:" in formatted[0]["content"]

    def test_empty_messages(self):
        assert format_messages_for_tracing([]) == []

    def test_messages_with_mixed_content(self):
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "First part"},
                    {"type": "text", "text": "Second part"},
                ],
            }
        ]
        formatted = format_messages_for_tracing(messages)
        assert len(formatted) == 1
        assert "First part\nSecond part" in formatted[0]["content"]


class TestGetStreamText:
    def test_chunk_with_delta_text(self):
        mock_delta = Mock()
        mock_delta.text = "Hello"

        mock_chunk = Mock()
        mock_chunk.delta = mock_delta

        assert get_stream_text(mock_chunk) == "Hello"

    def test_chunk_with_content_block_text(self):
        mock_content_block = Mock()
        mock_content_block.text = "World"

        mock_chunk = Mock()
        del mock_chunk.delta
        mock_chunk.content_block = mock_content_block

        assert get_stream_text(mock_chunk) == "World"

    def test_chunk_with_none_text(self):
        mock_delta = Mock()
        mock_delta.text = None

        mock_chunk = Mock()
        mock_chunk.delta = mock_delta

        assert get_stream_text(mock_chunk) == ""

    def test_chunk_without_text(self):
        mock_chunk = Mock()
        del mock_chunk.delta
        del mock_chunk.content_block

        assert get_stream_text(mock_chunk) == ""

    def test_none_chunk(self):
        assert get_stream_text(None) == ""


class TestIntegrationScenarios:
    def test_complete_response_processing(self):
        mock_content_block = Mock()
        mock_content_block.text = "The weather is sunny."
        mock_content_block.type = "text"

        mock_usage = Mock()
        mock_usage.input_tokens = 50
        mock_usage.output_tokens = 25

        mock_response = Mock()
        mock_response.content = [mock_content_block]
        mock_response.role = "assistant"
        mock_response.stop_reason = "end_turn"
        mock_response.usage = mock_usage
        mock_response.model = "claude-3-5-sonnet-20241022"
        mock_response.id = "msg_123"

        choices = get_completion_choices(mock_response)
        usage = get_usage_from_response(mock_response)
        finish_reason = get_finish_reason(mock_response)

        assert len(choices) == 1
        assert choices[0]["message"]["content"] == "The weather is sunny."
        assert usage["total_tokens"] == 75
        assert finish_reason == "end_turn"

    def test_streaming_chunk_processing(self):
        mock_delta = Mock()
        mock_delta.text = "Streaming"

        mock_chunk = Mock()
        mock_chunk.delta = mock_delta
        mock_chunk.type = "content_block_delta"

        text = get_stream_text(mock_chunk)
        assert text == "Streaming"
