from unittest.mock import Mock

from verse_sdk.integrations.google.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_finish_reason(self):
        mock_candidate = Mock()
        mock_candidate.finish_reason = "STOP"

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]

        assert get_finish_reason(mock_response) == "stop"

    def test_response_with_max_tokens_finish_reason(self):
        mock_candidate = Mock()
        mock_candidate.finish_reason = "MAX_TOKENS"

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]

        assert get_finish_reason(mock_response) == "length"

    def test_response_with_safety_finish_reason(self):
        mock_candidate = Mock()
        mock_candidate.finish_reason = "SAFETY"

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]

        assert get_finish_reason(mock_response) == "content_filter"

    def test_response_with_other_finish_reason(self):
        mock_candidate = Mock()
        mock_candidate.finish_reason = "OTHER"

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]

        assert get_finish_reason(mock_response) == "other"

    def test_response_without_candidates(self):
        mock_response = Mock()
        mock_response.candidates = []
        assert get_finish_reason(mock_response) is None

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


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

        mock_content = Mock()
        mock_content.parts = [mock_part]
        mock_content.role = "model"

        mock_candidate = Mock()
        mock_candidate.content = mock_content

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]

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

    def test_response_with_function_call(self):
        mock_function_call = Mock()
        mock_function_call.name = "get_weather"
        mock_function_call.args = {"location": "San Francisco"}

        mock_part = Mock()
        del mock_part.text
        mock_part.function_call = mock_function_call

        mock_content = Mock()
        mock_content.parts = [mock_part]
        mock_content.role = "model"

        mock_candidate = Mock()
        mock_candidate.content = mock_content

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]

        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]["function"]["name"] == "get_weather"

    def test_response_without_candidates(self):
        mock_response = Mock()
        mock_response.candidates = []
        assert get_completion_choices(mock_response) == []

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

    def test_response_with_multiple_text_parts(self):
        mock_part1 = Mock()
        mock_part1.text = "Hello, "

        mock_part2 = Mock()
        mock_part2.text = "world!"

        mock_content = Mock()
        mock_content.parts = [mock_part1, mock_part2]
        mock_content.role = "model"

        mock_candidate = Mock()
        mock_candidate.content = mock_content

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]

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


class TestGetUsageFromResponse:
    def test_response_with_usage(self):
        mock_usage = Mock()
        mock_usage.prompt_token_count = 100
        mock_usage.response_token_count = 50
        mock_usage.total_token_count = 150

        mock_response = Mock()
        mock_response.usage_metadata = 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.prompt_token_count = 100
        del mock_usage.response_token_count
        del mock_usage.total_token_count

        mock_response = Mock()
        mock_response.usage_metadata = 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_metadata
        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):
        mock_config = Mock()
        mock_config.max_output_tokens = 1024
        mock_config.temperature = 0.7
        mock_config.top_p = 0.9
        mock_config.top_k = 40

        kwargs = {
            "config": mock_config,
            "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 False

    def test_missing_parameters(self):
        mock_config = Mock()
        mock_config.max_output_tokens = 512
        del mock_config.temperature
        del mock_config.top_p
        del mock_config.top_k

        kwargs = {"config": mock_config}

        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_string_content(self):
        contents = "Hello, how are you?"
        formatted = format_messages_for_tracing(contents)
        assert len(formatted) == 1
        assert formatted[0]["role"] == "user"
        assert formatted[0]["content"] == "Hello, how are you?"

    def test_list_of_strings(self):
        contents = ["Hello", "How are you?"]
        formatted = format_messages_for_tracing(contents)
        assert len(formatted) == 2
        assert formatted[0]["content"] == "Hello"
        assert formatted[1]["content"] == "How are you?"

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

    def test_dict_with_function_call(self):
        contents = [
            {
                "role": "model",
                "parts": [{"function_call": {"name": "get_weather"}}],
            }
        ]
        formatted = format_messages_for_tracing(contents)
        assert len(formatted) == 1
        assert "[Function Call: get_weather]" in formatted[0]["content"]

    def test_object_with_role_and_parts(self):
        mock_part = Mock()
        mock_part.text = "Hello from object"

        mock_content = Mock()
        mock_content.role = "user"
        mock_content.parts = [mock_part]

        formatted = format_messages_for_tracing(mock_content)
        assert len(formatted) == 1
        assert formatted[0]["role"] == "user"
        assert formatted[0]["content"] == "Hello from object"

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

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


class TestGetStreamText:
    def test_chunk_with_text_attribute(self):
        mock_chunk = Mock()
        mock_chunk.text = "Hello"
        assert get_stream_text(mock_chunk) == "Hello"

    def test_chunk_with_candidates_and_parts(self):
        mock_part = Mock()
        mock_part.text = "World"

        mock_content = Mock()
        mock_content.parts = [mock_part]

        mock_candidate = Mock()
        mock_candidate.content = mock_content

        mock_chunk = Mock()
        del mock_chunk.text
        mock_chunk.candidates = [mock_candidate]

        assert get_stream_text(mock_chunk) == "World"

    def test_chunk_with_multiple_parts(self):
        mock_part1 = Mock()
        mock_part1.text = "Hello "

        mock_part2 = Mock()
        mock_part2.text = "World"

        mock_content = Mock()
        mock_content.parts = [mock_part1, mock_part2]

        mock_candidate = Mock()
        mock_candidate.content = mock_content

        mock_chunk = Mock()
        del mock_chunk.text
        mock_chunk.candidates = [mock_candidate]

        assert get_stream_text(mock_chunk) == "Hello World"

    def test_chunk_without_text(self):
        mock_chunk = Mock()
        del mock_chunk.text
        mock_chunk.candidates = []
        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_part = Mock()
        mock_part.text = "The weather is sunny."

        mock_content = Mock()
        mock_content.parts = [mock_part]
        mock_content.role = "model"

        mock_candidate = Mock()
        mock_candidate.content = mock_content
        mock_candidate.finish_reason = "STOP"

        mock_usage = Mock()
        mock_usage.prompt_token_count = 50
        mock_usage.response_token_count = 25
        mock_usage.total_token_count = 75

        mock_response = Mock()
        mock_response.candidates = [mock_candidate]
        mock_response.usage_metadata = mock_usage

        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 == "stop"

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

        mock_content = Mock()
        mock_content.parts = [mock_part]

        mock_candidate = Mock()
        mock_candidate.content = mock_content

        mock_chunk = Mock()
        del mock_chunk.text
        mock_chunk.candidates = [mock_candidate]

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