import json
from unittest.mock import Mock

from verse_sdk.contexts.generation import GenerationContext


class TestGenerationContextCompletions:
    def setup_method(self):
        self.mock_span = Mock()
        self.mock_span.get_span_context.return_value.is_valid = True
        self.mock_span.set_attribute = Mock()
        self.mock_span.add_event = Mock()

        self.context = GenerationContext.__new__(GenerationContext)
        self.context._span = self.mock_span
        self.context.set_attributes = Mock(return_value=self.context)

    def test_completions_empty_list(self):
        result = self.context.completions([])
        assert result is self.context
        self.mock_span.set_attribute.assert_not_called()

    def test_completions_single_text_choice(self):
        choice = {"text": "Hello, world!"}
        choices = [choice]

        result = self.context.completions(choices)
        assert result is self.context

        self.mock_span.set_attribute.assert_called_once()
        call_args = self.mock_span.set_attribute.call_args
        assert call_args[0][0] == "gen_ai.output.messages"

        messages = json.loads(call_args[0][1])
        assert len(messages) == 1
        assert messages[0]["role"] == "assistant"
        assert messages[0]["finish_reason"] == "stop"
        assert len(messages[0]["parts"]) == 1
        assert messages[0]["parts"][0]["type"] == "text"
        assert messages[0]["parts"][0]["content"] == "Hello, world!"

    def test_completions_multiple_text_choices(self):
        choices = [
            {"text": "First response"},
            {"text": "Second response"},
            {"text": "Third response"},
        ]

        result = self.context.completions(choices)

        assert result is self.context

        self.mock_span.set_attribute.assert_called_once()
        call_args = self.mock_span.set_attribute.call_args
        assert call_args[0][0] == "gen_ai.output.messages"

        messages = json.loads(call_args[0][1])
        assert len(messages) == 3
        assert messages[0]["parts"][0]["content"] == "First response"
        assert messages[1]["parts"][0]["content"] == "Second response"
        assert messages[2]["parts"][0]["content"] == "Third response"

    def test_completions_with_message_object(self):
        message = {"role": "assistant", "content": "This is a test message"}
        choice = {"message": message}
        choices = [choice]

        self.context.completions(choices)

        call_args = self.mock_span.set_attribute.call_args
        messages = json.loads(call_args[0][1])
        assert messages[0]["role"] == "assistant"
        assert messages[0]["parts"][0]["content"] == "This is a test message"

    def test_completions_with_function_call(self):
        message = {
            "role": "assistant",
            "content": "I'll call a function",
            "function_call": {
                "name": "get_weather",
                "arguments": '{"location": "New York"}',
            },
        }
        choice = {"message": message}
        choices = [choice]

        self.context.completions(choices)

        call_args = self.mock_span.set_attribute.call_args
        messages = json.loads(call_args[0][1])
        assert messages[0]["parts"][0]["type"] == "text"
        assert messages[0]["parts"][0]["content"] == "I'll call a function"
        assert messages[0]["parts"][1]["type"] == "tool_call"
        assert messages[0]["parts"][1]["name"] == "get_weather"
        assert messages[0]["parts"][1]["arguments"] == '{"location": "New York"}'

    def test_completions_with_tool_calls(self):
        message = {
            "role": "assistant",
            "content": "I'll use some tools",
            "tool_calls": [
                {
                    "id": "call_123",
                    "type": "function",
                    "function": {
                        "name": "search_database",
                        "arguments": '{"query": "test"}',
                    },
                },
                {
                    "id": "call_456",
                    "type": "function",
                    "function": {
                        "name": "send_email",
                        "arguments": '{"to": "user@example.com"}',
                    },
                },
            ],
        }
        choice = {"message": message}
        choices = [choice]

        self.context.completions(choices)

        call_args = self.mock_span.set_attribute.call_args
        messages = json.loads(call_args[0][1])
        assert messages[0]["parts"][0]["content"] == "I'll use some tools"
        assert messages[0]["parts"][1]["id"] == "call_123"
        assert messages[0]["parts"][1]["name"] == "search_database"
        assert messages[0]["parts"][2]["id"] == "call_456"
        assert messages[0]["parts"][2]["name"] == "send_email"

    def test_completions_with_delta_object(self):
        delta = {"role": "assistant", "content": "Streaming content"}
        choice = {"delta": delta}
        choices = [choice]

        self.context.completions(choices)

        call_args = self.mock_span.set_attribute.call_args
        messages = json.loads(call_args[0][1])
        assert messages[0]["parts"][0]["content"] == "Streaming content"

    def test_completions_mixed_choices(self):
        choices = [
            {"text": "Simple text response"},
            {"message": {"role": "assistant", "content": "Message response"}},
            {"delta": {"role": "assistant", "content": "Delta response"}},
        ]

        self.context.completions(choices)
        call_args = self.mock_span.set_attribute.call_args
        messages = json.loads(call_args[0][1])
        assert len(messages) == 3

    def test_completions_none_choice(self):
        choices = [None]
        self.context.completions(choices)
        self.mock_span.set_attribute.assert_not_called()

    def test_completions_empty_message(self):
        choice = {"text": ""}
        choices = [choice]
        self.context.completions(choices)
        call_args = self.mock_span.set_attribute.call_args
        messages = json.loads(call_args[0][1])
        assert messages[0]["parts"][0]["content"] == ""

    def test_completions_exception_handling(self):
        self.mock_span.set_attribute.side_effect = Exception("Test error")
        choices = [{"text": "Test"}]
        result = self.context.completions(choices)
        assert result is self.context

    def test_completions_with_complex_content(self):
        message = {
            "role": "assistant",
            "content": [
                {"type": "text", "text": "Hello"},
                {"type": "text", "text": "World"},
            ],
        }
        choice = {"message": message}
        choices = [choice]
        self.context.completions(choices)
        self.mock_span.set_attribute.assert_called_once()

    def test_completions_return_value_chaining(self):
        choices = [{"text": "Test"}]
        result = self.context.completions(choices)
        assert result is self.context
        chained_result = result.completions(choices)
        assert chained_result is self.context


class TestGenerationContextMessages:
    def setup_method(self):
        self.mock_span = Mock()
        self.mock_span.get_span_context.return_value.is_valid = True
        self.mock_span.set_attribute = Mock()
        self.mock_span.add_event = Mock()

        self.context = GenerationContext.__new__(GenerationContext)
        self.context._span = self.mock_span
        self.context.set_attributes = Mock(return_value=self.context)

    def test_messages_empty_list(self):
        self.context.messages([])
        self.mock_span.set_attribute.assert_not_called()

    def test_messages_single_simple_message(self):
        message = {
            "role": "user",
            "parts": [{"type": "text", "content": "Hello, how are you?"}],
        }
        messages = [message]

        self.context.messages(messages)

        self.mock_span.set_attribute.assert_called_once()
        call_args = self.mock_span.set_attribute.call_args
        assert call_args[0][0] == "gen_ai.input.messages"

        parsed_messages = json.loads(call_args[0][1])
        assert len(parsed_messages) == 1
        assert parsed_messages[0]["role"] == "user"
        assert parsed_messages[0]["parts"][0]["type"] == "text"
        assert parsed_messages[0]["parts"][0]["content"] == "Hello, how are you?"

    def test_messages_multiple_simple_messages(self):
        messages = [
            {"role": "user", "parts": [{"type": "text", "content": "First message"}]},
            {
                "role": "assistant",
                "parts": [{"type": "text", "content": "Second message"}],
            },
            {"role": "user", "parts": [{"type": "text", "content": "Third message"}]},
        ]

        self.context.messages(messages)

        call_args = self.mock_span.set_attribute.call_args
        parsed_messages = json.loads(call_args[0][1])
        assert len(parsed_messages) == 3
        assert parsed_messages[0]["parts"][0]["content"] == "First message"
        assert parsed_messages[1]["parts"][0]["content"] == "Second message"
        assert parsed_messages[2]["parts"][0]["content"] == "Third message"

    def test_messages_with_function_call(self):
        message = {
            "role": "assistant",
            "parts": [
                {"type": "text", "content": "I'll call a function"},
                {
                    "type": "tool_call",
                    "name": "calculate_sum",
                    "arguments": '{"a": 5, "b": 3}',
                },
            ],
        }

        messages = [message]
        self.context.messages(messages)

        call_args = self.mock_span.set_attribute.call_args
        parsed_messages = json.loads(call_args[0][1])
        assert parsed_messages[0]["parts"][0]["type"] == "text"
        assert parsed_messages[0]["parts"][0]["content"] == "I'll call a function"
        assert parsed_messages[0]["parts"][1]["type"] == "tool_call"
        assert parsed_messages[0]["parts"][1]["name"] == "calculate_sum"

    def test_messages_with_tool_calls(self):
        message = {
            "role": "assistant",
            "parts": [
                {"type": "text", "content": "I'll use tools"},
                {
                    "type": "tool_call",
                    "id": "call_abc",
                    "name": "search_web",
                    "arguments": '{"query": "python"}',
                },
            ],
        }
        messages = [message]

        self.context.messages(messages)

        call_args = self.mock_span.set_attribute.call_args
        parsed_messages = json.loads(call_args[0][1])
        assert parsed_messages[0]["parts"][0]["type"] == "text"
        assert parsed_messages[0]["parts"][1]["type"] == "tool_call"
        assert parsed_messages[0]["parts"][1]["name"] == "search_web"
        assert parsed_messages[0]["parts"][1]["id"] == "call_abc"

    def test_messages_with_complex_content(self):
        message = {
            "role": "user",
            "parts": [
                {"type": "text", "content": "Hello"},
                {"type": "text", "content": "World"},
                {"type": "text", "content": "Extra text"},
            ],
        }
        messages = [message]

        self.context.messages(messages)
        call_args = self.mock_span.set_attribute.call_args
        parsed_messages = json.loads(call_args[0][1])
        # With the new format, each part is separate
        assert len(parsed_messages[0]["parts"]) == 3
        assert parsed_messages[0]["parts"][0]["content"] == "Hello"
        assert parsed_messages[0]["parts"][1]["content"] == "World"
        assert parsed_messages[0]["parts"][2]["content"] == "Extra text"

    def test_messages_with_none_message(self):
        messages = [
            {"role": "user", "parts": [{"type": "text", "content": "Valid message"}]},
            None,
            {
                "role": "assistant",
                "parts": [{"type": "text", "content": "Another valid message"}],
            },
        ]

        self.context.messages(messages)
        call_args = self.mock_span.set_attribute.call_args
        parsed_messages = json.loads(call_args[0][1])
        assert len(parsed_messages) == 2

    def test_messages_exception_handling(self):
        self.mock_span.set_attribute.side_effect = Exception("Test error")

        messages = [{"role": "user", "parts": [{"type": "text", "content": "Test"}]}]
        self.context.messages(messages)

    def test_messages_return_value_chaining(self):
        messages = [{"role": "user", "parts": [{"type": "text", "content": "Test"}]}]
        result = self.context.messages(messages)

        assert result is self.context
        chained_result = result.messages(messages)
        assert chained_result is self.context

    def test_messages_mixed_content_types(self):
        messages = [
            {"role": "user", "parts": [{"type": "text", "content": "Simple string"}]},
            {
                "role": "assistant",
                "parts": [
                    {"type": "text", "content": "Complex"},
                    {"type": "text", "content": "content"},
                ],
            },
            {"role": "system", "parts": [{"type": "text", "content": "123"}]},
        ]

        self.context.messages(messages)
        call_args = self.mock_span.set_attribute.call_args
        parsed_messages = json.loads(call_args[0][1])
        assert len(parsed_messages) == 3

    def test_messages_with_multiple_tool_calls(self):
        message = {
            "role": "assistant",
            "parts": [
                {"type": "text", "content": "Multiple tools"},
                {
                    "type": "tool_call",
                    "id": "call_1",
                    "name": "tool1",
                    "arguments": "{}",
                },
                {
                    "type": "tool_call",
                    "id": "call_2",
                    "name": "tool2",
                    "arguments": "{}",
                },
            ],
        }
        messages = [message]

        self.context.messages(messages)
        call_args = self.mock_span.set_attribute.call_args
        parsed_messages = json.loads(call_args[0][1])

        assert parsed_messages[0]["parts"][0]["type"] == "text"
        assert parsed_messages[0]["parts"][0]["content"] == "Multiple tools"
        assert parsed_messages[0]["parts"][1]["type"] == "tool_call"
        assert parsed_messages[0]["parts"][1]["id"] == "call_1"
        assert parsed_messages[0]["parts"][1]["name"] == "tool1"
        assert parsed_messages[0]["parts"][2]["type"] == "tool_call"
        assert parsed_messages[0]["parts"][2]["id"] == "call_2"
        assert parsed_messages[0]["parts"][2]["name"] == "tool2"


class TestGenerationContextProvider:
    def setup_method(self):
        self.mock_span = Mock()
        self.mock_span.get_span_context.return_value.is_valid = True
        self.mock_span.set_attribute = Mock()

        self.context = GenerationContext.__new__(GenerationContext)
        self.context._span = self.mock_span

    def test_provider_sets_attribute(self):
        result = self.context.provider("anthropic")
        assert result is self.context
        self.mock_span.set_attribute.assert_called_once_with(
            "gen_ai.provider.name", "anthropic"
        )

    def test_provider_chaining(self):
        result = self.context.provider("openai").provider("anthropic")
        assert result is self.context
        assert self.mock_span.set_attribute.call_count == 2


class TestGenerationContextOperation:
    def setup_method(self):
        self.mock_span = Mock()
        self.mock_span.get_span_context.return_value.is_valid = True
        self.mock_span.set_attribute = Mock()

        self.context = GenerationContext.__new__(GenerationContext)
        self.context._span = self.mock_span

    def test_operation_sets_attribute(self):
        result = self.context.operation("chat")
        assert result is self.context
        self.mock_span.set_attribute.assert_called_once_with(
            "gen_ai.operation.name", "chat"
        )

    def test_operation_chaining(self):
        result = self.context.operation("chat").operation("completion")
        assert result is self.context
        assert self.mock_span.set_attribute.call_count == 2


class TestGenerationContextRequestParams:
    def setup_method(self):
        self.mock_span = Mock()
        self.mock_span.get_span_context.return_value.is_valid = True
        self.mock_span.set_attribute = Mock()

        self.context = GenerationContext.__new__(GenerationContext)
        self.context._span = self.mock_span

    def test_request_params_sets_all_attributes(self):
        usage = {
            "max_tokens": 1024,
            "temperature": 0.7,
            "top_p": 0.9,
            "top_k": 40,
        }
        result = self.context.request_params(usage)
        assert result is self.context
        assert self.mock_span.set_attribute.call_count == 4
        self.mock_span.set_attribute.assert_any_call("gen_ai.request.max_tokens", 1024)
        self.mock_span.set_attribute.assert_any_call("gen_ai.request.temperature", 0.7)
        self.mock_span.set_attribute.assert_any_call("gen_ai.request.top_p", 0.9)
        self.mock_span.set_attribute.assert_any_call("gen_ai.request.top_k", 40)

    def test_request_params_skips_none_values(self):
        usage = {
            "max_tokens": 1024,
            "temperature": None,
            "top_p": 0.9,
        }
        result = self.context.request_params(usage)
        assert result is self.context
        assert self.mock_span.set_attribute.call_count == 2
        self.mock_span.set_attribute.assert_any_call("gen_ai.request.max_tokens", 1024)
        self.mock_span.set_attribute.assert_any_call("gen_ai.request.top_p", 0.9)

    def test_request_params_empty_dict(self):
        result = self.context.request_params({})
        assert result is self.context
        self.mock_span.set_attribute.assert_not_called()

    def test_request_params_chaining(self):
        usage1 = {"max_tokens": 512}
        usage2 = {"temperature": 0.8}
        result = self.context.request_params(usage1).request_params(usage2)
        assert result is self.context
        assert self.mock_span.set_attribute.call_count == 2


class TestGenerationContextResponseId:
    def setup_method(self):
        self.mock_span = Mock()
        self.mock_span.get_span_context.return_value.is_valid = True
        self.mock_span.set_attribute = Mock()

        self.context = GenerationContext.__new__(GenerationContext)
        self.context._span = self.mock_span

    def test_response_id_sets_attribute(self):
        result = self.context.response_id("msg_123")
        assert result is self.context
        self.mock_span.set_attribute.assert_called_once_with(
            "gen_ai.response.id", "msg_123"
        )

    def test_response_id_chaining(self):
        result = self.context.response_id("msg_123").response_id("msg_456")
        assert result is self.context
        assert self.mock_span.set_attribute.call_count == 2
