"""Tests for SwarmMessage types — focusing on usage field handling."""

from swarm.chunk import TokenUsage
from swarm.message import SwarmAssistantMessage


class TestSwarmAssistantMessageUsage:
    """Verify usage field is excluded from OpenAI dict but present in model_dump."""

    def test_to_openai_dict_excludes_usage(self) -> None:
        msg = SwarmAssistantMessage(
            id="msg-1",
            content=[{"type": "text", "text": "hello"}],
            usage=TokenUsage(model="gpt-4o", input_tokens=10, output_tokens=5),
        )
        openai_dict = msg.to_openai_dict()
        assert "usage" not in openai_dict

    def test_to_openai_dict_includes_tool_calls(self) -> None:
        from swarm.providers.types import ToolCallDelta, ToolCallDeltaFunction

        tc = ToolCallDelta(
            id="call-1",
            type="function",
            function=ToolCallDeltaFunction(name="my_tool", arguments='{"x": 1}'),
            index=0,
        )
        msg = SwarmAssistantMessage(
            id="msg-2",
            content=[],
            tool_calls=[tc],
            usage=TokenUsage(model="gpt-4o", input_tokens=10, output_tokens=5),
        )
        openai_dict = msg.to_openai_dict()
        assert "usage" not in openai_dict
        assert len(openai_dict["tool_calls"]) == 1
        assert openai_dict["tool_calls"][0]["function"]["name"] == "my_tool"

    def test_model_dump_with_usage_serializes_full_object(self) -> None:
        usage = TokenUsage(model="gpt-4o", input_tokens=10, output_tokens=5)
        msg = SwarmAssistantMessage(
            id="msg-3",
            content=[{"type": "text", "text": "hi"}],
            usage=usage,
        )
        dumped = msg.model_dump(mode="json")
        assert dumped["usage"]["model"] == "gpt-4o"
        assert dumped["usage"]["input_tokens"] == 10
        assert dumped["usage"]["output_tokens"] == 5

    def test_model_dump_without_usage(self) -> None:
        msg = SwarmAssistantMessage(
            id="msg-4",
            content=[{"type": "text", "text": "hi"}],
        )
        dumped = msg.model_dump(mode="json")
        # usage is None when not set — acceptable in model_dump
        assert dumped["usage"] is None

    def test_to_openai_dict_without_usage_has_no_key(self) -> None:
        msg = SwarmAssistantMessage(
            id="msg-5",
            content=[{"type": "text", "text": "hi"}],
        )
        openai_dict = msg.to_openai_dict()
        assert "usage" not in openai_dict
