# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from unittest.mock import call, patch

from openai.types.chat import chat_completion_chunk

from haystack.components.generators.utils import _convert_streaming_chunks_to_chat_message, print_streaming_chunk
from haystack.dataclasses import (
    ComponentInfo,
    ReasoningContent,
    StreamingChunk,
    ToolCall,
    ToolCallDelta,
    ToolCallResult,
)


def test_convert_streaming_chunks_to_chat_message_tool_calls_in_any_chunk():
    chunks = [
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": None,
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.910076",
            },
            component_info=ComponentInfo(name="test", type="test"),
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0,
                        id="call_ZOj5l67zhZOx6jqjg7ATQwb6",
                        function=chat_completion_chunk.ChoiceDeltaToolCallFunction(
                            arguments="", name="rag_pipeline_tool"
                        ),
                        type="function",
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.913919",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            start=True,
            tool_calls=[
                ToolCallDelta(id="call_ZOj5l67zhZOx6jqjg7ATQwb6", tool_name="rag_pipeline_tool", arguments="", index=0)
            ],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='{"qu')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.914439",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments='{"qu', index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='ery":')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.924146",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments='ery":', index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments=' "Wher')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.924420",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments=' "Wher', index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="e do")
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.944398",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments="e do", index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="es Ma")
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.944958",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments="es Ma", index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="rk liv")
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.945507",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments="rk liv", index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='e?"}')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.946018",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments='e?"}', index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=1,
                        id="call_STxsYY69wVOvxWqopAt3uWTB",
                        function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments="", name="get_weather"),
                        type="function",
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.946578",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=1,
            start=True,
            tool_calls=[
                ToolCallDelta(id="call_STxsYY69wVOvxWqopAt3uWTB", tool_name="get_weather", arguments="", index=1)
            ],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='{"ci')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.946981",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=1,
            tool_calls=[ToolCallDelta(arguments='{"ci', index=1)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='ty": ')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.947411",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=1,
            tool_calls=[ToolCallDelta(arguments='ty": ', index=1)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='"Berli')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.947643",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=1,
            tool_calls=[ToolCallDelta(arguments='"Berli', index=1)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=1, function=chat_completion_chunk.ChoiceDeltaToolCallFunction(arguments='n"}')
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.947939",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=1,
            tool_calls=[ToolCallDelta(arguments='n"}', index=1)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": None,
                "finish_reason": "tool_calls",
                "received_at": "2025-02-19T16:02:55.948772",
            },
            component_info=ComponentInfo(name="test", type="test"),
            finish_reason="tool_calls",
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": None,
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.948772",
                "usage": {
                    "completion_tokens": 42,
                    "prompt_tokens": 282,
                    "total_tokens": 324,
                    "completion_tokens_details": {
                        "accepted_prediction_tokens": 0,
                        "audio_tokens": 0,
                        "reasoning_tokens": 0,
                        "rejected_prediction_tokens": 0,
                    },
                    "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0},
                },
            },
            component_info=ComponentInfo(name="test", type="test"),
        ),
    ]

    # Convert chunks to a chat message
    result = _convert_streaming_chunks_to_chat_message(chunks=chunks)

    assert not result.texts
    assert not result.text

    # Verify both tool calls were found and processed
    assert len(result.tool_calls) == 2
    assert result.tool_calls[0].id == "call_ZOj5l67zhZOx6jqjg7ATQwb6"
    assert result.tool_calls[0].tool_name == "rag_pipeline_tool"
    assert result.tool_calls[0].arguments == {"query": "Where does Mark live?"}
    assert result.tool_calls[1].id == "call_STxsYY69wVOvxWqopAt3uWTB"
    assert result.tool_calls[1].tool_name == "get_weather"
    assert result.tool_calls[1].arguments == {"city": "Berlin"}

    # Verify meta information
    assert result.meta["model"] == "gpt-4o-mini-2024-07-18"
    assert result.meta["finish_reason"] == "tool_calls"
    assert result.meta["index"] == 0
    assert result.meta["completion_start_time"] == "2025-02-19T16:02:55.910076"
    assert result.meta["usage"] == {
        "completion_tokens": 42,
        "prompt_tokens": 282,
        "total_tokens": 324,
        "completion_tokens_details": {
            "accepted_prediction_tokens": 0,
            "audio_tokens": 0,
            "reasoning_tokens": 0,
            "rejected_prediction_tokens": 0,
        },
        "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0},
    }


def test_convert_streaming_chunk_to_chat_message_two_tool_calls_in_same_chunk():
    chunks = [
        StreamingChunk(
            content="",
            meta={
                "model": "mistral-small-latest",
                "index": 0,
                "tool_calls": None,
                "finish_reason": None,
                "usage": None,
            },
            component_info=ComponentInfo(
                type="haystack_integrations.components.generators.mistral.chat.chat_generator.MistralChatGenerator",
                name=None,
            ),
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "mistral-small-latest",
                "index": 0,
                "finish_reason": "tool_calls",
                "usage": {
                    "completion_tokens": 35,
                    "prompt_tokens": 77,
                    "total_tokens": 112,
                    "completion_tokens_details": None,
                    "prompt_tokens_details": None,
                },
            },
            component_info=ComponentInfo(
                type="haystack_integrations.components.generators.mistral.chat.chat_generator.MistralChatGenerator",
                name=None,
            ),
            index=0,
            tool_calls=[
                ToolCallDelta(index=0, tool_name="weather", arguments='{"city": "Paris"}', id="FL1FFlqUG"),
                ToolCallDelta(index=1, tool_name="weather", arguments='{"city": "Berlin"}', id="xSuhp66iB"),
            ],
            start=True,
            finish_reason="tool_calls",
        ),
    ]

    # Convert chunks to a chat message
    result = _convert_streaming_chunks_to_chat_message(chunks=chunks)

    assert not result.texts
    assert not result.text

    # Verify both tool calls were found and processed
    assert len(result.tool_calls) == 2
    assert result.tool_calls[0].id == "FL1FFlqUG"
    assert result.tool_calls[0].tool_name == "weather"
    assert result.tool_calls[0].arguments == {"city": "Paris"}
    assert result.tool_calls[1].id == "xSuhp66iB"
    assert result.tool_calls[1].tool_name == "weather"
    assert result.tool_calls[1].arguments == {"city": "Berlin"}


def test_convert_streaming_chunk_to_chat_message_empty_tool_call_delta():
    chunks = [
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": None,
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.910076",
            },
            component_info=ComponentInfo(name="test", type="test"),
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0,
                        id="call_ZOj5l67zhZOx6jqjg7ATQwb6",
                        function=chat_completion_chunk.ChoiceDeltaToolCallFunction(
                            arguments='{"query":', name="rag_pipeline_tool"
                        ),
                        type="function",
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.913919",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            start=True,
            tool_calls=[
                ToolCallDelta(
                    id="call_ZOj5l67zhZOx6jqjg7ATQwb6", tool_name="rag_pipeline_tool", arguments='{"query":', index=0
                )
            ],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0,
                        function=chat_completion_chunk.ChoiceDeltaToolCallFunction(
                            arguments=' "Where does Mark live?"}'
                        ),
                    )
                ],
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.924420",
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            tool_calls=[ToolCallDelta(arguments=' "Where does Mark live?"}', index=0)],
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": [
                    chat_completion_chunk.ChoiceDeltaToolCall(
                        index=0, function=chat_completion_chunk.ChoiceDeltaToolCallFunction()
                    )
                ],
                "finish_reason": "tool_calls",
                "received_at": "2025-02-19T16:02:55.948772",
            },
            tool_calls=[ToolCallDelta(index=0)],
            component_info=ComponentInfo(name="test", type="test"),
            finish_reason="tool_calls",
            index=0,
        ),
        StreamingChunk(
            content="",
            meta={
                "model": "gpt-4o-mini-2024-07-18",
                "index": 0,
                "tool_calls": None,
                "finish_reason": None,
                "received_at": "2025-02-19T16:02:55.948772",
                "usage": {
                    "completion_tokens": 42,
                    "prompt_tokens": 282,
                    "total_tokens": 324,
                    "completion_tokens_details": {
                        "accepted_prediction_tokens": 0,
                        "audio_tokens": 0,
                        "reasoning_tokens": 0,
                        "rejected_prediction_tokens": 0,
                    },
                    "prompt_tokens_details": {"audio_tokens": 0, "cached_tokens": 0},
                },
            },
            component_info=ComponentInfo(name="test", type="test"),
        ),
    ]

    # Convert chunks to a chat message
    result = _convert_streaming_chunks_to_chat_message(chunks=chunks)

    assert not result.texts
    assert not result.text

    # Verify both tool calls were found and processed
    assert len(result.tool_calls) == 1
    assert result.tool_calls[0].id == "call_ZOj5l67zhZOx6jqjg7ATQwb6"
    assert result.tool_calls[0].tool_name == "rag_pipeline_tool"
    assert result.tool_calls[0].arguments == {"query": "Where does Mark live?"}
    assert result.meta["finish_reason"] == "tool_calls"


def test_convert_streaming_chunk_to_chat_message_with_empty_tool_call_arguments():
    chunks = [
        # Message start with input tokens
        StreamingChunk(
            content="",
            meta={
                "type": "message_start",
                "message": {
                    "id": "msg_123",
                    "type": "message",
                    "role": "assistant",
                    "content": [],
                    "model": "claude-sonnet-4-20250514",
                    "stop_reason": None,
                    "stop_sequence": None,
                    "usage": {"input_tokens": 25, "output_tokens": 0},
                },
            },
            index=0,
            tool_calls=[],
            tool_call_result=None,
            start=True,
            finish_reason=None,
        ),
        # Initial text content
        StreamingChunk(
            content="",
            meta={"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
            index=1,
            tool_calls=[],
            tool_call_result=None,
            start=True,
            finish_reason=None,
        ),
        StreamingChunk(
            content="Let me check",
            meta={"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Let me check"}},
            index=2,
            tool_calls=[],
            tool_call_result=None,
            start=False,
            finish_reason=None,
        ),
        StreamingChunk(
            content=" the weather",
            meta={"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " the weather"}},
            index=3,
            tool_calls=[],
            tool_call_result=None,
            start=False,
            finish_reason=None,
        ),
        # Tool use content
        StreamingChunk(
            content="",
            meta={
                "type": "content_block_start",
                "index": 1,
                "content_block": {"type": "tool_use", "id": "toolu_123", "name": "weather", "input": {}},
            },
            index=5,
            tool_calls=[ToolCallDelta(index=1, id="toolu_123", tool_name="weather", arguments=None)],
            tool_call_result=None,
            start=True,
            finish_reason=None,
        ),
        StreamingChunk(
            content="",
            meta={"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": ""}},
            index=7,
            tool_calls=[ToolCallDelta(index=1, id=None, tool_name=None, arguments="")],
            tool_call_result=None,
            start=False,
            finish_reason=None,
        ),
        # Final message delta
        StreamingChunk(
            content="",
            meta={
                "type": "message_delta",
                "delta": {"stop_reason": "tool_use", "stop_sequence": None},
                "usage": {"completion_tokens": 40},
            },
            index=8,
            tool_calls=[],
            tool_call_result=None,
            start=False,
            finish_reason="tool_calls",
        ),
    ]

    message = _convert_streaming_chunks_to_chat_message(chunks=chunks)

    assert message.texts == ["Let me check the weather"]
    assert len(message.tool_calls) == 1
    assert message.tool_calls[0].arguments == {}
    assert message.tool_calls[0].id == "toolu_123"
    assert message.tool_calls[0].tool_name == "weather"


def test_print_streaming_chunk_content_only():
    chunk = StreamingChunk(
        content="Hello, world!",
        meta={"model": "test-model"},
        component_info=ComponentInfo(name="test", type="test"),
        start=True,
    )
    with patch("builtins.print") as mock_print:
        print_streaming_chunk(chunk)
        expected_calls = [call("[ASSISTANT]\n", flush=True, end=""), call("Hello, world!", flush=True, end="")]
        mock_print.assert_has_calls(expected_calls)


def test_print_streaming_chunk_tool_call():
    chunk = StreamingChunk(
        content="",
        meta={"model": "test-model"},
        component_info=ComponentInfo(name="test", type="test"),
        start=True,
        index=0,
        tool_calls=[ToolCallDelta(id="call_123", tool_name="test_tool", arguments='{"param": "value"}', index=0)],
    )
    with patch("builtins.print") as mock_print:
        print_streaming_chunk(chunk)
        expected_calls = [
            call("[TOOL CALL]\nTool: test_tool \nArguments: ", flush=True, end=""),
            call('{"param": "value"}', flush=True, end=""),
        ]
        mock_print.assert_has_calls(expected_calls)


def test_print_streaming_chunk_tool_call_result():
    chunk = StreamingChunk(
        content="",
        meta={"model": "test-model"},
        component_info=ComponentInfo(name="test", type="test"),
        index=0,
        tool_call_result=ToolCallResult(
            result="Tool execution completed successfully",
            origin=ToolCall(id="call_123", tool_name="test_tool", arguments={}),
            error=False,
        ),
    )
    with patch("builtins.print") as mock_print:
        print_streaming_chunk(chunk)
        expected_calls = [call("[TOOL RESULT]\nTool execution completed successfully", flush=True, end="")]
        mock_print.assert_has_calls(expected_calls)


def test_print_streaming_chunk_with_finish_reason():
    chunk = StreamingChunk(
        content="Final content.",
        meta={"model": "test-model"},
        component_info=ComponentInfo(name="test", type="test"),
        start=True,
        finish_reason="stop",
    )
    with patch("builtins.print") as mock_print:
        print_streaming_chunk(chunk)
        expected_calls = [
            call("[ASSISTANT]\n", flush=True, end=""),
            call("Final content.", flush=True, end=""),
            call("\n\n", flush=True, end=""),
        ]
        mock_print.assert_has_calls(expected_calls)


def test_print_streaming_chunk_empty_chunk():
    chunk = StreamingChunk(
        content="", meta={"model": "test-model"}, component_info=ComponentInfo(name="test", type="test")
    )
    with patch("builtins.print") as mock_print:
        print_streaming_chunk(chunk)
        mock_print.assert_not_called()


def test_convert_streaming_chunks_to_chat_message_usage_not_in_last_chunk():
    """
    Test that usage info is correctly extracted even when it's not in the last chunk.
    This can happen with some API providers like Qwen3 where usage info may be returned
    in a different chunk than the final one.
    """
    chunks = [
        StreamingChunk(
            content="",
            meta={"model": "qwen-plus", "index": 0, "finish_reason": None, "received_at": "2025-01-01T00:00:00.000000"},
            component_info=ComponentInfo(name="test", type="test"),
        ),
        StreamingChunk(
            content="Hello",
            meta={"model": "qwen-plus", "index": 0, "finish_reason": None, "received_at": "2025-01-01T00:00:00.100000"},
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
            start=True,
        ),
        StreamingChunk(
            content=" world",
            meta={"model": "qwen-plus", "index": 0, "finish_reason": None, "received_at": "2025-01-01T00:00:00.200000"},
            component_info=ComponentInfo(name="test", type="test"),
            index=0,
        ),
        # Chunk with usage info (not the last chunk)
        StreamingChunk(
            content="",
            meta={
                "model": "qwen-plus",
                "received_at": "2025-01-01T00:00:00.300000",
                "usage": {"completion_tokens": 10, "prompt_tokens": 20, "total_tokens": 30},
            },
            component_info=ComponentInfo(name="test", type="test"),
            index=None,
        ),
        # Final chunk with finish_reason but no usage (simulating Qwen3 behavior)
        StreamingChunk(
            content="",
            meta={
                "model": "qwen-plus",
                "index": 0,
                "finish_reason": "stop",
                "received_at": "2025-01-01T00:00:00.400000",
                "usage": None,  # No usage info in final chunk
            },
            component_info=ComponentInfo(name="test", type="test"),
            finish_reason="stop",
        ),
    ]

    result = _convert_streaming_chunks_to_chat_message(chunks=chunks)

    assert result.text == "Hello world"
    assert result.meta["model"] == "qwen-plus"
    assert result.meta["finish_reason"] == "stop"
    # Usage should be extracted from the chunk that has it, not just the last chunk
    assert result.meta["usage"] == {"completion_tokens": 10, "prompt_tokens": 20, "total_tokens": 30}


def test_convert_streaming_chunks_to_chat_message_with_reasoning():
    """Test that reasoning content is correctly accumulated from streaming chunks."""
    chunks = [
        StreamingChunk(
            content="",
            meta={"model": "test-model", "received_at": "2025-01-01T00:00:00"},
            component_info=ComponentInfo(name="test", type="test"),
            reasoning=ReasoningContent(reasoning_text="Let me think about this..."),
            index=0,
        ),
        StreamingChunk(
            content="",
            meta={"model": "test-model", "received_at": "2025-01-01T00:00:01"},
            component_info=ComponentInfo(name="test", type="test"),
            reasoning=ReasoningContent(reasoning_text=" The capital of France is Paris."),
            index=0,
        ),
        StreamingChunk(
            content="Paris",
            meta={"model": "test-model", "received_at": "2025-01-01T00:00:02"},
            component_info=ComponentInfo(name="test", type="test"),
        ),
        StreamingChunk(
            content="",
            meta={"model": "test-model", "received_at": "2025-01-01T00:00:03"},
            component_info=ComponentInfo(name="test", type="test"),
            finish_reason="stop",
        ),
    ]

    result = _convert_streaming_chunks_to_chat_message(chunks=chunks)

    assert result.text == "Paris"
    assert result.reasoning is not None
    assert isinstance(result.reasoning, ReasoningContent)
    assert result.reasoning.reasoning_text == "Let me think about this... The capital of France is Paris."
    assert result.meta["finish_reason"] == "stop"


def test_convert_streaming_chunks_to_chat_message_without_reasoning():
    """Test that messages without reasoning work correctly (backward compatibility)."""
    chunks = [
        StreamingChunk(
            content="Hello",
            meta={"model": "test-model", "received_at": "2025-01-01T00:00:00"},
            component_info=ComponentInfo(name="test", type="test"),
        ),
        StreamingChunk(
            content=" world",
            meta={"model": "test-model", "received_at": "2025-01-01T00:00:01"},
            component_info=ComponentInfo(name="test", type="test"),
            finish_reason="stop",
        ),
    ]

    result = _convert_streaming_chunks_to_chat_message(chunks=chunks)

    assert result.text == "Hello world"
    assert result.reasoning is None


def test_print_streaming_chunk_with_reasoning():
    """Test that print_streaming_chunk handles reasoning content correctly."""
    chunk = StreamingChunk(
        content="",
        meta={"model": "test-model"},
        component_info=ComponentInfo(name="test", type="test"),
        start=True,
        reasoning=ReasoningContent(reasoning_text="I am thinking about this question."),
        index=0,
    )
    with patch("builtins.print") as mock_print:
        print_streaming_chunk(chunk)
        expected_calls = [
            call("[REASONING]\n", flush=True, end=""),
            call("I am thinking about this question.", flush=True, end=""),
        ]
        mock_print.assert_has_calls(expected_calls)


def test_print_streaming_chunk_with_reasoning_continuation():
    """Test that print_streaming_chunk handles reasoning continuation correctly."""
    chunk = StreamingChunk(
        content="",
        meta={"model": "test-model"},
        component_info=ComponentInfo(name="test", type="test"),
        start=False,  # Not the first chunk
        reasoning=ReasoningContent(reasoning_text="continued reasoning..."),
        index=0,
    )
    with patch("builtins.print") as mock_print:
        print_streaming_chunk(chunk)
        # Should only print the reasoning text without the header since it's a continuation
        expected_calls = [call("continued reasoning...", flush=True, end="")]
        mock_print.assert_has_calls(expected_calls)
