"""Tests for message extraction functions in llm_scanner."""

import json

import pytest
from inspect_ai._util.content import (
    ContentData,
    ContentDocument,
    ContentImage,
    ContentReasoning,
    ContentText,
)
from inspect_ai.model import (
    ChatMessage,
    ChatMessageAssistant,
    ChatMessageSystem,
    ChatMessageTool,
    ChatMessageUser,
)
from inspect_ai.tool import ToolCall, ToolCallError
from inspect_scout._scanner.extract import (
    MessagesPreprocessor,
    _extract_references,
    message_as_str,
    message_numbering,
    messages_as_str,
)
from inspect_scout._scanner.result import Reference
from inspect_scout._transcript.types import Transcript


@pytest.mark.parametrize(
    "message,expected",
    [
        # Basic user message
        (
            ChatMessageUser(content="Hello, world!"),
            "USER:\nHello, world!\n",
        ),
        # Basic assistant message
        (
            ChatMessageAssistant(content="I can help with that."),
            "ASSISTANT:\nI can help with that.\n",
        ),
        # Basic tool message
        (
            ChatMessageTool(content="Operation successful", function="write_file"),
            "TOOL:\nOperation successful\n",
        ),
        # Message with whitespace gets stripped
        (
            ChatMessageUser(content="  \n  Hello  \n  "),
            "USER:\n  \n  Hello  \n  \n",
        ),
        # Message with empty content
        (
            ChatMessageUser(content=""),
            "USER:\n\n",
        ),
    ],
)
def test_simple_messages(message: ChatMessage, expected: str) -> None:
    """Simple messages produce expected formatted output."""
    result = message_as_str(message)
    assert result == expected


@pytest.mark.parametrize(
    "content,function,error_message,expected",
    [
        # Tool error with function name
        (
            "Error occurred",
            "get_weather",
            "Connection timeout",
            "TOOL:\nError occurred\n\nError in tool call 'get_weather':\nConnection timeout\n\n",
        ),
        # Tool error without function name
        (
            "Something went wrong",
            None,
            "Unknown error",
            "TOOL:\nSomething went wrong\n\nError in tool call 'unknown':\nUnknown error\n\n",
        ),
        # Tool error with empty content
        (
            "",
            "read_file",
            "Access denied",
            "TOOL:\n\n\nError in tool call 'read_file':\nAccess denied\n\n",
        ),
    ],
)
def test_tool_messages_with_errors(
    content: str,
    function: str | None,
    error_message: str,
    expected: str,
) -> None:
    """Tool messages with errors include error information."""
    error = ToolCallError(type="unknown", message=error_message)
    message = ChatMessageTool(content=content, function=function, error=error)
    result = message_as_str(message)
    assert result == expected


@pytest.mark.parametrize(
    "message,expected",
    [
        # Single tool call with dict arguments
        (
            ChatMessageAssistant(
                content="Let me check the weather.",
                tool_calls=[
                    ToolCall(
                        id="call_1",
                        function="get_weather",
                        arguments={"location": "San Francisco", "units": "celsius"},
                    )
                ],
            ),
            "ASSISTANT:\nLet me check the weather.\n\n"
            "Tool Call: get_weather\nArguments:\n"
            "location: San Francisco\nunits: celsius\n",
        ),
        # Multiple tool calls
        (
            ChatMessageAssistant(
                content="Let me check both.",
                tool_calls=[
                    ToolCall(
                        id="call_1",
                        function="get_weather",
                        arguments={"location": "San Francisco"},
                    ),
                    ToolCall(
                        id="call_2",
                        function="get_time",
                        arguments={"timezone": "PST"},
                    ),
                ],
            ),
            "ASSISTANT:\nLet me check both.\n\n"
            "Tool Call: get_weather\nArguments:\nlocation: San Francisco\n\n"
            "Tool Call: get_time\nArguments:\ntimezone: PST\n",
        ),
        # Tool call with empty arguments
        (
            ChatMessageAssistant(
                content="Calling with no args.",
                tool_calls=[
                    ToolCall(id="call_1", function="no_args_function", arguments={})
                ],
            ),
            "ASSISTANT:\nCalling with no args.\n\nTool Call: no_args_function\nArguments:\n\n",
        ),
        # Tool call with empty content
        (
            ChatMessageAssistant(
                content="",
                tool_calls=[
                    ToolCall(
                        id="call_1", function="search", arguments={"query": "test"}
                    )
                ],
            ),
            "ASSISTANT:\n\n\nTool Call: search\nArguments:\nquery: test\n",
        ),
        # Tool call with nested arguments
        (
            ChatMessageAssistant(
                content="Running complex call.",
                tool_calls=[
                    ToolCall(
                        id="call_1",
                        function="complex_function",
                        arguments={
                            "config": {"retry": True, "timeout": 30},
                            "items": ["item1", "item2"],
                        },
                    )
                ],
            ),
            "ASSISTANT:\nRunning complex call.\n\n"
            "Tool Call: complex_function\nArguments:\n"
            "config: {'retry': True, 'timeout': 30}\n"
            "items: ['item1', 'item2']\n",
        ),
    ],
)
def test_assistant_messages_with_tool_calls(
    message: ChatMessageAssistant, expected: str
) -> None:
    """Assistant messages with tool calls format calls and arguments."""
    result = message_as_str(message)
    assert result == expected


@pytest.mark.parametrize(
    "message,expected",
    [
        # Multiple text content parts joined with newlines
        (
            ChatMessageUser(
                content=[
                    ContentText(text="First text"),
                    ContentText(text="Second text"),
                ]
            ),
            "USER:\nFirst text\nSecond text\n",
        ),
        # Mixed content - text extracted, non-text ignored
        (
            ChatMessageUser(
                content=[
                    ContentText(text="Text before image"),
                    ContentImage(image="data:image/png;base64,abc"),
                    ContentText(text="Text after image"),
                ]
            ),
            "USER:\nText before image\n<image />\nText after image\n",
        ),
        # Only non-text content produces empty text
        (
            ChatMessageUser(content=[ContentImage(image="data:image/png;base64,abc")]),
            "USER:\n<image />\n",
        ),
        # Reasoning + text (reasoning is not type='text' so excluded)
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="Let me think..."),
                    ContentText(text="The answer is 42"),
                ]
            ),
            "ASSISTANT:\n\n<thinking>Let me think...</thinking>\nThe answer is 42\n",
        ),
        # Text mixed with various content types
        (
            ChatMessageUser(
                content=[
                    ContentText(text="Start"),
                    ContentData(data={"key": "value"}),
                    ContentText(text="Middle"),
                    ContentDocument(document="data:application/pdf;base64,xyz"),
                    ContentText(text="End"),
                ]
            ),
            "USER:\nStart\n<data />\nMiddle\n<document />\nEnd\n",
        ),
        # Content with whitespace in text parts - strips outer, not inner
        (
            ChatMessageUser(
                content=[
                    ContentText(text="  Text with spaces  "),
                    ContentText(text="\n\nText with newlines\n\n"),
                ]
            ),
            "USER:\n  Text with spaces  \n\n\nText with newlines\n\n\n",
        ),
        # Tool message with list content
        (
            ChatMessageTool(
                content=[
                    ContentText(text="Tool result"),
                    ContentImage(image="data:image/png;base64,abc"),
                ],
                function="analyze",
            ),
            "TOOL:\nTool result\n<image />\n",
        ),
        # Assistant with list content and tool calls
        (
            ChatMessageAssistant(
                content=[
                    ContentText(text="Let me search"),
                    ContentImage(image="data:image/png;base64,abc"),
                    ContentText(text="for that"),
                ],
                tool_calls=[
                    ToolCall(
                        id="call_1", function="search", arguments={"query": "test"}
                    )
                ],
            ),
            "ASSISTANT:\nLet me search\n<image />\nfor that\n\n"
            "Tool Call: search\nArguments:\nquery: test\n",
        ),
        # Tool message with list content and error
        (
            ChatMessageTool(
                content=[
                    ContentText(text="Attempted operation"),
                    ContentData(data={"status": "failed"}),
                ],
                function="slow_operation",
                error=ToolCallError(type="timeout", message="Operation timed out"),
            ),
            "TOOL:\nAttempted operation\n<data />\n\nError in tool call 'slow_operation':\nOperation timed out\n\n",
        ),
    ],
)
def test_messages_with_list_content(message: ChatMessage, expected: str) -> None:
    """Messages with list[Content] extract only text content parts."""
    result = message_as_str(message)
    assert result == expected


@pytest.mark.parametrize(
    "message,exclude_tool_usage,exclude_reasoning,expected",
    [
        # Tool calls only: excluded with string content
        (
            ChatMessageAssistant(
                content="Let me check the weather.",
                tool_calls=[
                    ToolCall(
                        id="call_1",
                        function="get_weather",
                        arguments={"location": "San Francisco"},
                    )
                ],
            ),
            True,
            False,
            "ASSISTANT:\nLet me check the weather.\n",
        ),
        # Tool calls only: excluded with multiple calls
        (
            ChatMessageAssistant(
                content="Let me check both.",
                tool_calls=[
                    ToolCall(id="call_1", function="get_weather", arguments={}),
                    ToolCall(id="call_2", function="get_time", arguments={}),
                ],
            ),
            True,
            False,
            "ASSISTANT:\nLet me check both.\n",
        ),
        # Tool calls only: included (default)
        (
            ChatMessageAssistant(
                content="Using tools.",
                tool_calls=[
                    ToolCall(id="call_1", function="search", arguments={"q": "test"})
                ],
            ),
            False,
            False,
            "ASSISTANT:\nUsing tools.\n\nTool Call: search\nArguments:\nq: test\n",
        ),
        # Tool calls only: excluded with empty content
        (
            ChatMessageAssistant(
                content="",
                tool_calls=[
                    ToolCall(id="call_1", function="search", arguments={"q": "test"})
                ],
            ),
            True,
            False,
            "ASSISTANT:\n\n",
        ),
        # Tool calls only: excluded with list content
        (
            ChatMessageAssistant(
                content=[
                    ContentText(text="Searching"),
                    ContentImage(image="data:image/png;base64,abc"),
                ],
                tool_calls=[
                    ToolCall(id="call_1", function="search", arguments={"q": "test"})
                ],
            ),
            True,
            False,
            "ASSISTANT:\nSearching\n<image />\n",
        ),
        # Reasoning only: included (exclude=False)
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="Let me think..."),
                    ContentText(text="The answer is 42"),
                ]
            ),
            False,
            False,
            "ASSISTANT:\n\n<thinking>Let me think...</thinking>\nThe answer is 42\n",
        ),
        # Reasoning only: excluded (default, exclude=True)
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="Let me think..."),
                    ContentText(text="The answer is 42"),
                ]
            ),
            False,
            True,
            "ASSISTANT:\nThe answer is 42\n",
        ),
        # Reasoning only: only reasoning content, included
        (
            ChatMessageAssistant(content=[ContentReasoning(reasoning="Thinking...")]),
            False,
            False,
            "ASSISTANT:\n\n<thinking>Thinking...</thinking>\n",
        ),
        # Reasoning only: only reasoning content, excluded (default)
        (
            ChatMessageAssistant(content=[ContentReasoning(reasoning="Thinking...")]),
            False,
            True,
            "ASSISTANT:\n\n",
        ),
        # Reasoning only: multiple reasoning parts, included
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="First thought"),
                    ContentText(text="Answer"),
                    ContentReasoning(reasoning="Second thought"),
                ]
            ),
            False,
            False,
            "ASSISTANT:\n\n<thinking>First thought</thinking>\nAnswer\n\n<thinking>Second thought</thinking>\n",
        ),
        # Reasoning only: multiple reasoning parts, excluded (default)
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="First thought"),
                    ContentText(text="Answer"),
                    ContentReasoning(reasoning="Second thought"),
                ]
            ),
            False,
            True,
            "ASSISTANT:\nAnswer\n",
        ),
        # Redacted reasoning with summary: included
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(
                        reasoning="", redacted=True, summary="Summary of thinking"
                    ),
                    ContentText(text="The answer"),
                ]
            ),
            False,
            False,
            "ASSISTANT:\n\n<thinking_summary>Summary of thinking</thinking_summary>\nThe answer\n",
        ),
        # Redacted reasoning with summary: excluded
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(
                        reasoning="", redacted=True, summary="Summary of thinking"
                    ),
                    ContentText(text="The answer"),
                ]
            ),
            False,
            True,
            "ASSISTANT:\nThe answer\n",
        ),
        # Fully redacted reasoning (no summary): included
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="", redacted=True),
                    ContentText(text="The answer"),
                ]
            ),
            False,
            False,
            "ASSISTANT:\n\n<thinking_redacted/>\nThe answer\n",
        ),
        # Fully redacted reasoning (no summary): excluded
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="", redacted=True),
                    ContentText(text="The answer"),
                ]
            ),
            False,
            True,
            "ASSISTANT:\nThe answer\n",
        ),
        # Combined: both excluded
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="Let me think"),
                    ContentText(text="Answer"),
                ],
                tool_calls=[
                    ToolCall(id="call_1", function="search", arguments={"q": "test"})
                ],
            ),
            True,
            True,
            "ASSISTANT:\nAnswer\n",
        ),
        # Combined: both included
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="Let me think"),
                    ContentText(text="Answer"),
                ],
                tool_calls=[
                    ToolCall(id="call_1", function="search", arguments={"q": "test"})
                ],
            ),
            False,
            False,
            "ASSISTANT:\n\n<thinking>Let me think</thinking>\nAnswer\n\n"
            "Tool Call: search\nArguments:\nq: test\n",
        ),
        # Combined: only reasoning included (tool calls excluded)
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="Thinking"),
                    ContentText(text="Done"),
                ],
                tool_calls=[ToolCall(id="call_1", function="act", arguments={})],
            ),
            True,
            False,
            "ASSISTANT:\n\n<thinking>Thinking</thinking>\nDone\n",
        ),
        # Combined: only tool calls included (reasoning excluded, default)
        (
            ChatMessageAssistant(
                content=[
                    ContentReasoning(reasoning="Thinking"),
                    ContentText(text="Done"),
                ],
                tool_calls=[ToolCall(id="call_1", function="act", arguments={})],
            ),
            False,
            True,
            "ASSISTANT:\nDone\n\nTool Call: act\nArguments:\n\n",
        ),
        # Combined: both excluded, only reasoning and tool calls in content
        (
            ChatMessageAssistant(
                content=[ContentReasoning(reasoning="Only thinking")],
                tool_calls=[ToolCall(id="call_1", function="act", arguments={})],
            ),
            True,
            True,
            "ASSISTANT:\n\n",
        ),
    ],
)
def test_exclude_parameters(
    message: ChatMessageAssistant,
    exclude_tool_usage: bool,
    exclude_reasoning: bool,
    expected: str,
) -> None:
    """Test exclude_tool_usage and exclude_reasoning parameters control output."""
    result = message_as_str(
        message,
        MessagesPreprocessor(
            exclude_tool_usage=exclude_tool_usage, exclude_reasoning=exclude_reasoning
        ),
    )
    assert result == expected


@pytest.mark.parametrize(
    "message,expected",
    [
        # Assistant with tool calls + prefill metadata
        (
            ChatMessageAssistant(
                content="Answer",
                tool_calls=[
                    ToolCall(id="call_1", function="search", arguments={"q": "test"})
                ],
                metadata={"prefill": True},
            ),
            "<prefill>\nASSISTANT:\nAnswer\n\nTool Call: search\nArguments:\nq: test\n</prefill>",
        ),
        # Assistant without tool calls + prefill metadata
        (
            ChatMessageAssistant(content="Answer", metadata={"prefill": True}),
            "<prefill>\nASSISTANT:\nAnswer\n</prefill>",
        ),
        # Assistant with prefill=False — not wrapped
        (
            ChatMessageAssistant(content="Answer", metadata={"prefill": False}),
            "ASSISTANT:\nAnswer\n",
        ),
        # Assistant with no metadata — not wrapped
        (
            ChatMessageAssistant(content="Answer"),
            "ASSISTANT:\nAnswer\n",
        ),
    ],
)
def test_prefill_wrapping(message: ChatMessageAssistant, expected: str) -> None:
    """Test that assistant messages with prefill metadata are wrapped in <prefill> tags."""
    result = message_as_str(message)
    assert result == expected


@pytest.mark.parametrize(
    "messages,filter,include_ids,expected_result,expected_ids",
    [
        # Basic case without IDs
        (
            [
                ChatMessageUser(content="Hello", id="msg1"),
                ChatMessageAssistant(content="Hi", id="msg2"),
            ],
            None,
            False,
            "USER:\nHello\n\nASSISTANT:\nHi\n",
            None,
        ),
        # Basic case with IDs
        (
            [
                ChatMessageUser(content="Hello", id="msg1"),
                ChatMessageAssistant(content="Hi", id="msg2"),
            ],
            None,
            True,
            "[M1] USER:\nHello\n\n[M2] ASSISTANT:\nHi\n",
            {"M1": "msg1", "M2": "msg2"},
        ),
        # System messages excluded by default
        (
            [
                ChatMessageSystem(content="System", id="sys1"),
                ChatMessageUser(content="Hello", id="msg1"),
            ],
            MessagesPreprocessor(),
            True,
            "[M1] USER:\nHello\n",
            {"M1": "msg1"},
        ),
        # Tool messages excluded when filter set
        (
            [
                ChatMessageUser(content="Hello", id="msg1"),
                ChatMessageTool(content="Tool result", function="test", id="tool1"),
                ChatMessageAssistant(content="Done", id="msg2"),
            ],
            MessagesPreprocessor(exclude_tool_usage=True),
            True,
            "[M1] USER:\nHello\n\n[M2] ASSISTANT:\nDone\n",
            {"M1": "msg1", "M2": "msg2"},
        ),
        # Empty list
        ([], None, True, "", {}),
        # Single message
        (
            [ChatMessageUser(content="Single", id="msg1")],
            None,
            True,
            "[M1] USER:\nSingle\n",
            {"M1": "msg1"},
        ),
    ],
)
@pytest.mark.asyncio
async def test_messages_as_str(
    messages: list[ChatMessage],
    filter: MessagesPreprocessor[list[ChatMessage]] | None,
    include_ids: bool,
    expected_result: str,
    expected_ids: dict[str, str] | None,
) -> None:
    """Test messages_as_str with various configurations."""
    if include_ids:
        result, extract_references = await messages_as_str(
            messages, preprocessor=filter, include_ids=True
        )
        assert result == expected_result
        # Test that the extract_references function works correctly
        if expected_ids:
            for ordinal_id, msg_id in expected_ids.items():
                refs = extract_references(f"See [{ordinal_id}]")
                assert len(refs) == 1
                assert refs[0].id == msg_id
                assert refs[0].cite == f"[{ordinal_id}]"
    else:
        result = await messages_as_str(messages, preprocessor=filter)
        assert result == expected_result


@pytest.mark.asyncio
async def test_messages_as_str_with_preprocessor() -> None:
    """Test messages_as_str with async message preprocessor."""

    async def keep_only_user_messages(
        messages: list[ChatMessage],
    ) -> list[ChatMessage]:
        return [m for m in messages if m.role == "user"]

    messages: list[ChatMessage] = [
        ChatMessageUser(content="User 1", id="msg1"),
        ChatMessageAssistant(content="Assistant 1", id="msg2"),
        ChatMessageUser(content="User 2", id="msg3"),
    ]

    result, extract_references = await messages_as_str(
        messages,
        preprocessor=MessagesPreprocessor(transform=keep_only_user_messages),
        include_ids=True,
    )

    # Test the formatted string
    assert "[M1] USER:\nUser 1\n" in result
    assert "[M2] USER:\nUser 2\n" in result
    assert "Assistant" not in result

    # Test the extract_references function
    refs = extract_references("See [M1] and [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg3"


@pytest.mark.asyncio
async def test_messages_as_str_with_transcript() -> None:
    """Test messages_as_str accepts Transcript objects."""
    transcript = Transcript(
        transcript_id="test-123",
        source_id="eval-456",
        source_type="test",
        source_uri="file:///test.log",
        messages=[
            ChatMessageUser(content="Hello", id="msg1"),
            ChatMessageAssistant(content="Hi there", id="msg2"),
        ],
    )

    # Without IDs
    result = await messages_as_str(transcript)
    assert result == "USER:\nHello\n\nASSISTANT:\nHi there\n"

    # With IDs
    result_with_ids, extract_references = await messages_as_str(
        transcript, include_ids=True
    )
    assert result_with_ids == "[M1] USER:\nHello\n\n[M2] ASSISTANT:\nHi there\n"

    # Test extract_references
    refs = extract_references("[M1] and [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg2"


@pytest.mark.asyncio
async def test_messages_as_str_with_transcript_and_preprocessor() -> None:
    """Test messages_as_str with Transcript and preprocessor."""

    async def keep_only_user_messages(
        transcript: Transcript,
    ) -> list[ChatMessage]:
        return [m for m in transcript.messages if m.role == "user"]

    transcript = Transcript(
        transcript_id="test-789",
        source_id="eval-abc",
        source_type="test",
        source_uri="file:///test2.log",
        messages=[
            ChatMessageUser(content="First user", id="msg1"),
            ChatMessageAssistant(content="First assistant", id="msg2"),
            ChatMessageUser(content="Second user", id="msg3"),
            ChatMessageAssistant(content="Second assistant", id="msg4"),
        ],
    )

    result, extract_references = await messages_as_str(
        transcript,
        preprocessor=MessagesPreprocessor(transform=keep_only_user_messages),
        include_ids=True,
    )

    # Only user messages should be present
    assert "[M1] USER:\nFirst user\n" in result
    assert "[M2] USER:\nSecond user\n" in result
    assert "assistant" not in result

    # Test extract_references
    refs = extract_references("[M1] [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg3"


@pytest.mark.asyncio
async def test_type_checking_with_generic_preprocessor() -> None:
    """Test that generic MessagesPreprocessor enforces type safety.

    Valid combinations:
    - list[ChatMessage] input + list[ChatMessage] transform
    - Transcript input + Transcript transform
    - Either input + no transform

    Invalid combinations (caught by mypy):
    - Transcript input + list[ChatMessage] transform (type error)
    - list[ChatMessage] input + Transcript transform (type error)
    """

    async def list_transform(messages: list[ChatMessage]) -> list[ChatMessage]:
        return [m for m in messages if m.role == "user"]

    async def transcript_transform(transcript: Transcript) -> list[ChatMessage]:
        return transcript.messages

    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi", id="msg2"),
    ]
    transcript = Transcript(
        transcript_id="t1",
        source_type="test",
        source_id="s1",
        source_uri="test://uri",
        messages=messages,
    )

    # Valid: matching types
    result1 = await messages_as_str(
        messages, preprocessor=MessagesPreprocessor(transform=list_transform)
    )
    assert "Hello" in result1 and "Hi" not in result1

    result2 = await messages_as_str(
        transcript, preprocessor=MessagesPreprocessor(transform=transcript_transform)
    )
    assert "Hello" in result2

    # Valid: no transform works with both
    result3 = await messages_as_str(
        messages, preprocessor=MessagesPreprocessor[list[ChatMessage]]()
    )
    result4 = await messages_as_str(
        transcript, preprocessor=MessagesPreprocessor[Transcript]()
    )
    assert "Hello" in result3 and "Hi" in result3
    assert "Hello" in result4 and "Hi" in result4


@pytest.mark.parametrize(
    "text,id_map,expected",
    [
        # Single message reference
        (
            "The issue is in [M1]",
            {"M1": "msg-123"},
            [Reference(type="message", cite="[M1]", id="msg-123")],
        ),
        # Single event reference
        (
            "The error occurred at [E1]",
            {"E1": "evt-456"},
            [Reference(type="event", cite="[E1]", id="evt-456")],
        ),
        # Multiple message references
        (
            "See [M1] and [M2] for details",
            {"M1": "msg-123", "M2": "msg-456"},
            [
                Reference(type="message", cite="[M1]", id="msg-123"),
                Reference(type="message", cite="[M2]", id="msg-456"),
            ],
        ),
        # Multiple event references
        (
            "Events [E1] and [E3] are related",
            {"E1": "evt-111", "E2": "evt-222", "E3": "evt-333"},
            [
                Reference(type="event", cite="[E1]", id="evt-111"),
                Reference(type="event", cite="[E3]", id="evt-333"),
            ],
        ),
        # Mixed message and event references
        (
            "Message [M1] triggered event [E2]",
            {"M1": "msg-123", "E2": "evt-456"},
            [
                Reference(type="message", cite="[M1]", id="msg-123"),
                Reference(type="event", cite="[E2]", id="evt-456"),
            ],
        ),
        # Multiple mixed references in complex text
        (
            "Analysis: [M1] shows error, [E1] logged it, then [M2] confirmed, see [E3]",
            {"M1": "msg-a", "M2": "msg-b", "E1": "evt-x", "E3": "evt-z"},
            [
                Reference(type="message", cite="[M1]", id="msg-a"),
                Reference(type="event", cite="[E1]", id="evt-x"),
                Reference(type="message", cite="[M2]", id="msg-b"),
                Reference(type="event", cite="[E3]", id="evt-z"),
            ],
        ),
        # Duplicate references - should only appear once
        (
            "Both [M1] and [M1] show the same issue",
            {"M1": "msg-123"},
            [Reference(type="message", cite="[M1]", id="msg-123")],
        ),
        # Reference not in id_map - should be ignored
        (
            "Reference [M1] is valid but [M2] is not",
            {"M1": "msg-123"},
            [Reference(type="message", cite="[M1]", id="msg-123")],
        ),
        # No references in text
        (
            "This text has no references at all",
            {"M1": "msg-123"},
            [],
        ),
        # Empty text
        (
            "",
            {"M1": "msg-123"},
            [],
        ),
        # Empty id_map
        (
            "Text with [M1] reference",
            {},
            [],
        ),
        # References with different number of digits
        (
            "From [M1] to [M10] to [M100]",
            {"M1": "msg-1", "M10": "msg-10", "M100": "msg-100"},
            [
                Reference(type="message", cite="[M1]", id="msg-1"),
                Reference(type="message", cite="[M10]", id="msg-10"),
                Reference(type="message", cite="[M100]", id="msg-100"),
            ],
        ),
        # Adjacent references
        (
            "[M1][M2][E1]",
            {"M1": "msg-1", "M2": "msg-2", "E1": "evt-1"},
            [
                Reference(type="message", cite="[M1]", id="msg-1"),
                Reference(type="message", cite="[M2]", id="msg-2"),
                Reference(type="event", cite="[E1]", id="evt-1"),
            ],
        ),
        # References in various contexts
        (
            "At start [M1] and end [E1].\nNew line [M2]",
            {"M1": "msg-a", "M2": "msg-b", "E1": "evt-x"},
            [
                Reference(type="message", cite="[M1]", id="msg-a"),
                Reference(type="event", cite="[E1]", id="evt-x"),
                Reference(type="message", cite="[M2]", id="msg-b"),
            ],
        ),
        # Case sensitivity - M vs m, E vs e (only uppercase should match)
        (
            "[M1] is valid but [m1] and [e1] are not, [E1] is valid",
            {"M1": "msg-1", "E1": "evt-1", "m1": "msg-invalid", "e1": "evt-invalid"},
            [
                Reference(type="message", cite="[M1]", id="msg-1"),
                Reference(type="event", cite="[E1]", id="evt-1"),
            ],
        ),
        # Bare references are not matched (brackets required)
        (
            "See M1 and M2",
            {"M1": "msg-1", "M2": "msg-2"},
            [],
        ),
        # Range inside brackets (only tokens present in text are matched)
        (
            "Early responses [M1-M3] are fine",
            {"M1": "msg-1", "M2": "msg-2", "M3": "msg-3"},
            [
                Reference(type="message", cite="[M1]", id="msg-1"),
                Reference(type="message", cite="[M3]", id="msg-3"),
            ],
        ),
        # Comma-separated inside brackets
        (
            "See [M2, M4] for details",
            {"M2": "msg-2", "M4": "msg-4"},
            [
                Reference(type="message", cite="[M2]", id="msg-2"),
                Reference(type="message", cite="[M4]", id="msg-4"),
            ],
        ),
        # Mixed messages and events inside brackets
        (
            "Related [M1, E2] show the issue",
            {"M1": "msg-1", "E2": "evt-2"},
            [
                Reference(type="message", cite="[M1]", id="msg-1"),
                Reference(type="event", cite="[E2]", id="evt-2"),
            ],
        ),
        # Multiple bracketed groups in prose (ranges match endpoints only)
        (
            "Concerns [M1-M3] escalate, responses [M5-M7] are problematic",
            {"M1": "m1", "M3": "m3", "M5": "m5", "M7": "m7"},
            [
                Reference(type="message", cite="[M1]", id="m1"),
                Reference(type="message", cite="[M3]", id="m3"),
                Reference(type="message", cite="[M5]", id="m5"),
                Reference(type="message", cite="[M7]", id="m7"),
            ],
        ),
        # Only refs present in id_map are returned
        (
            "[M1-M3] but M2 not in map",
            {"M1": "msg-1", "M3": "msg-3"},
            [
                Reference(type="message", cite="[M1]", id="msg-1"),
                Reference(type="message", cite="[M3]", id="msg-3"),
            ],
        ),
    ],
)
def test_extract_references(
    text: str, id_map: dict[str, str], expected: list[Reference]
) -> None:
    """Test extract_references with various message and event references."""
    result = _extract_references(text, id_map)
    assert result == expected


@pytest.mark.asyncio
async def test_messages_as_str_json_without_ids() -> None:
    """Test format="json" without include_ids."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi there", id="msg2"),
    ]

    result = await messages_as_str(messages, format="json")
    parsed = json.loads(result)

    assert len(parsed) == 2
    assert parsed[0] == {"role": "user", "content": "USER:\nHello\n"}
    assert parsed[1] == {"role": "assistant", "content": "ASSISTANT:\nHi there\n"}
    # No id field when include_ids is not set
    assert "id" not in parsed[0]
    assert "id" not in parsed[1]


@pytest.mark.asyncio
async def test_messages_as_str_json_with_ids() -> None:
    """Test format="json" with include_ids=True."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi there", id="msg2"),
    ]

    result, extract_references = await messages_as_str(
        messages, format="json", include_ids=True
    )
    parsed = json.loads(result)

    assert len(parsed) == 2
    assert parsed[0] == {"id": "M1", "role": "user", "content": "USER:\nHello\n"}
    assert parsed[1] == {
        "id": "M2",
        "role": "assistant",
        "content": "ASSISTANT:\nHi there\n",
    }

    # extract_references should still work
    refs = extract_references("[M1] and [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg2"


@pytest.mark.asyncio
async def test_messages_as_str_json_with_filtered_messages() -> None:
    """Test format="json" with system messages excluded."""
    messages: list[ChatMessage] = [
        ChatMessageSystem(content="System prompt", id="sys1"),
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi", id="msg2"),
    ]

    result = await messages_as_str(
        messages, preprocessor=MessagesPreprocessor[list[ChatMessage]](), format="json"
    )
    parsed = json.loads(result)

    # System message excluded by default
    assert len(parsed) == 2
    assert parsed[0]["role"] == "user"
    assert parsed[1]["role"] == "assistant"


@pytest.mark.asyncio
async def test_messages_as_str_json_empty_list() -> None:
    """Test format="json" with empty message list."""
    messages: list[ChatMessage] = []

    result = await messages_as_str(messages, format="json")
    parsed = json.loads(result)

    assert parsed == []


# --- format="list" tests ---


@pytest.mark.asyncio
async def test_messages_as_str_list_without_ids() -> None:
    """Test format="list" returns list of formatted strings."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi there", id="msg2"),
    ]

    result = await messages_as_str(messages, format="list")

    assert isinstance(result, list)
    assert len(result) == 2
    assert result[0] == "USER:\nHello\n"
    assert result[1] == "ASSISTANT:\nHi there\n"


@pytest.mark.asyncio
async def test_messages_as_str_list_with_ids() -> None:
    """Test format="list" with include_ids=True."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi there", id="msg2"),
    ]

    result, extract_references = await messages_as_str(
        messages, format="list", include_ids=True
    )

    assert isinstance(result, list)
    assert len(result) == 2
    assert result[0] == "[M1] USER:\nHello\n"
    assert result[1] == "[M2] ASSISTANT:\nHi there\n"

    refs = extract_references("[M1] and [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg2"


@pytest.mark.asyncio
async def test_messages_as_str_list_with_filtered_messages() -> None:
    """Test format="list" with system messages excluded."""
    messages: list[ChatMessage] = [
        ChatMessageSystem(content="System prompt", id="sys1"),
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi", id="msg2"),
    ]

    result = await messages_as_str(
        messages, preprocessor=MessagesPreprocessor[list[ChatMessage]](), format="list"
    )

    assert isinstance(result, list)
    assert len(result) == 2
    assert result[0] == "USER:\nHello\n"
    assert result[1] == "ASSISTANT:\nHi\n"


# --- message_numbering() tests ---


@pytest.mark.asyncio
async def test_message_numbering_single_call() -> None:
    """Single call to message_numbering's messages_as_str produces M1..Mn."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(content="Hi there", id="msg2"),
        ChatMessageUser(content="How are you?", id="msg3"),
    ]

    messages_as_str_fn, extract_references = message_numbering()
    result = await messages_as_str_fn(messages)

    assert "[M1] USER:\nHello\n" in result
    assert "[M2] ASSISTANT:\nHi there\n" in result
    assert "[M3] USER:\nHow are you?\n" in result

    refs = extract_references("See [M1] and [M3]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[0].cite == "[M1]"
    assert refs[1].id == "msg3"
    assert refs[1].cite == "[M3]"


@pytest.mark.asyncio
async def test_message_numbering_continues_across_calls() -> None:
    """Multiple calls to messages_as_str continue numbering from prior calls."""
    batch1: list[ChatMessage] = [
        ChatMessageUser(content="First", id="msg1"),
        ChatMessageAssistant(content="Second", id="msg2"),
    ]
    batch2: list[ChatMessage] = [
        ChatMessageUser(content="Third", id="msg3"),
        ChatMessageAssistant(content="Fourth", id="msg4"),
        ChatMessageUser(content="Fifth", id="msg5"),
    ]

    messages_as_str_fn, extract_references = message_numbering()

    result1 = await messages_as_str_fn(batch1)
    assert "[M1]" in result1
    assert "[M2]" in result1

    result2 = await messages_as_str_fn(batch2)
    assert "[M3]" in result2
    assert "[M4]" in result2
    assert "[M5]" in result2
    assert "[M1]" not in result2


@pytest.mark.asyncio
async def test_message_numbering_extract_references_across_calls() -> None:
    """extract_references resolves citations from any prior messages_as_str call."""
    batch1: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg-a"),
        ChatMessageAssistant(content="Hi", id="msg-b"),
    ]
    batch2: list[ChatMessage] = [
        ChatMessageUser(content="More", id="msg-c"),
    ]

    messages_as_str_fn, extract_references = message_numbering()

    await messages_as_str_fn(batch1)
    await messages_as_str_fn(batch2)

    # References from batch1
    refs = extract_references("[M1] and [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg-a"
    assert refs[1].id == "msg-b"

    # Reference from batch2
    refs = extract_references("[M3]")
    assert len(refs) == 1
    assert refs[0].id == "msg-c"

    # Mixed references across batches
    refs = extract_references("See [M1] and [M3]")
    assert len(refs) == 2
    assert refs[0].id == "msg-a"
    assert refs[1].id == "msg-c"


@pytest.mark.asyncio
async def test_message_numbering_empty_list_no_advance() -> None:
    """Empty message list does not advance the counter."""
    messages_as_str_fn, extract_references = message_numbering()

    result_empty = await messages_as_str_fn([])
    assert result_empty == ""

    batch: list[ChatMessage] = [
        ChatMessageUser(content="First", id="msg1"),
    ]
    result = await messages_as_str_fn(batch)
    assert "[M1]" in result

    refs = extract_references("[M1]")
    assert len(refs) == 1
    assert refs[0].id == "msg1"


@pytest.mark.asyncio
async def test_message_numbering_default_excludes_system() -> None:
    """Default (no preprocessor) excludes system messages."""
    messages: list[ChatMessage] = [
        ChatMessageSystem(content="System prompt", id="sys1"),
        ChatMessageUser(content="Hello", id="msg1"),
    ]

    messages_as_str_fn, _ = message_numbering()
    result = await messages_as_str_fn(messages)

    assert "System" not in result
    assert "[M1] USER:\nHello\n" in result


@pytest.mark.asyncio
async def test_message_numbering_with_preprocessor_exclude_reasoning() -> None:
    """Preprocessor exclude_reasoning option filters reasoning content."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageAssistant(
            content=[
                ContentReasoning(reasoning="Let me think..."),
                ContentText(text="The answer is 42"),
            ],
            id="msg2",
        ),
    ]

    messages_as_str_fn, _ = message_numbering(
        preprocessor=MessagesPreprocessor(exclude_reasoning=True)
    )
    result = await messages_as_str_fn(messages)

    assert "[M1]" in result
    assert "[M2]" in result
    assert "think" not in result
    assert "42" in result


@pytest.mark.asyncio
async def test_message_numbering_exclude_tool_usage() -> None:
    """Preprocessor exclude_tool_usage excludes tool messages from numbering."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Hello", id="msg1"),
        ChatMessageTool(content="Tool result", function="test", id="tool1"),
        ChatMessageAssistant(content="Done", id="msg2"),
    ]

    messages_as_str_fn, extract_references = message_numbering(
        preprocessor=MessagesPreprocessor(exclude_tool_usage=True)
    )
    result = await messages_as_str_fn(messages)

    assert "[M1] USER:\nHello\n" in result
    assert "[M2] ASSISTANT:\nDone\n" in result
    assert "Tool" not in result

    refs = extract_references("[M1] [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg2"


@pytest.mark.asyncio
async def test_message_numbering_with_transform() -> None:
    """Preprocessor with custom transform filters messages."""

    async def keep_only_user(messages: list[ChatMessage]) -> list[ChatMessage]:
        return [m for m in messages if m.role == "user"]

    messages: list[ChatMessage] = [
        ChatMessageUser(content="User 1", id="msg1"),
        ChatMessageAssistant(content="Assistant 1", id="msg2"),
        ChatMessageUser(content="User 2", id="msg3"),
    ]

    messages_as_str_fn, extract_references = message_numbering(
        preprocessor=MessagesPreprocessor(transform=keep_only_user)
    )
    result = await messages_as_str_fn(messages)

    assert "[M1] USER:\nUser 1\n" in result
    assert "[M2] USER:\nUser 2\n" in result
    assert "Assistant" not in result

    refs = extract_references("[M1] and [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg3"


@pytest.mark.asyncio
async def test_message_numbering_filtered_messages_continue_numbering() -> None:
    """Filtered messages don't consume numbers; numbering continues across calls."""
    batch1: list[ChatMessage] = [
        ChatMessageSystem(content="System", id="sys1"),
        ChatMessageUser(content="Hello", id="msg1"),
    ]
    batch2: list[ChatMessage] = [
        ChatMessageAssistant(content="Reply", id="msg2"),
    ]

    messages_as_str_fn, extract_references = message_numbering()
    await messages_as_str_fn(batch1)
    await messages_as_str_fn(batch2)

    # System excluded, so batch1 has M1 only, batch2 has M2
    refs = extract_references("[M1] [M2]")
    assert len(refs) == 2
    assert refs[0].id == "msg1"
    assert refs[1].id == "msg2"


# --- metadata["role_label"] tests ---


@pytest.mark.parametrize(
    "message,expected",
    [
        # Already-uppercase label
        (
            ChatMessageUser(content="Hello", metadata={"role_label": "PROPOSER"}),
            "PROPOSER:\nHello\n",
        ),
        # Mixed-case label is uppercased to match the default convention
        (
            ChatMessageAssistant(content="Reply", metadata={"role_label": "Critic"}),
            "CRITIC:\nReply\n",
        ),
        # Lowercase label is uppercased
        (
            ChatMessageTool(
                content="ok",
                function="write_file",
                metadata={"role_label": "worker"},
            ),
            "WORKER:\nok\n",
        ),
        # Empty string label falls back to canonical role
        (
            ChatMessageUser(content="Hello", metadata={"role_label": ""}),
            "USER:\nHello\n",
        ),
        # Missing key falls back to canonical role
        (
            ChatMessageUser(content="Hello", metadata={"other": "x"}),
            "USER:\nHello\n",
        ),
        # Non-string label falls back to canonical role
        (
            ChatMessageUser(content="Hello", metadata={"role_label": 123}),
            "USER:\nHello\n",
        ),
    ],
)
def test_message_as_str_role_label(message: ChatMessage, expected: str) -> None:
    """`metadata["role_label"]` overrides the role prefix when set to a non-empty string."""
    assert message_as_str(message) == expected


def test_message_as_str_role_label_assistant_with_tool_calls() -> None:
    """Custom role label applies to the assistant-with-tool_calls branch."""
    message = ChatMessageAssistant(
        content="Calling tool",
        metadata={"role_label": "proposer"},
        tool_calls=[
            ToolCall(id="tc1", function="get_weather", arguments={"city": "Paris"})
        ],
    )

    expected = (
        "PROPOSER:\nCalling tool\n\nTool Call: get_weather\nArguments:\ncity: Paris\n"
    )
    assert message_as_str(message) == expected


def test_message_as_str_role_label_tool_with_error() -> None:
    """Custom role label applies to the ChatMessageTool error branch."""
    error = ToolCallError(type="unknown", message="boom")
    message = ChatMessageTool(
        content="Failed",
        function="get_weather",
        error=error,
        metadata={"role_label": "worker"},
    )

    expected = "WORKER:\nFailed\n\nError in tool call 'get_weather':\nboom\n\n"
    assert message_as_str(message) == expected


@pytest.mark.asyncio
async def test_messages_as_str_text_uses_role_label() -> None:
    """Default text format honors `metadata["role_label"]`."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Q1", metadata={"role_label": "PROPOSER"}),
        ChatMessageUser(content="Q2", metadata={"role_label": "CRITIC"}),
    ]

    result = await messages_as_str(messages)
    assert result == "PROPOSER:\nQ1\n\nCRITIC:\nQ2\n"


@pytest.mark.asyncio
async def test_messages_as_str_json_uses_role_label() -> None:
    """JSON format uppercases `metadata["role_label"]` for the `role` field."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Q1", metadata={"role_label": "proposer"}),
        ChatMessageAssistant(content="A1", metadata={"role_label": "Critic"}),
        ChatMessageUser(content="Q2"),  # falls back to lowercase canonical role
    ]

    result = await messages_as_str(messages, format="json")
    parsed = json.loads(result)

    assert parsed[0] == {"role": "PROPOSER", "content": "PROPOSER:\nQ1\n"}
    assert parsed[1] == {"role": "CRITIC", "content": "CRITIC:\nA1\n"}
    assert parsed[2] == {"role": "user", "content": "USER:\nQ2\n"}


@pytest.mark.asyncio
async def test_messages_as_str_list_uses_role_label() -> None:
    """List format honors `metadata["role_label"]` in each item."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Q1", metadata={"role_label": "PROPOSER"}),
        ChatMessageUser(content="Q2", metadata={"role_label": "CRITIC"}),
    ]

    result = await messages_as_str(messages, format="list")
    assert result == ["PROPOSER:\nQ1\n", "CRITIC:\nQ2\n"]


@pytest.mark.asyncio
async def test_message_numbering_uses_role_label() -> None:
    """`message_numbering` rendering uses `metadata["role_label"]` after the [M#] prefix."""
    messages: list[ChatMessage] = [
        ChatMessageUser(content="Q1", id="msg1", metadata={"role_label": "PROPOSER"}),
        ChatMessageUser(content="Q2", id="msg2", metadata={"role_label": "CRITIC"}),
    ]

    messages_as_str_fn, _ = message_numbering()
    result = await messages_as_str_fn(messages)

    assert result == "[M1] PROPOSER:\nQ1\n\n[M2] CRITIC:\nQ2\n"
