# Copyright (c) Microsoft. All rights reserved.

import base64
import json
from collections.abc import AsyncIterable, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Literal, cast

import pytest
from pydantic import BaseModel, Field, ValidationError
from pytest import fixture, mark, raises

from agent_framework import (
    AgentResponse,
    AgentResponseUpdate,
    Annotation,
    ChatOptions,
    ChatResponse,
    ChatResponseUpdate,
    Content,
    FunctionTool,
    Message,
    ResponseStream,
    TextSpanRegion,
    ToolMode,
    UsageDetails,
    detect_media_type_from_base64,
    merge_chat_options,
    tool,
)
from agent_framework._compaction import (
    GROUP_ANNOTATION_KEY,
    GROUP_HAS_REASONING_KEY,
    GROUP_ID_KEY,
    GROUP_TOKEN_COUNT_KEY,
)
from agent_framework._types import (
    _get_data_bytes,
    _get_data_bytes_as_str,
    _parse_content_list,
    _parse_structured_response_value,
    _process_update,
    _validate_uri,
    add_usage_details,
    map_chat_to_agent_update,
    validate_tool_mode,
)
from agent_framework.exceptions import AdditionItemMismatch, ContentError


@fixture
def ai_tool() -> FunctionTool:
    """Returns a generic FunctionTool."""

    @tool
    def generic_tool(name: str) -> str:
        """A generic tool that echoes the name."""
        return f"Hello, {name}"

    return generic_tool


@fixture
def tool_tool() -> FunctionTool:
    """Returns a executable FunctionTool."""

    @tool
    def simple_function(x: int, y: int) -> int:
        """A simple function that adds two numbers."""
        return x + y

    return simple_function


# region TextContent


def test_text_content_positional():
    """Test the TextContent class to ensure it initializes correctly and inherits from Content."""
    # Create an instance of TextContent
    content = Content.from_text(
        "Hello, world!", raw_representation="Hello, world!", additional_properties={"version": 1}
    )

    # Check the type and content
    assert content.type == "text"
    assert content.text == "Hello, world!"
    assert content.raw_representation == "Hello, world!"
    assert content.additional_properties["version"] == 1
    # Ensure the instance is of type BaseContent
    assert isinstance(content, Content)
    # Note: No longer using Pydantic validation, so type assignment should work
    content.type = "text"  # This should work fine now


def test_text_content_keyword():
    """Test the TextContent class to ensure it initializes correctly and inherits from Content."""
    # Create an instance of TextContent
    content = Content.from_text(
        text="Hello, world!", raw_representation="Hello, world!", additional_properties={"version": 1}
    )

    # Check the type and content
    assert content.type == "text"
    assert content.text == "Hello, world!"
    assert content.raw_representation == "Hello, world!"
    assert content.additional_properties["version"] == 1
    # Ensure the instance is of type BaseContent
    assert isinstance(content, Content)
    # Note: No longer using Pydantic validation, so type assignment should work
    content.type = "text"  # This should work fine now


# region DataContent


def test_data_content_bytes():
    """Test the DataContent class to ensure it initializes correctly."""
    # Create an instance of DataContent
    content = Content.from_data(
        data=b"test", media_type="application/octet-stream", additional_properties={"version": 1}
    )

    # Check the type and content
    assert content.type == "data"
    assert content.uri == "data:application/octet-stream;base64,dGVzdA=="
    assert content.media_type.startswith("application/") is True  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
    assert content.media_type.startswith("image/") is False  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
    assert content.additional_properties["version"] == 1

    # Ensure the instance is of type BaseContent
    assert isinstance(content, Content)


def test_data_content_uri():
    """Test the Content.from_uri class to ensure it initializes correctly with a URI."""
    # Create an instance of Content.from_uri with a URI and explicit media_type
    content = Content.from_uri(
        uri="data:application/octet-stream;base64,dGVzdA==",
        media_type="application/octet-stream",
        additional_properties={"version": 1},
    )

    # Check the type and content
    assert content.type == "data"
    assert content.uri == "data:application/octet-stream;base64,dGVzdA=="
    # media_type must be explicitly provided
    assert content.media_type == "application/octet-stream"
    assert content.media_type.startswith("application/") is True
    assert content.additional_properties["version"] == 1

    # Ensure the instance is of type BaseContent
    assert isinstance(content, Content)


def test_data_content_invalid():
    """Test the DataContent class to ensure it raises an error for invalid initialization."""
    with pytest.raises(ContentError):
        Content.from_uri(uri="invalid_uri", media_type="text/plain")


def test_data_content_empty():
    """Test the DataContent class to ensure it raises an error for empty data."""
    data = Content.from_data(data=b"", media_type="application/octet-stream")
    assert data.uri == "data:application/octet-stream;base64,"
    assert data.media_type == "application/octet-stream"


def test_data_content_detect_image_format_from_base64():
    """Test the detect_image_format_from_base64 static method."""
    # Test each supported format
    png_data = b"\x89PNG\r\n\x1a\n" + b"fake_data"
    assert detect_media_type_from_base64(data_bytes=png_data) == "image/png"
    assert detect_media_type_from_base64(data_str=base64.b64encode(png_data).decode()) == "image/png"

    jpeg_data = b"\xff\xd8\xff\xe0" + b"fake_data"
    assert detect_media_type_from_base64(data_bytes=jpeg_data) == "image/jpeg"
    assert detect_media_type_from_base64(data_str=base64.b64encode(jpeg_data).decode()) == "image/jpeg"

    webp_data = b"RIFF" + b"1234" + b"WEBP" + b"fake_data"
    assert detect_media_type_from_base64(data_str=base64.b64encode(webp_data).decode()) == "image/webp"
    gif_data = b"GIF89a" + b"fake_data"
    assert detect_media_type_from_base64(data_str=base64.b64encode(gif_data).decode()) == "image/gif"

    # Test fallback behavior
    unknown_data = b"UNKNOWN_FORMAT"
    assert detect_media_type_from_base64(data_str=base64.b64encode(unknown_data).decode()) is None
    assert (
        detect_media_type_from_base64(
            data_uri=f"data:application/octet-stream;base64,{base64.b64encode(unknown_data).decode()}"
        )
        is None
    )
    assert detect_media_type_from_base64(data_bytes=unknown_data) is None
    # Test error handling
    with pytest.raises(ValueError, match="Invalid base64 data provided."):
        detect_media_type_from_base64(data_str="invalid_base64!")
        detect_media_type_from_base64(data_str="")

    with pytest.raises(ValueError, match="Provide exactly one of data_bytes, data_str, or data_uri."):
        detect_media_type_from_base64()
        detect_media_type_from_base64(
            data_bytes=b"data", data_str="data", data_uri="data:application/octet-stream;base64,AAA"
        )
        detect_media_type_from_base64(data_bytes=b"data", data_str="data")
        detect_media_type_from_base64(data_bytes=b"data", data_uri="data:application/octet-stream;base64,AAA")
        detect_media_type_from_base64(data_str="data", data_uri="data:application/octet-stream;base64,AAA")


def test_data_content_create_data_uri_from_base64():
    """Test the create_data_uri_from_base64 class method."""
    # Test with PNG data
    png_data = b"\x89PNG\r\n\x1a\n" + b"fake_data"
    content = Content.from_data(png_data, media_type=detect_media_type_from_base64(data_bytes=png_data))  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    assert content.uri == f"data:image/png;base64,{base64.b64encode(png_data).decode()}"
    assert content.media_type == "image/png"

    # Test with different format
    jpeg_data = b"\xff\xd8\xff\xe0" + b"fake_data"
    jpeg_base64 = base64.b64encode(jpeg_data).decode()
    content = Content.from_data(jpeg_data, media_type=detect_media_type_from_base64(data_bytes=jpeg_data))  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    assert content.uri == f"data:image/jpeg;base64,{jpeg_base64}"
    assert content.media_type == "image/jpeg"


# region UriContent


def test_uri_content():
    """Test the UriContent class to ensure it initializes correctly."""
    content = Content.from_uri(uri="http://example.com", media_type="image/jpg", additional_properties={"version": 1})

    # Check the type and content
    assert content.type == "uri"
    assert content.uri == "http://example.com"
    assert content.media_type == "image/jpg"
    assert content.media_type.startswith("image/") is True
    assert content.media_type.startswith("application/") is False
    assert content.additional_properties["version"] == 1
    assert isinstance(content, Content)


# region: HostedFileContent


def test_hosted_file_content():
    """Test the HostedFileContent class to ensure it initializes correctly."""
    content = Content.from_hosted_file(file_id="file-123", additional_properties={"version": 1})

    # Check the type and content
    assert content.type == "hosted_file"
    assert content.file_id == "file-123"
    assert content.additional_properties["version"] == 1
    assert isinstance(content, Content)


def test_hosted_file_content_minimal():
    """Test the HostedFileContent class with minimal parameters."""
    content = Content.from_hosted_file(file_id="file-456")

    # Check the type and content
    assert content.type == "hosted_file"
    assert content.file_id == "file-456"
    assert content.additional_properties == {}
    assert content.raw_representation is None
    assert isinstance(content, Content)


def test_hosted_file_content_optional_fields():
    """HostedFileContent should capture optional media type and name."""
    content = Content.from_hosted_file(file_id="file-789", media_type="image/png", name="plot.png")

    assert content.media_type == "image/png"
    assert content.name == "plot.png"
    assert content.media_type.startswith("image/")
    assert content.media_type.startswith("application/") is False


# region: CodeInterpreter content


def test_code_interpreter_tool_call_content_parses_inputs():
    call = Content.from_code_interpreter_tool_call(
        call_id="call-1",
        inputs=[Content.from_text(text="print('hi')")],
    )

    assert call.type == "code_interpreter_tool_call"
    assert call.call_id == "call-1"
    assert call.inputs and call.inputs[0].type == "text"
    assert call.inputs[0].text == "print('hi')"


def test_code_interpreter_tool_result_content_outputs():
    result = Content.from_code_interpreter_tool_result(
        call_id="call-2",
        outputs=[
            Content.from_text(text="log output"),
            Content.from_uri(uri="https://example.com/file.png", media_type="image/png"),
        ],
    )

    assert result.type == "code_interpreter_tool_result"
    assert result.call_id == "call-2"
    assert result.outputs is not None
    assert result.outputs[0].type == "text"
    assert result.outputs[1].type == "uri"


# region: Image generation content


def test_image_generation_tool_contents():
    call = Content.from_image_generation_tool_call(image_id="img-1")
    outputs = [Content.from_data(data=b"1234", media_type="image/png")]
    result = Content.from_image_generation_tool_result(image_id="img-1", outputs=outputs)

    assert call.type == "image_generation_tool_call"
    assert call.image_id == "img-1"
    assert result.type == "image_generation_tool_result"
    assert result.image_id == "img-1"
    assert result.outputs and result.outputs[0].type == "data"


# region: MCP server tool content


def test_mcp_server_tool_call_and_result():
    call = Content.from_mcp_server_tool_call(call_id="c-1", tool_name="tool", server_name="server", arguments={"x": 1})
    assert call.type == "mcp_server_tool_call"
    assert call.arguments == {"x": 1}
    assert call.informational_only is True

    result = Content.from_mcp_server_tool_result(call_id="c-1", output=[{"type": "text", "text": "done"}])
    assert result.type == "mcp_server_tool_result"
    assert result.output

    # Empty call_id is allowed, validation happens elsewhere
    call2 = Content.from_mcp_server_tool_call(call_id="", tool_name="tool", server_name="server")
    assert call2.call_id == ""


def test_mcp_server_tool_call_is_always_informational_only():
    direct = Content("mcp_server_tool_call", call_id="c-1", tool_name="tool", informational_only=False)
    assert direct.informational_only is True

    serialized = direct.to_dict()
    assert "informational_only" not in serialized

    restored = Content.from_dict({**serialized, "informational_only": False})
    assert restored.informational_only is True


# region: Shell tool content


def test_shell_tool_call_content_creation():
    call = Content.from_shell_tool_call(
        call_id="shell-1",
        commands=["ls -la", "pwd"],
        timeout_ms=60000,
        max_output_length=4096,
        status="completed",
    )

    assert call.type == "shell_tool_call"
    assert call.call_id == "shell-1"
    assert call.commands == ["ls -la", "pwd"]
    assert call.timeout_ms == 60000
    assert call.max_output_length == 4096
    assert call.status == "completed"


def test_shell_tool_call_content_minimal():
    call = Content.from_shell_tool_call(call_id="shell-2")

    assert call.type == "shell_tool_call"
    assert call.call_id == "shell-2"
    assert call.commands is None
    assert call.timeout_ms is None
    assert call.max_output_length is None
    assert call.status is None


def test_shell_tool_result_content_creation():
    result = Content.from_shell_tool_result(
        call_id="shell-1",
        outputs=[
            Content.from_shell_command_output(stdout="hello world\n", stderr=None, exit_code=0, timed_out=False),
            Content.from_shell_command_output(stderr="error msg", exit_code=1, timed_out=False),
        ],
        max_output_length=4096,
    )

    assert result.type == "shell_tool_result"
    assert result.call_id == "shell-1"
    assert result.outputs is not None
    assert len(result.outputs) == 2
    assert result.outputs[0].type == "shell_command_output"
    assert result.outputs[0].stdout == "hello world\n"
    assert result.outputs[0].exit_code == 0
    assert result.outputs[0].timed_out is False
    assert result.outputs[1].type == "shell_command_output"
    assert result.outputs[1].stderr == "error msg"
    assert result.outputs[1].exit_code == 1
    assert result.max_output_length == 4096


def test_shell_tool_result_with_timeout():
    result = Content.from_shell_tool_result(
        call_id="shell-t",
        outputs=[Content.from_shell_command_output(stdout="partial", timed_out=True)],
    )

    assert result.type == "shell_tool_result"
    assert result.outputs is not None
    assert result.outputs[0].timed_out is True
    assert result.outputs[0].exit_code is None


def test_shell_command_output_content_creation():
    output = Content.from_shell_command_output(
        stdout="hello\n",
        stderr="warn\n",
        exit_code=0,
        timed_out=False,
    )

    assert output.type == "shell_command_output"
    assert output.stdout == "hello\n"
    assert output.stderr == "warn\n"
    assert output.exit_code == 0
    assert output.timed_out is False


def test_shell_content_serialization_roundtrip():
    call = Content.from_shell_tool_call(
        call_id="shell-r",
        commands=["echo hello"],
        timeout_ms=30000,
        status="completed",
    )
    call_dict = call.to_dict()
    restored_call = Content.from_dict(call_dict)
    assert restored_call.type == "shell_tool_call"
    assert restored_call.call_id == "shell-r"
    assert restored_call.commands == ["echo hello"]
    assert restored_call.timeout_ms == 30000
    assert restored_call.status == "completed"

    result = Content.from_shell_tool_result(
        call_id="shell-r",
        outputs=[Content.from_shell_command_output(stdout="hello\n", exit_code=0, timed_out=False)],
        max_output_length=4096,
    )
    result_dict = result.to_dict()
    restored_result = Content.from_dict(result_dict)
    assert restored_result.type == "shell_tool_result"
    assert restored_result.call_id == "shell-r"
    assert restored_result.outputs is not None
    assert len(restored_result.outputs) == 1
    assert restored_result.outputs[0].type == "shell_command_output"
    assert restored_result.outputs[0].stdout == "hello\n"
    assert restored_result.outputs[0].exit_code == 0
    assert restored_result.max_output_length == 4096


# region: HostedVectorStoreContent


def test_hosted_vector_store_content():
    """Test the HostedVectorStoreContent class to ensure it initializes correctly."""
    content = Content.from_hosted_vector_store(vector_store_id="vs-789", additional_properties={"version": 1})

    # Check the type and content
    assert content.type == "hosted_vector_store"
    assert content.vector_store_id == "vs-789"
    assert content.additional_properties["version"] == 1

    # Ensure the instance is of type BaseContent
    assert isinstance(content, Content)
    assert content.type == "hosted_vector_store"
    assert isinstance(content, Content)


def test_hosted_vector_store_content_minimal():
    """Test the HostedVectorStoreContent class with minimal parameters."""
    content = Content.from_hosted_vector_store(vector_store_id="vs-101112")

    # Check the type and content
    assert content.type == "hosted_vector_store"
    assert content.vector_store_id == "vs-101112"
    assert content.additional_properties == {}
    assert content.raw_representation is None


# region FunctionCallContent


def test_function_call_content():
    """Test the FunctionCallContent class to ensure it initializes correctly."""
    content = Content.from_function_call(call_id="1", name="example_function", arguments={"param1": "value1"})

    # Check the type and content
    assert content.type == "function_call"
    assert content.name == "example_function"
    assert content.arguments == {"param1": "value1"}
    assert content.informational_only is False

    # Ensure the instance is of type BaseContent
    assert isinstance(content, Content)


def test_function_call_content_informational_only_serialization():
    content = Content.from_function_call(
        call_id="1",
        name="example_function",
        arguments={"param1": "value1"},
        informational_only=True,
    )

    assert content.informational_only is True
    assert content.to_dict()["informational_only"] is True
    assert Content.from_dict(content.to_dict()).informational_only is True
    assert "informational_only" not in Content.from_function_call(call_id="1", name="f").to_dict()
    assert "informational_only" not in Content.from_text("hello").to_dict(exclude_none=False)


def test_function_call_content_parse_arguments():
    c1 = Content.from_function_call(call_id="1", name="f", arguments='{"a": 1, "b": 2}')
    assert c1.parse_arguments() == {"a": 1, "b": 2}
    c2 = Content.from_function_call(call_id="1", name="f", arguments="not json")
    assert c2.parse_arguments() == {"raw": "not json"}
    c3 = Content.from_function_call(call_id="1", name="f", arguments={"x": None})
    assert c3.parse_arguments() == {"x": None}


def test_function_call_content_add_merging_and_errors():
    # str + str concatenation
    a = Content.from_function_call(call_id="1", name="f", arguments="abc")
    b = Content.from_function_call(call_id="1", name="f", arguments="def")
    c = a + b
    assert isinstance(c.arguments, str) and c.arguments == "abcdef"

    # dict + dict merge
    a = Content.from_function_call(call_id="1", name="f", arguments={"x": 1})
    b = Content.from_function_call(call_id="1", name="f", arguments={"y": 2})
    c = a + b
    assert c.arguments == {"x": 1, "y": 2}

    # informational_only is preserved across streamed chunks
    a = Content.from_function_call(call_id="1", name="f", arguments='{"x":', informational_only=True)
    b = Content.from_function_call(call_id="1", name="f", arguments="1}")
    c = a + b
    assert c.informational_only is True

    # incompatible argument types
    a = Content.from_function_call(call_id="1", name="f", arguments="abc")
    b = Content.from_function_call(call_id="1", name="f", arguments={"y": 2})
    with raises(TypeError):
        _ = a + b

    # incompatible call ids
    a = Content.from_function_call(call_id="1", name="f", arguments="abc")
    b = Content.from_function_call(call_id="2", name="f", arguments="def")

    with raises(ContentError):
        _ = a + b


# region FunctionResultContent


def test_function_result_content():
    """Test the FunctionResultContent class to ensure it initializes correctly."""
    content = Content.from_function_result(call_id="1", result={"param1": "value1"})

    # Check the type and content
    assert content.type == "function_result"
    # Dict results are stringified and stored as text items
    assert "param1" in content.result
    assert "value1" in content.result
    assert content.items is not None
    assert len(content.items) == 1
    assert content.items[0].type == "text"

    # Ensure the instance is of type BaseContent
    assert isinstance(content, Content)


# region UsageDetails


def test_usage_details():
    usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15)
    assert usage["input_token_count"] == 5
    assert usage["output_token_count"] == 10
    assert usage["total_token_count"] == 15


def test_usage_details_addition():
    usage1 = UsageDetails(  # type: ignore[typeddict-unknown-key]
        input_token_count=5,
        output_token_count=10,
        total_token_count=15,
        test1=10,
        test2=20,
    )
    usage2 = UsageDetails(  # type: ignore[typeddict-unknown-key]
        input_token_count=3,
        output_token_count=6,
        total_token_count=9,
        test1=10,
        test3=30,
    )

    combined_usage = add_usage_details(usage1, usage2)
    assert combined_usage["input_token_count"] == 8
    assert combined_usage["output_token_count"] == 16
    assert combined_usage["total_token_count"] == 24
    assert combined_usage["test1"] == 20  # type: ignore[typeddict-item]
    assert combined_usage["test2"] == 20  # type: ignore[typeddict-item]
    assert combined_usage["test3"] == 30  # type: ignore[typeddict-item]


def test_usage_details_fail():
    # TypedDict doesn't validate types at runtime, so this test no longer applies
    # Creating UsageDetails with wrong types won't raise ValueError
    usage = cast(
        UsageDetails,
        {"input_token_count": 5, "output_token_count": 10, "total_token_count": 15, "wrong_type": "42.923"},
    )
    assert usage["wrong_type"] == "42.923"  # type: ignore[typeddict-item]


def test_usage_details_additional_counts():
    usage = UsageDetails(input_token_count=5, output_token_count=10, total_token_count=15, **{"test": 1})  # type: ignore[call-arg, typeddict-unknown-key]
    assert usage.get("test") == 1


def test_usage_details_add_with_none_and_type_errors():
    u = UsageDetails(input_token_count=1)
    # add_usage_details with None returns the non-None value
    v = add_usage_details(u, None)
    assert v == u
    # add_usage_details with None on left
    v2 = add_usage_details(None, u)
    assert v2 == u
    # TypedDict doesn't support + operator, use add_usage_details


def test_usage_details_add_skips_non_int():
    u1 = cast(UsageDetails, {"input_token_count": 10, "other": "test"})
    u2 = cast(UsageDetails, {"input_token_count": 10, "another": "test"})
    u3 = add_usage_details(u1, u2)
    assert len(u3.keys()) == 1
    assert "input_token_count" in u3
    assert u3["input_token_count"] == 20


# region UserInputRequest and Response


def test_function_approval_request_and_response_creation():
    """Test creating a FunctionApprovalRequestContent and producing a response."""
    fc = Content.from_function_call(call_id="call-1", name="do_something", arguments={"a": 1})
    req = Content.from_function_approval_request(id="req-1", function_call=fc)

    assert req.type == "function_approval_request"
    assert req.function_call == fc
    assert req.id == "req-1"
    assert isinstance(req, Content)

    resp = req.to_function_approval_response(True)

    assert isinstance(resp, Content)
    assert resp.type == "function_approval_response"
    assert resp.approved is True
    assert resp.function_call == fc
    assert resp.id == "req-1"


def test_function_approval_serialization_roundtrip():
    fc = Content.from_function_call(call_id="c2", name="f", arguments='{"x":1}')
    req = Content.from_function_approval_request(id="id-2", function_call=fc, additional_properties={"meta": 1})

    dumped = req.to_dict()
    loaded = Content.from_dict(dumped)

    # Test that the basic properties match
    assert loaded.id == req.id
    assert loaded.additional_properties == req.additional_properties
    assert loaded.function_call.call_id == req.function_call.call_id  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
    assert loaded.function_call.name == req.function_call.name  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
    assert loaded.function_call.arguments == req.function_call.arguments  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]

    # Skip the BaseModel validation test since we're no longer using Pydantic
    # The Content union will need to be handled differently when we fully migrate


def test_function_approval_request_function_call_none_guard():
    """Test that accessing function_call attributes is safe when function_call is None."""
    # Construct a Content with type "function_approval_request" but no function_call.
    # This verifies the None-guard pattern used in samples to prevent AttributeError.
    content = Content("function_approval_request", id="req-none")
    assert content.function_call is None

    # A proper approval request always has function_call set
    fc = Content.from_function_call(call_id="call-1", name="do_something", arguments={"a": 1})
    req = Content.from_function_approval_request(id="req-1", function_call=fc)
    assert req.function_call is not None
    assert req.function_call.name == "do_something"
    assert req.function_call.arguments == {"a": 1}


def test_function_approval_accepts_mcp_call():
    """Ensure FunctionApprovalRequestContent supports MCP server tool calls."""
    mcp_call = Content.from_mcp_server_tool_call(
        call_id="c-mcp", tool_name="tool", server_name="srv", arguments={"x": 1}
    )
    req = Content.from_function_approval_request(id="req-mcp", function_call=mcp_call)

    assert isinstance(req.function_call, Content)
    assert req.function_call.call_id == "c-mcp"
    assert req.function_call.informational_only is True


# region BaseContent Serialization


@mark.parametrize(
    "args",
    [
        {"type": "text", "text": "Hello, world!"},
        {"type": "uri", "uri": "http://example.com", "media_type": "text/html"},
        {"type": "function_call", "call_id": "1", "name": "example_function", "arguments": {}},
        {"type": "function_result", "call_id": "1", "result": {}},
        {"type": "file", "file_id": "file-123"},
        {"type": "vector_store", "vector_store_id": "vs-789"},
    ],
)
def test_ai_content_serialization(args: dict):
    content = Content(**args)
    serialized = content.to_dict()
    deserialized = Content.from_dict(serialized)
    assert content == deserialized


# region Message


def test_chat_message_text():
    """Test the Message class to ensure it initializes correctly with text content."""
    # Create a Message with a role and text content
    message = Message(role="user", contents=["Hello, how are you?"])

    # Check the type and content
    assert message.role == "user"
    assert len(message.contents) == 1
    assert message.contents[0].type == "text"
    assert message.contents[0].text == "Hello, how are you?"
    assert message.text == "Hello, how are you?"

    # Ensure the instance is of type BaseContent
    assert isinstance(message.contents[0], Content)


def test_chat_message_contents():
    """Test the Message class to ensure it initializes correctly with contents."""
    # Create a Message with a role and multiple contents
    content1 = Content.from_text("Hello, how are you?")
    content2 = Content.from_text("I'm fine, thank you!")
    message = Message(role="user", contents=[content1, content2])

    # Check the type and content
    assert message.role == "user"
    assert len(message.contents) == 2
    assert message.contents[0].type == "text"
    assert message.contents[1].type == "text"
    assert message.contents[0].text == "Hello, how are you?"
    assert message.contents[1].text == "I'm fine, thank you!"
    assert message.text == "Hello, how are you? I'm fine, thank you!"


def test_chat_message_with_chatrole_instance():
    m = Message(role="user", contents=["hi"])
    assert m.role == "user"
    assert m.text == "hi"


# region ChatResponse


def test_chat_response():
    """Test the ChatResponse class to ensure it initializes correctly with a message."""
    # Create a Message
    message = Message(role="assistant", contents=["I'm doing well, thank you!"])

    # Create a ChatResponse with the message
    response = ChatResponse(messages=message)

    # Check the type and content
    assert response.messages[0].role == "assistant"
    assert response.messages[0].text == "I'm doing well, thank you!"
    assert isinstance(response.messages[0], Message)
    # __str__ returns text
    assert str(response) == response.text


def test_chat_response_accepts_model_alias() -> None:
    """Test ChatResponse accepts model and exposes it through model alias."""
    response = ChatResponse(messages=Message(role="assistant", contents=["Hello"]), model="claude-test")

    assert response.model == "claude-test"
    assert response.model == "claude-test"


class OutputModel(BaseModel):
    response: str


def test_chat_response_with_format():
    """Test the ChatResponse class to ensure it initializes correctly with a message."""
    # Create a Message
    message = Message(role="assistant", contents=['{"response": "Hello"}'])

    # Create a ChatResponse with the message
    response = ChatResponse(messages=message, response_format=OutputModel)

    # Check the type and content
    assert response.messages[0].role == "assistant"
    assert response.messages[0].text == '{"response": "Hello"}'
    assert isinstance(response.messages[0], Message)
    assert response.text == '{"response": "Hello"}'
    assert response.value is not None
    assert response.value.response == "Hello"


def test_chat_response_with_format_init():
    """Test the ChatResponse class to ensure it initializes correctly with a message."""
    # Create a Message
    message = Message(role="assistant", contents=['{"response": "Hello"}'])

    # Create a ChatResponse with the message
    response = ChatResponse(messages=message, response_format=OutputModel)

    # Check the type and content
    assert response.messages[0].role == "assistant"
    assert response.messages[0].text == '{"response": "Hello"}'
    assert isinstance(response.messages[0], Message)
    assert response.text == '{"response": "Hello"}'
    assert response.value is not None
    assert response.value.response == "Hello"


def test_chat_response_with_mapping_response_format() -> None:
    """ChatResponse.value should parse JSON when response_format is a mapping."""
    message = Message(role="assistant", contents=['{"response": "Hello"}'])
    response = ChatResponse(
        messages=message,
        response_format={"type": "object", "properties": {"response": {"type": "string"}}},
    )

    assert response.value is not None
    assert isinstance(response.value, dict)
    assert response.value["response"] == "Hello"


def test_chat_response_value_parses_final_message_with_response_format() -> None:
    """ChatResponse.value should ignore intermediate messages when parsing structured output."""
    response = ChatResponse(
        messages=[
            Message(role="assistant", contents=['{"skill_name": "building-permit-compliance"}']),
            Message(role="assistant", contents=['{"response": "Hello"}']),
        ],
        response_format=OutputModel,
    )

    assert response.text == '{"skill_name": "building-permit-compliance"}\n{"response": "Hello"}'
    assert response.value is not None
    assert response.value.response == "Hello"


def test_agent_response_value_parses_final_message_with_response_format() -> None:
    """AgentResponse.value should ignore intermediate messages when parsing structured output."""
    response = AgentResponse(
        messages=[
            Message(role="assistant", contents=['{"skill_name": "building-permit-compliance"}']),
            Message(role="assistant", contents=['{"response": "Hello"}']),
        ],
        response_format=OutputModel,
    )

    assert response.text == '{"skill_name": "building-permit-compliance"}{"response": "Hello"}'
    assert response.value is not None
    assert response.value.response == "Hello"


def test_agent_response_mapping_value_parses_final_message() -> None:
    """AgentResponse.value should parse the final message for JSON schema mappings."""
    response = AgentResponse(
        messages=[
            Message(role="assistant", contents=['{"skill_name": "building-permit-compliance"}']),
            Message(role="assistant", contents=['{"response": "Hello"}']),
        ],
        response_format={"type": "object", "properties": {"response": {"type": "string"}}},
    )

    assert response.value is not None
    assert isinstance(response.value, dict)
    assert response.value["response"] == "Hello"


def test_chat_response_value_ignores_trailing_non_assistant_message() -> None:
    """ChatResponse.value should parse the final assistant message when later tool output exists."""
    response = ChatResponse(
        messages=[
            Message(role="assistant", contents=['{"response": "Hello"}']),
            Message(role="tool", contents=["tool output is not structured JSON"]),
        ],
        response_format=OutputModel,
    )

    assert response.value is not None
    assert response.value.response == "Hello"


def test_agent_response_value_ignores_trailing_non_assistant_message() -> None:
    """AgentResponse.value should parse the final assistant message when later tool output exists."""
    response = AgentResponse(
        messages=[
            Message(role="assistant", contents=['{"response": "Hello"}']),
            Message(role="tool", contents=["tool output is not structured JSON"]),
        ],
        response_format=OutputModel,
    )

    assert response.value is not None
    assert response.value.response == "Hello"


def test_parse_structured_response_value_empty_text_with_pydantic_model() -> None:
    """Empty text should return None instead of raising when response_format is a Pydantic model."""
    result = _parse_structured_response_value("", OutputModel)
    assert result is None


def test_parse_structured_response_value_empty_text_with_mapping() -> None:
    """Empty text should return None instead of raising when response_format is a mapping."""
    result = _parse_structured_response_value("", {"type": "object"})
    assert result is None


def test_chat_response_value_with_empty_text_and_response_format() -> None:
    """ChatResponse.value should return None when text is empty and response_format is set."""
    message = Message(role="assistant", contents=[""])
    response = ChatResponse(messages=message, response_format=OutputModel)
    assert response.value is None


def test_agent_response_value_with_empty_text_and_response_format() -> None:
    """AgentResponse.value should return None when text is empty and response_format is set."""
    message = Message(role="assistant", contents=[""])
    response = AgentResponse(messages=message, response_format=OutputModel)
    assert response.value is None


def test_chat_response_value_raises_on_invalid_schema():
    """Test that value property raises ValidationError with field constraint details."""

    class StrictSchema(BaseModel):
        id: Literal[5]
        name: str = Field(min_length=10)
        score: int = Field(gt=0, le=100)

    message = Message(role="assistant", contents=['{"id": 1, "name": "test", "score": -5}'])
    response = ChatResponse(messages=message, response_format=StrictSchema)

    with raises(ValidationError) as exc_info:
        _ = response.value

    errors = exc_info.value.errors()
    error_fields = {e["loc"][0] for e in errors}
    assert "id" in error_fields, "Expected 'id' Literal constraint error"
    assert "name" in error_fields, "Expected 'name' min_length constraint error"
    assert "score" in error_fields, "Expected 'score' gt constraint error"


def test_agent_response_value_raises_on_invalid_schema():
    """Test that AgentResponse.value property raises ValidationError with field constraint details."""

    class StrictSchema(BaseModel):
        id: Literal[5]
        name: str = Field(min_length=10)
        score: int = Field(gt=0, le=100)

    message = Message(role="assistant", contents=['{"id": 1, "name": "test", "score": -5}'])
    response = AgentResponse(messages=message, response_format=StrictSchema)

    with raises(ValidationError) as exc_info:
        _ = response.value

    errors = exc_info.value.errors()
    error_fields = {e["loc"][0] for e in errors}
    assert "id" in error_fields, "Expected 'id' Literal constraint error"
    assert "name" in error_fields, "Expected 'name' min_length constraint error"
    assert "score" in error_fields, "Expected 'score' gt constraint error"


# region ChatResponseUpdate


def test_chat_response_update():
    """Test the ChatResponseUpdate class to ensure it initializes correctly with a message."""
    # Create a Message
    message = Content.from_text(text="I'm doing well, thank you!")

    # Create a ChatResponseUpdate with the message
    response_update = ChatResponseUpdate(contents=[message])

    # Check the type and content
    assert response_update.contents[0].text == "I'm doing well, thank you!"
    assert response_update.contents[0].type == "text"
    assert response_update.text == "I'm doing well, thank you!"


def test_chat_response_update_accepts_model_alias() -> None:
    """Test ChatResponseUpdate accepts model and exposes it through model alias."""
    response_update = ChatResponseUpdate(contents=[Content.from_text("Hello")], model="claude-test")

    assert response_update.model == "claude-test"
    assert response_update.model == "claude-test"


def test_chat_response_updates_to_chat_response_one():
    """Test converting ChatResponseUpdate to ChatResponse."""
    # Create a Message
    message1 = Content.from_text("I'm doing well, ")
    message2 = Content.from_text("thank you!")

    # Create a ChatResponseUpdate with the message
    response_updates = [
        ChatResponseUpdate(contents=[message1], message_id="1"),
        ChatResponseUpdate(contents=[message2], message_id="1"),
    ]

    # Convert to ChatResponse
    chat_response = ChatResponse.from_updates(response_updates)

    # Check the type and content
    assert len(chat_response.messages) == 1
    assert chat_response.text == "I'm doing well, thank you!"
    assert isinstance(chat_response.messages[0], Message)
    assert len(chat_response.messages[0].contents) == 1
    assert chat_response.messages[0].message_id == "1"


def test_chat_response_updates_to_chat_response_two():
    """Test converting ChatResponseUpdate to ChatResponse."""
    # Create a Message
    message1 = Content.from_text("I'm doing well, ")
    message2 = Content.from_text("thank you!")

    # Create a ChatResponseUpdate with the message
    response_updates = [
        ChatResponseUpdate(contents=[message1], message_id="1"),
        ChatResponseUpdate(contents=[message2], message_id="2"),
    ]

    # Convert to ChatResponse
    chat_response = ChatResponse.from_updates(response_updates)

    # Check the type and content
    assert len(chat_response.messages) == 2
    assert chat_response.text == "I'm doing well, \nthank you!"
    assert isinstance(chat_response.messages[0], Message)
    assert chat_response.messages[0].message_id == "1"
    assert isinstance(chat_response.messages[1], Message)
    assert chat_response.messages[1].message_id == "2"


def test_chat_response_updates_to_chat_response_multiple():
    """Test converting ChatResponseUpdate to ChatResponse."""
    # Create a Message
    message1 = Content.from_text("I'm doing well, ")
    message2 = Content.from_text("thank you!")

    # Create a ChatResponseUpdate with the message
    response_updates = [
        ChatResponseUpdate(contents=[message1], message_id="1"),
        ChatResponseUpdate(contents=[Content.from_text_reasoning(text="Additional context")], message_id="1"),
        ChatResponseUpdate(contents=[message2], message_id="1"),
    ]

    # Convert to ChatResponse
    chat_response = ChatResponse.from_updates(response_updates)

    # Check the type and content
    assert len(chat_response.messages) == 1
    assert chat_response.text == "I'm doing well,  thank you!"
    assert isinstance(chat_response.messages[0], Message)
    assert len(chat_response.messages[0].contents) == 3
    assert chat_response.messages[0].message_id == "1"


def test_chat_response_updates_to_chat_response_multiple_multiple():
    """Test converting ChatResponseUpdate to ChatResponse."""
    # Create a Message
    message1 = Content.from_text("I'm doing well, ", raw_representation="I'm doing well, ")
    message2 = Content.from_text("thank you!")

    # Create a ChatResponseUpdate with the message
    response_updates = [
        ChatResponseUpdate(contents=[message1], message_id="1"),
        ChatResponseUpdate(contents=[message2], message_id="1"),
        ChatResponseUpdate(contents=[Content.from_text_reasoning(text="Additional context")], message_id="1"),
        ChatResponseUpdate(contents=[Content.from_text(text="More context")], message_id="1"),
        ChatResponseUpdate(contents=[Content.from_text("Final part")], message_id="1"),
    ]

    # Convert to ChatResponse
    chat_response = ChatResponse.from_updates(response_updates)

    # Check the type and content
    assert len(chat_response.messages) == 1
    assert isinstance(chat_response.messages[0], Message)
    assert chat_response.messages[0].message_id == "1"
    assert chat_response.messages[0].contents[0].raw_representation is not None

    assert len(chat_response.messages[0].contents) == 3
    assert chat_response.messages[0].contents[0].type == "text"
    assert chat_response.messages[0].contents[0].text == "I'm doing well, thank you!"
    assert chat_response.messages[0].contents[1].type == "text_reasoning"
    assert chat_response.messages[0].contents[1].text == "Additional context"
    assert chat_response.messages[0].contents[2].type == "text"
    assert chat_response.messages[0].contents[2].text == "More contextFinal part"

    assert chat_response.text == "I'm doing well, thank you! More contextFinal part"


async def test_chat_response_from_async_generator():
    async def gen() -> AsyncIterable[ChatResponseUpdate]:
        yield ChatResponseUpdate(contents=[Content.from_text("Hello")], message_id="1")
        yield ChatResponseUpdate(contents=[Content.from_text(" world")], message_id="1")

    resp = await ChatResponse.from_update_generator(gen())
    assert resp.text == "Hello world"


async def test_chat_response_from_async_generator_output_format():
    async def gen() -> AsyncIterable[ChatResponseUpdate]:
        yield ChatResponseUpdate(contents=[Content.from_text('{ "respon')], message_id="1")
        yield ChatResponseUpdate(contents=[Content.from_text('se": "Hello" }')], message_id="1")

    resp = await ChatResponse.from_update_generator(gen(), output_format_type=OutputModel)
    assert resp.text == '{ "response": "Hello" }'
    assert resp.value is not None
    assert resp.value.response == "Hello"


async def test_chat_response_from_async_generator_output_format_in_method():
    async def gen() -> AsyncIterable[ChatResponseUpdate]:
        yield ChatResponseUpdate(contents=[Content.from_text('{ "respon')], message_id="1")
        yield ChatResponseUpdate(contents=[Content.from_text('se": "Hello" }')], message_id="1")

    resp = await ChatResponse.from_update_generator(gen(), output_format_type=OutputModel)
    assert resp.text == '{ "response": "Hello" }'
    assert resp.value is not None
    assert resp.value.response == "Hello"


async def test_chat_response_from_async_generator_mapping_response_format() -> None:
    async def gen() -> AsyncIterable[ChatResponseUpdate]:
        yield ChatResponseUpdate(contents=[Content.from_text('{ "respon')], message_id="1")
        yield ChatResponseUpdate(contents=[Content.from_text('se": "Hello" }')], message_id="1")

    resp = await ChatResponse.from_update_generator(
        gen(),
        output_format_type={"type": "object", "properties": {"response": {"type": "string"}}},
    )

    assert resp.text == '{ "response": "Hello" }'
    assert resp.value is not None
    assert isinstance(resp.value, dict)
    assert resp.value["response"] == "Hello"


def test_chat_response_from_streaming_updates_parses_final_assistant_message() -> None:
    """Combined streaming updates should parse the final assistant message, not trailing tool output."""
    updates = [
        ChatResponseUpdate(
            role="assistant",
            message_id="skill-message",
            contents=[Content.from_text('{"skill_name": "building-permit-compliance"}')],
        ),
        ChatResponseUpdate(
            role="assistant",
            message_id="final-message",
            contents=[Content.from_text('{"respon')],
        ),
        ChatResponseUpdate(
            message_id="final-message",
            contents=[Content.from_text('se": "Hello"}')],
        ),
        ChatResponseUpdate(
            role="tool",
            message_id="tool-message",
            contents=[Content.from_text("tool output is not structured JSON")],
        ),
    ]

    response = ChatResponse.from_updates(updates, output_format_type=OutputModel)

    assert [message.role for message in response.messages] == ["assistant", "assistant", "tool"]
    assert response.value is not None
    assert response.value.response == "Hello"


# region ToolMode


def test_chat_tool_mode():
    """Test the ToolMode class to ensure it initializes correctly."""
    # Create instances of ToolMode
    auto_mode: ToolMode = {"mode": "auto"}
    required_any: ToolMode = {"mode": "required"}
    required_mode: ToolMode = {"mode": "required", "required_function_name": "example_function"}
    none_mode: ToolMode = {"mode": "none"}
    allowed_mode: ToolMode = {"mode": "auto", "allowed_tools": ["get_weather", "search_docs"]}

    # Check the type and content
    assert auto_mode["mode"] == "auto"
    assert "required_function_name" not in auto_mode
    assert "allowed_tools" not in auto_mode
    assert required_any["mode"] == "required"
    assert "required_function_name" not in required_any
    assert required_mode["mode"] == "required"
    assert required_mode["required_function_name"] == "example_function"
    assert none_mode["mode"] == "none"
    assert "required_function_name" not in none_mode
    assert allowed_mode["mode"] == "auto"
    assert allowed_mode["allowed_tools"] == ["get_weather", "search_docs"]

    # equality of dicts
    assert {"mode": "required", "required_function_name": "example_function"} == {
        "mode": "required",
        "required_function_name": "example_function",
    }


def test_chat_tool_mode_from_dict():
    """Test creating ToolMode from a dictionary."""
    mode: ToolMode = {"mode": "required", "required_function_name": "example_function"}

    # Check the type and content
    assert mode["mode"] == "required"
    assert mode["required_function_name"] == "example_function"


# region ChatOptions


def test_chat_options_init() -> None:
    """Test that ChatOptions can be created as a TypedDict."""
    options: ChatOptions = {}
    assert options.get("model") is None

    # With values
    options_with_model: ChatOptions = {"model": "gpt-4o", "temperature": 0.7}
    assert options_with_model.get("model") == "gpt-4o"
    assert options_with_model.get("temperature") == 0.7


def test_chat_options_tool_choice_validation():
    """Test validate_tool_mode utility function."""
    # Valid string values
    assert validate_tool_mode("auto") == {"mode": "auto"}
    assert validate_tool_mode("required") == {"mode": "required"}
    assert validate_tool_mode("none") == {"mode": "none"}

    # Valid ToolMode dict values
    assert validate_tool_mode({"mode": "auto"}) == {"mode": "auto"}
    assert validate_tool_mode({"mode": "required"}) == {"mode": "required"}
    assert validate_tool_mode({"mode": "required", "required_function_name": "example_function"}) == {
        "mode": "required",
        "required_function_name": "example_function",
    }
    assert validate_tool_mode({"mode": "none"}) == {"mode": "none"}

    # None should remain unset
    assert validate_tool_mode(None) is None

    with raises(ContentError):
        validate_tool_mode("invalid_mode")  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
    with raises(ContentError):
        validate_tool_mode({"mode": "invalid_mode"})  # type: ignore[arg-type, typeddict-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
    with raises(ContentError):
        validate_tool_mode({"mode": "auto", "required_function_name": "should_not_be_here"})

    # Valid allowed_tools
    assert validate_tool_mode({"mode": "auto", "allowed_tools": ["get_weather"]}) == {
        "mode": "auto",
        "allowed_tools": ["get_weather"],
    }
    assert validate_tool_mode({"mode": "auto", "allowed_tools": ["get_weather", "search_docs"]}) == {
        "mode": "auto",
        "allowed_tools": ["get_weather", "search_docs"],
    }

    # allowed_tools valid with required mode
    assert validate_tool_mode({"mode": "required", "allowed_tools": ["get_weather"]}) == {
        "mode": "required",
        "allowed_tools": ["get_weather"],
    }

    # allowed_tools invalid with none mode
    with raises(ContentError):
        validate_tool_mode({"mode": "none", "allowed_tools": ["get_weather"]})

    # allowed_tools must be a non-string sequence of strings
    with raises(ContentError):
        validate_tool_mode({"mode": "auto", "allowed_tools": "get_weather"})  # type: ignore[arg-type, typeddict-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
    with raises(ContentError):
        validate_tool_mode({"mode": "auto", "allowed_tools": 123})  # type: ignore[arg-type, typeddict-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
    with raises(ContentError):
        validate_tool_mode({"mode": "auto", "allowed_tools": ["get_weather", 123]})  # type: ignore[arg-type, list-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    # Empty list is valid (caller explicitly allows no tools)
    assert validate_tool_mode({"mode": "auto", "allowed_tools": []}) == {
        "mode": "auto",
        "allowed_tools": [],
    }

    # Tuple is normalized to list
    result = validate_tool_mode({"mode": "auto", "allowed_tools": ("get_weather",)})  # type: ignore[arg-type, typeddict-item]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
    assert result is not None
    assert result["allowed_tools"] == ["get_weather"]


def test_chat_options_merge(tool_tool, ai_tool) -> None:
    """Test merge_chat_options utility function."""
    options1: ChatOptions = {
        "model": "gpt-4o",
        "tools": [tool_tool],
        "logit_bias": {"x": 1},
        "metadata": {"a": "b"},
    }
    options2: ChatOptions = {"model": "gpt-4.1", "tools": [ai_tool]}
    assert options1 != options2

    # Merge options - override takes precedence for non-collection fields
    options3 = merge_chat_options(options1, options2)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    assert options3.get("model") == "gpt-4.1"
    assert options3.get("tools") == [tool_tool, ai_tool]  # tools are combined
    assert options3.get("logit_bias") == {"x": 1}  # base value preserved
    assert options3.get("metadata") == {"a": "b"}  # base value preserved


def test_chat_options_and_tool_choice_override() -> None:
    """Test that tool_choice from other takes precedence in ChatOptions merge."""
    # Agent-level defaults to "auto"
    agent_options: ChatOptions = {"model": "gpt-4o", "tool_choice": "auto"}
    # Run-level specifies "required"
    run_options: ChatOptions = {"tool_choice": "required"}

    merged = merge_chat_options(agent_options, run_options)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    # Run-level should override agent-level
    assert merged.get("tool_choice") == "required"
    assert merged.get("model") == "gpt-4o"  # Other fields preserved


def test_chat_options_and_tool_choice_none_in_other_uses_self() -> None:
    """Test that when other.tool_choice is None, self.tool_choice is used."""
    agent_options: ChatOptions = {"tool_choice": "auto"}
    run_options: ChatOptions = {"model": "gpt-4.1"}  # tool_choice is None

    merged = merge_chat_options(agent_options, run_options)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    # Should keep agent-level tool_choice since run-level is None
    assert merged.get("tool_choice") == "auto"
    assert merged.get("model") == "gpt-4.1"


def test_chat_options_and_tool_choice_with_tool_mode() -> None:
    """Test ChatOptions merge with ToolMode objects."""
    agent_options: ChatOptions = {"tool_choice": "auto"}
    run_options: ChatOptions = {"tool_choice": "required"}

    merged = merge_chat_options(agent_options, run_options)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    assert merged.get("tool_choice") == "required"
    assert merged.get("tool_choice") == "required"


def test_chat_options_and_tool_choice_required_specific_function() -> None:
    """Test ChatOptions merge with required specific function."""
    agent_options: ChatOptions = {"tool_choice": "auto"}
    run_options: ChatOptions = {"tool_choice": {"mode": "required", "required_function_name": "get_weather"}}

    merged = merge_chat_options(agent_options, run_options)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

    tool_choice = merged.get("tool_choice")
    assert isinstance(tool_choice, dict)
    assert tool_choice == {"mode": "required", "required_function_name": "get_weather"}
    assert tool_choice["required_function_name"] == "get_weather"  # pyrefly: ignore[unsupported-operation]


# region Agent Response Fixtures


@fixture
def chat_message() -> Message:
    return Message(role="user", contents=["Hello"])


@fixture
def text_content() -> Content:
    return Content.from_text(text="Test content")


@fixture
def agent_response(chat_message: Message) -> AgentResponse:
    return AgentResponse(messages=chat_message)


@fixture
def agent_response_update(text_content: Content) -> AgentResponseUpdate:
    return AgentResponseUpdate(role="assistant", contents=[text_content])


# region AgentResponse


def test_agent_run_response_init_single_message(chat_message: Message) -> None:
    response = AgentResponse(messages=chat_message)
    assert response.messages == [chat_message]


def test_agent_run_response_init_list_messages(chat_message: Message) -> None:
    response = AgentResponse(messages=[chat_message, chat_message])
    assert len(response.messages) == 2
    assert response.messages[0] == chat_message


def test_agent_run_response_init_none_messages() -> None:
    response = AgentResponse()
    assert response.messages == []


def test_agent_run_response_text_property(chat_message: Message) -> None:
    response = AgentResponse(messages=[chat_message, chat_message])
    assert response.text == "HelloHello"


def test_agent_run_response_text_property_empty() -> None:
    response = AgentResponse()
    assert response.text == ""


def test_agent_run_response_from_updates(agent_response_update: AgentResponseUpdate) -> None:
    updates = [agent_response_update, agent_response_update]
    response = AgentResponse.from_updates(updates)
    assert len(response.messages) > 0
    assert response.text == "Test contentTest content"


def test_agent_run_response_str_method(chat_message: Message) -> None:
    response = AgentResponse(messages=chat_message)
    assert str(response) == "Hello"


# region AgentResponseUpdate


def test_agent_run_response_update_init_content_list(text_content: Content) -> None:
    update = AgentResponseUpdate(contents=[text_content, text_content])
    assert len(update.contents) == 2
    assert update.contents[0] == text_content


def test_agent_run_response_update_init_none_content() -> None:
    update = AgentResponseUpdate()
    assert update.contents == []


def test_agent_run_response_update_text_property(text_content: Content) -> None:
    update = AgentResponseUpdate(contents=[text_content, text_content])
    assert update.text == "Test contentTest content"


def test_agent_run_response_update_text_property_empty() -> None:
    update = AgentResponseUpdate()
    assert update.text == ""


def test_agent_run_response_update_str_method(text_content: Content) -> None:
    update = AgentResponseUpdate(contents=[text_content])
    assert str(update) == "Test content"


def test_agent_run_response_update_created_at() -> None:
    """Test that AgentResponseUpdate properly handles created_at timestamps."""
    # Test with a properly formatted UTC timestamp
    utc_timestamp = "2024-12-01T00:31:30.000000Z"
    update = AgentResponseUpdate(
        contents=[Content.from_text(text="test")],
        role="assistant",
        created_at=utc_timestamp,
    )
    assert update.created_at == utc_timestamp
    assert update.created_at.endswith("Z"), "Timestamp should end with 'Z' for UTC"

    # Verify that we can generate a proper UTC timestamp
    now_utc = datetime.now(tz=timezone.utc)
    formatted_utc = now_utc.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
    update_with_now = AgentResponseUpdate(
        contents=[Content.from_text(text="test")],
        role="assistant",
        created_at=formatted_utc,
    )
    assert update_with_now.created_at == formatted_utc
    assert update_with_now.created_at is not None
    assert update_with_now.created_at.endswith("Z")


def test_agent_run_response_created_at() -> None:
    """Test that AgentResponse properly handles created_at timestamps."""
    # Test with a properly formatted UTC timestamp
    utc_timestamp = "2024-12-01T00:31:30.000000Z"
    response = AgentResponse(
        messages=[Message(role="assistant", contents=["Hello"])],
        created_at=utc_timestamp,
    )
    assert response.created_at == utc_timestamp
    assert response.created_at.endswith("Z"), "Timestamp should end with 'Z' for UTC"

    # Verify that we can generate a proper UTC timestamp
    now_utc = datetime.now(tz=timezone.utc)
    formatted_utc = now_utc.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
    response_with_now = AgentResponse(
        messages=[Message(role="assistant", contents=["Hello"])],
        created_at=formatted_utc,
    )
    assert response_with_now.created_at == formatted_utc
    assert response_with_now.created_at is not None
    assert response_with_now.created_at.endswith("Z")


# region ErrorContent


def test_error_content_str():
    e1 = Content.from_error(message="Oops", error_code="E1")
    assert str(e1) == "Error E1: Oops"
    e2 = Content.from_error(message="Oops")
    assert str(e2) == "Oops"
    e3 = Content.from_error()
    assert str(e3) == "Unknown error"


# region Annotation


def test_annotations_models_and_roundtrip():
    span = TextSpanRegion(type="text_span", start_index=0, end_index=5)
    cit = Annotation(
        type="citation", title="Doc", url="http://example.com", snippet="Snippet", annotated_regions=[span]
    )

    # Attach to content
    content = Content.from_text(text="hello", additional_properties={"v": 1})
    content.annotations = [cit]

    dumped = content.to_dict()
    loaded = Content.from_dict(dumped)
    assert isinstance(loaded.annotations, list)
    assert len(loaded.annotations) == 1
    # After migration from Pydantic, annotations are now TypedDicts (dicts at runtime)
    assert isinstance(loaded.annotations[0], dict)
    # Check the annotation properties
    loaded_cit = loaded.annotations[0]
    assert loaded_cit["type"] == "citation"
    assert loaded_cit["title"] == "Doc"
    assert loaded_cit["url"] == "http://example.com"
    assert loaded_cit["snippet"] == "Snippet"
    # Check the annotated_regions
    assert isinstance(loaded_cit["annotated_regions"], list)
    assert len(loaded_cit["annotated_regions"]) == 1
    assert isinstance(loaded_cit["annotated_regions"][0], dict)
    assert loaded_cit["annotated_regions"][0]["type"] == "text_span"
    assert loaded_cit["annotated_regions"][0]["start_index"] == 0
    assert loaded_cit["annotated_regions"][0]["end_index"] == 5


def test_function_call_merge_in_process_update_and_usage_aggregation():
    # Two function call chunks with same call_id should merge
    u1 = ChatResponseUpdate(
        contents=[Content.from_function_call(call_id="c1", name="f", arguments="{")], message_id="m"
    )
    u2 = ChatResponseUpdate(
        contents=[Content.from_function_call(call_id="c1", name="f", arguments="}")], message_id="m"
    )
    # plus usage
    u3 = ChatResponseUpdate(contents=[Content.from_usage(UsageDetails(input_token_count=1, output_token_count=2))])

    resp = ChatResponse.from_updates([u1, u2, u3])
    assert len(resp.messages) == 1
    last_contents = resp.messages[0].contents
    assert any(c.type == "function_call" for c in last_contents)
    fcs = [c for c in last_contents if c.type == "function_call"]
    assert len(fcs) == 1
    assert fcs[0].arguments == "{}"
    assert resp.usage_details is not None
    assert resp.usage_details["input_token_count"] == 1
    assert resp.usage_details["output_token_count"] == 2


def test_function_call_incompatible_ids_are_not_merged():
    u1 = ChatResponseUpdate(contents=[Content.from_function_call(call_id="a", name="f", arguments="x")], message_id="m")
    u2 = ChatResponseUpdate(contents=[Content.from_function_call(call_id="b", name="f", arguments="y")], message_id="m")

    resp = ChatResponse.from_updates([u1, u2])
    fcs = [c for c in resp.messages[0].contents if c.type == "function_call"]
    assert len(fcs) == 2


# region Role & FinishReason basics


def test_chat_role_str_and_repr():
    # Role is now a NewType of str, so it's just a plain string
    assert "user" == "user"
    assert repr("user") == "'user'"


def test_chat_finish_reason_constants():
    # FinishReason is now a NewType of str, so it's just a plain string
    assert "stop" == "stop"


def test_response_update_propagates_fields_and_metadata():
    upd = ChatResponseUpdate(
        contents=[Content.from_text("hello")],
        role="assistant",
        author_name="bot",
        response_id="rid",
        message_id="mid",
        conversation_id="cid",
        model="model-x",
        created_at="t0",
        finish_reason="stop",
        additional_properties={"k": "v"},
    )
    resp = ChatResponse.from_updates([upd])
    assert resp.response_id == "rid"
    assert resp.created_at == "t0"
    assert resp.conversation_id == "cid"
    assert resp.model == "model-x"
    assert resp.finish_reason == "stop"
    assert resp.additional_properties and resp.additional_properties["k"] == "v"
    assert resp.messages[0].role == "assistant"
    assert resp.messages[0].author_name == "bot"
    assert resp.messages[0].message_id == "mid"


def test_text_coalescing_preserves_first_properties():
    t1 = Content.from_text("A", raw_representation={"r": 1}, additional_properties={"p": 1})
    t2 = Content.from_text("B")
    upd1 = ChatResponseUpdate(contents=[t1], message_id="x")
    upd2 = ChatResponseUpdate(contents=[t2], message_id="x")
    resp = ChatResponse.from_updates([upd1, upd2])
    # After coalescing there should be a single TextContent with merged text and preserved props from first
    items = [c for c in resp.messages[0].contents if c.type == "text"]
    assert len(items) >= 1
    assert items[0].text == "AB"
    assert items[0].raw_representation == {"r": 1}
    assert items[0].additional_properties == {"p": 1}


def test_function_call_content_parse_numeric_or_list():
    c_num = Content.from_function_call(call_id="1", name="f", arguments="123")
    assert c_num.parse_arguments() == {"raw": 123}
    c_list = Content.from_function_call(call_id="1", name="f", arguments="[1,2]")
    assert c_list.parse_arguments() == {"raw": [1, 2]}


def test_chat_tool_mode_eq_with_string():
    assert {"mode": "auto"} == {"mode": "auto"}


# region AgentResponse


@fixture
def agent_run_response_async() -> AgentResponse:
    return AgentResponse(messages=[Message(role="user", contents=["Hello"])])


async def test_agent_run_response_from_async_generator():
    async def gen():
        yield AgentResponseUpdate(contents=[Content.from_text("A")])
        yield AgentResponseUpdate(contents=[Content.from_text("B")])

    r = await AgentResponse.from_update_generator(gen())
    assert r.text == "AB"


# region Additional Coverage Tests for Serialization and Arithmetic Methods


def test_text_content_add_comprehensive_coverage():
    """Test TextContent __add__ method with various combinations to improve coverage."""

    # Test with None raw_representation
    t1 = Content.from_text("Hello", raw_representation=None, annotations=None)
    t2 = Content.from_text(" World", raw_representation=None, annotations=None)
    result = t1 + t2
    assert result.text == "Hello World"
    assert result.raw_representation is None
    assert result.annotations is None

    # Test first has raw_representation, second has None
    t1 = Content.from_text("Hello", raw_representation="raw1", annotations=None)
    t2 = Content.from_text(" World", raw_representation=None, annotations=None)
    result = t1 + t2
    assert result.text == "Hello World"
    assert result.raw_representation == "raw1"

    # Test first has None, second has raw_representation
    t1 = Content.from_text("Hello", raw_representation=None, annotations=None)
    t2 = Content.from_text(" World", raw_representation="raw2", annotations=None)
    result = t1 + t2
    assert result.text == "Hello World"
    assert result.raw_representation == "raw2"

    # Test both have raw_representation (non-list)
    t1 = Content.from_text("Hello", raw_representation="raw1", annotations=None)
    t2 = Content.from_text(" World", raw_representation="raw2", annotations=None)
    result = t1 + t2
    assert result.text == "Hello World"
    assert result.raw_representation == ["raw1", "raw2"]

    # Test first has list raw_representation, second has single
    t1 = Content.from_text("Hello", raw_representation=["raw1", "raw2"], annotations=None)
    t2 = Content.from_text(" World", raw_representation="raw3", annotations=None)
    result = t1 + t2
    assert result.text == "Hello World"
    assert result.raw_representation == ["raw1", "raw2", "raw3"]

    # Test both have list raw_representation
    t1 = Content.from_text("Hello", raw_representation=["raw1", "raw2"], annotations=None)
    t2 = Content.from_text(" World", raw_representation=["raw3", "raw4"], annotations=None)
    result = t1 + t2
    assert result.text == "Hello World"
    assert result.raw_representation == ["raw1", "raw2", "raw3", "raw4"]

    # Test first has single raw_representation, second has list
    t1 = Content.from_text("Hello", raw_representation="raw1", annotations=None)
    t2 = Content.from_text(" World", raw_representation=["raw2", "raw3"], annotations=None)
    result = t1 + t2
    assert result.text == "Hello World"
    assert result.raw_representation == ["raw1", "raw2", "raw3"]


def test_text_content_iadd_coverage():
    """Test TextContent += operator for better coverage."""

    t1 = Content.from_text("Hello", raw_representation="raw1", additional_properties={"key1": "val1"})
    t2 = Content.from_text(" World", raw_representation="raw2", additional_properties={"key2": "val2"})

    t1 += t2

    # Content doesn't implement __iadd__, so += creates a new object via __add__
    assert t1.text == "Hello World"
    assert t1.raw_representation == ["raw1", "raw2"]
    assert t1.additional_properties == {"key1": "val1", "key2": "val2"}


def test_text_reasoning_content_add_coverage():
    """Test TextReasoningContent __add__ method for better coverage."""

    t1 = Content.from_text_reasoning(text="Thinking 1")
    t2 = Content.from_text_reasoning(text=" Thinking 2")

    result = t1 + t2
    assert result.text == "Thinking 1 Thinking 2"


def test_text_reasoning_content_iadd_coverage():
    """Test TextReasoningContent += operator for better coverage."""

    t1 = Content.from_text_reasoning(text="Thinking 1")
    t2 = Content.from_text_reasoning(text=" Thinking 2")

    t1 += t2

    # Content doesn't implement __iadd__, so += creates a new object via __add__
    assert t1.text == "Thinking 1 Thinking 2"


def test_text_reasoning_content_add_preserves_id():
    """Test that coalescing text_reasoning Content preserves the id field."""

    t1 = Content.from_text_reasoning(id="rs_abc123", text="Thinking part 1")
    t2 = Content.from_text_reasoning(id="rs_abc123", text=" part 2")

    result = t1 + t2
    assert result.text == "Thinking part 1 part 2"
    assert result.id == "rs_abc123"


def test_text_reasoning_content_add_id_fallback_to_other():
    """Test that coalescing falls back to other's id when self has no id."""

    t1 = Content.from_text_reasoning(text="Thinking part 1")
    t2 = Content.from_text_reasoning(id="rs_abc123", text=" part 2")

    result = t1 + t2
    assert result.id == "rs_abc123"


def test_text_reasoning_content_add_preserves_id_with_encrypted_content():
    """Test that id and encrypted_content both survive coalescing for round-trip."""

    t1 = Content.from_text_reasoning(
        id="rs_abc123",
        text="Thinking",
        additional_properties={"encrypted_content": "enc_blob_data"},
    )
    t2 = Content.from_text_reasoning(id="rs_abc123", text=" more")

    result = t1 + t2
    assert result.text == "Thinking more"
    assert result.id == "rs_abc123"
    assert result.additional_properties.get("encrypted_content") == "enc_blob_data"


def test_text_reasoning_content_add_conflicting_ids_raises():
    """Test that coalescing text_reasoning Content with different ids raises AdditionItemMismatch."""

    t1 = Content.from_text_reasoning(id="rs_abc123", text="Thinking part 1")
    t2 = Content.from_text_reasoning(id="rs_xyz789", text=" part 2")

    with pytest.raises(AdditionItemMismatch, match="different ids"):
        _ = t1 + t2


def test_text_reasoning_content_add_neither_has_id():
    """Test that coalescing text_reasoning Content when neither has an id results in None id."""

    t1 = Content.from_text_reasoning(text="Thinking part 1")
    t2 = Content.from_text_reasoning(text=" part 2")

    result = t1 + t2
    assert result.text == "Thinking part 1 part 2"
    assert result.id is None


def test_coalesce_text_reasoning_with_different_ids():
    """Test that _coalesce_text_content keeps separate text_reasoning items when IDs differ.

    Regression test: streaming responses can produce multiple text_reasoning
    segments with distinct IDs. These must not be merged into one.
    """
    from agent_framework._types import _coalesce_text_content

    contents = [
        Content.from_text_reasoning(id="rs_aaa", text="Thinking A1"),
        Content.from_text_reasoning(id="rs_aaa", text=" A2"),
        Content.from_text_reasoning(id="rs_bbb", text="Thinking B1"),
        Content.from_text_reasoning(id="rs_bbb", text=" B2"),
    ]

    _coalesce_text_content(contents, "text_reasoning")

    assert len(contents) == 2
    assert contents[0].id == "rs_aaa"
    assert contents[0].text == "Thinking A1 A2"
    assert contents[1].id == "rs_bbb"
    assert contents[1].text == "Thinking B1 B2"


def test_comprehensive_to_dict_exclude_options():
    """Test to_dict methods with various exclude options for better coverage."""

    # Test TextContent with exclude_none
    text_content = Content.from_text("Hello", raw_representation=None, additional_properties={"prop": "val"})
    text_dict = text_content.to_dict(exclude_none=True)
    assert "raw_representation" not in text_dict
    assert text_dict["additional_properties"]["prop"] == "val"

    # Test with custom exclude set
    text_dict_exclude = text_content.to_dict(exclude={"additional_properties"})
    assert "additional_properties" not in text_dict_exclude
    assert "text" in text_dict_exclude

    # Test UsageDetails - it's a TypedDict now, not a class with to_dict
    usage = UsageDetails(input_token_count=5, custom_count=10)  # type: ignore[typeddict-unknown-key]
    assert usage["input_token_count"] == 5
    assert usage["custom_count"] == 10  # type: ignore[typeddict-item]

    # Test UsageDetails exclude_none behavior isn't applicable to TypedDict
    # TypedDict doesn't have a to_dict method


def test_usage_details_iadd_edge_cases():
    """Test UsageDetails addition with edge cases for better coverage."""
    # Test with None values
    u1 = UsageDetails(input_token_count=None, output_token_count=5, custom1=10)  # type: ignore[typeddict-unknown-key]
    u2 = UsageDetails(input_token_count=3, output_token_count=None, custom2=20)  # type: ignore[typeddict-unknown-key]

    result = add_usage_details(u1, u2)
    assert result["input_token_count"] == 3
    assert result["output_token_count"] == 5
    assert result.get("custom1") == 10
    assert result.get("custom2") == 20

    # Test merging additional counts
    u3 = UsageDetails(input_token_count=1, shared_count=5)  # type: ignore[typeddict-unknown-key]
    u4 = UsageDetails(input_token_count=2, shared_count=15)  # type: ignore[typeddict-unknown-key]

    result2 = add_usage_details(u3, u4)
    assert result2["input_token_count"] == 3
    assert result2.get("shared_count") == 20


def test_chat_message_from_dict_with_mixed_content():
    """Test Message from_dict with mixed content types for better coverage."""

    message_data = {
        "role": "assistant",
        "contents": [
            {"type": "text", "text": "Hello"},
            {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {"arg": "val"}},
            {"type": "function_result", "call_id": "call1", "result": "success"},
        ],
    }

    message = Message.from_dict(message_data)
    assert len(message.contents) == 3  # Unknown type is ignored
    assert message.contents[0].type == "text"
    assert message.contents[1].type == "function_call"
    assert message.contents[2].type == "function_result"

    # Test round-trip
    message_dict = message.to_dict()
    assert len(message_dict["contents"]) == 3


def test_text_content_add_type_error():
    """Test TextContent __add__ raises TypeError for incompatible types."""
    t1 = Content.from_text("Hello")

    with raises(TypeError, match="Incompatible type"):
        t1 + "not a TextContent"  # type: ignore[operator]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[unsupported-operator]


def test_comprehensive_serialization_methods():
    """Test from_dict and to_dict methods for various content types."""

    # Test TextContent with all fields
    text_data = {
        "type": "text",
        "text": "Hello world",
        "raw_representation": {"key": "value"},
        "additional_properties": {"prop": "val"},
        "annotations": None,
    }
    text_content = Content.from_dict(text_data)
    assert text_content.text == "Hello world"
    assert text_content.raw_representation == {"key": "value"}
    assert text_content.additional_properties == {"prop": "val"}

    # Test round-trip
    text_dict = text_content.to_dict()
    assert text_dict["text"] == "Hello world"
    assert text_dict["additional_properties"] == {"prop": "val"}
    # Note: raw_representation is always excluded from to_dict() output

    # Test with exclude_none
    text_dict_no_none = text_content.to_dict(exclude_none=True)
    assert "annotations" not in text_dict_no_none

    # Test FunctionResultContent
    result_data = {
        "type": "function_result",
        "call_id": "call123",
        "result": "success",
        "additional_properties": {"meta": "data"},
    }
    result_content = Content.from_dict(result_data)
    assert result_content.call_id == "call123"
    assert result_content.result == "success"


def test_chat_message_complex_content_serialization():
    """Test Message serialization with various content types."""

    # Create a message with multiple content types
    contents = [
        Content.from_text("Hello"),
        Content.from_function_call(call_id="call1", name="func", arguments={"arg": "val"}),
        Content.from_function_result(call_id="call1", result="success"),
    ]

    message = Message(role="assistant", contents=contents)

    # Test to_dict
    message_dict = message.to_dict()
    assert len(message_dict["contents"]) == 3
    assert message_dict["contents"][0]["type"] == "text"
    assert message_dict["contents"][1]["type"] == "function_call"
    assert message_dict["contents"][2]["type"] == "function_result"

    # Test from_dict round-trip
    reconstructed = Message.from_dict(message_dict)
    assert len(reconstructed.contents) == 3
    assert reconstructed.contents[0].type == "text"
    assert reconstructed.contents[1].type == "function_call"
    assert reconstructed.contents[2].type == "function_result"


def test_message_roundtrip_preserves_compaction_annotation_dict() -> None:
    message = Message(
        role="assistant",
        contents=[Content.from_text("Hello")],
        additional_properties={
            GROUP_ANNOTATION_KEY: {
                "id": "group_1",
                "kind": "assistant_text",
                "index": 1,
                "has_reasoning": False,
                "token_count": 42,
            }
        },
    )

    restored = Message.from_dict(message.to_dict())
    annotation = restored.additional_properties.get(GROUP_ANNOTATION_KEY)

    assert isinstance(annotation, dict)
    assert annotation[GROUP_ID_KEY] == "group_1"
    assert annotation[GROUP_TOKEN_COUNT_KEY] == 42


def test_content_roundtrip_preserves_compaction_annotation_dict() -> None:
    content = Content.from_text(
        text="Hello",
        additional_properties={
            GROUP_ANNOTATION_KEY: {
                "id": "group_2",
                "kind": "assistant_text",
                "index": 2,
                "has_reasoning": False,
                "token_count": None,
            }
        },
    )

    restored = Content.from_dict(content.to_dict())
    annotation = restored.additional_properties.get(GROUP_ANNOTATION_KEY)

    assert isinstance(annotation, dict)
    assert annotation[GROUP_ID_KEY] == "group_2"
    assert annotation[GROUP_TOKEN_COUNT_KEY] is None


def test_content_from_dict_via_json() -> None:
    """Test Content.from_dict with data parsed from a JSON string."""
    data = json.loads(json.dumps({"type": "text", "text": "Hello world"}))
    content = Content.from_dict(data)
    assert content.type == "text"
    assert content.text == "Hello world"


def test_content_from_dict_roundtrip_via_json() -> None:
    """Test Content.from_dict roundtrip via to_dict and json.dumps."""
    original = Content.from_function_call(call_id="call1", name="my_func", arguments={"key": "value"})
    data = json.loads(json.dumps(original.to_dict()))
    restored = Content.from_dict(data)
    assert restored.type == "function_call"
    assert restored.call_id == "call1"
    assert restored.name == "my_func"
    assert restored.arguments == {"key": "value"}


def test_content_to_dict_exclude_none() -> None:
    """Test Content.to_dict excludes None fields by default."""
    content = Content.from_text("Hello")
    d = content.to_dict()
    parsed = json.loads(json.dumps(d))
    assert "uri" not in parsed

    d_with_none = content.to_dict(exclude_none=False)
    parsed_with_none = json.loads(json.dumps(d_with_none))
    assert "uri" in parsed_with_none
    assert parsed_with_none["uri"] is None


def test_content_to_dict_exclude_fields() -> None:
    """Test Content.to_dict with explicit field exclusion."""
    content = Content.from_text("Hello")
    d = content.to_dict(exclude={"text"})
    parsed = json.loads(json.dumps(d))
    assert "text" not in parsed
    assert parsed["type"] == "text"


def test_chat_response_roundtrip_preserves_compaction_annotation_dict() -> None:
    response = ChatResponse(
        messages=[
            Message(
                role="assistant",
                contents=[Content.from_text("Hello")],
                additional_properties={
                    GROUP_ANNOTATION_KEY: {
                        "id": "group_3",
                        "kind": "assistant_text",
                        "index": 3,
                        "has_reasoning": True,
                        "token_count": 15,
                    }
                },
            )
        ]
    )

    restored = ChatResponse.from_dict(response.to_dict())
    annotation = restored.messages[0].additional_properties.get(GROUP_ANNOTATION_KEY)

    assert isinstance(annotation, dict)
    assert annotation[GROUP_ID_KEY] == "group_3"
    assert annotation[GROUP_HAS_REASONING_KEY] is True


def test_usage_content_serialization_with_details():
    """Test UsageContent from_dict and to_dict with UsageDetails conversion."""

    # Test from_dict with details as dict
    usage_data = {
        "type": "usage",
        "usage_details": {
            "type": "usage_details",
            "input_token_count": 10,
            "output_token_count": 20,
            "total_token_count": 30,
            "custom_count": 5,
        },
    }
    usage_content = Content(**usage_data)  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
    assert isinstance(usage_content.usage_details, dict)
    assert usage_content.usage_details["input_token_count"] == 10
    assert usage_content.usage_details["custom_count"] == 5  # type: ignore[typeddict-item]  # Custom fields go directly in UsageDetails

    # Test to_dict with UsageDetails object
    usage_dict = usage_content.to_dict()
    assert isinstance(usage_dict["usage_details"], dict)
    assert usage_dict["usage_details"]["input_token_count"] == 10


def test_function_approval_response_content_serialization():
    """Test FunctionApprovalResponseContent from_dict and to_dict with function_call conversion."""

    # Test from_dict with function_call as dict
    response_data = {
        "type": "function_approval_response",
        "id": "response123",
        "approved": True,
        "function_call": {
            "type": "function_call",
            "call_id": "call123",
            "name": "test_func",
            "arguments": {"param": "value"},
        },
    }
    response_content = Content.from_dict(response_data)
    assert response_content.function_call.type == "function_call"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
    assert response_content.function_call.call_id == "call123"  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]

    # Test to_dict with FunctionCallContent object
    response_dict = response_content.to_dict()
    assert isinstance(response_dict["function_call"], dict)
    assert response_dict["function_call"]["call_id"] == "call123"


def test_chat_response_complex_serialization():
    """Test ChatResponse from_dict and to_dict with complex nested objects."""

    # Test from_dict with messages, finish_reason, and usage_details as dicts
    response_data = {
        "messages": [
            {"role": "user", "contents": [{"type": "text", "text": "Hello"}]},
            {"role": "assistant", "contents": [{"type": "text", "text": "Hi there"}]},
        ],
        "finish_reason": "stop",
        "usage_details": {
            "type": "usage_details",
            "input_token_count": 5,
            "output_token_count": 8,
            "total_token_count": 13,
        },
        "model": "gpt-4",  # Test alias handling
    }

    response = ChatResponse.from_dict(response_data)
    assert len(response.messages) == 2
    assert isinstance(response.messages[0], Message)
    assert isinstance(response.finish_reason, str)  # FinishReason is now a NewType of str
    assert isinstance(response.usage_details, dict)
    assert response.model == "gpt-4"  # Should be stored as model

    # Test to_dict with complex objects
    response_dict = response.to_dict()
    assert len(response_dict["messages"]) == 2
    assert isinstance(response_dict["messages"][0], dict)
    assert isinstance(response_dict["finish_reason"], str)  # FinishReason serializes to string
    assert isinstance(response_dict["usage_details"], dict)
    assert response_dict["model"] == "gpt-4"  # Should serialize as model


def test_chat_response_update_all_content_types():
    """Test ChatResponseUpdate from_dict with all supported content types."""

    update_data = {
        "contents": [
            {"type": "text", "text": "Hello"},
            {"type": "data", "data": b"base64data", "media_type": "text/plain"},
            {"type": "uri", "uri": "http://example.com", "media_type": "text/html"},
            {"type": "error", "message": "An error occurred"},
            {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}},
            {"type": "function_result", "call_id": "call1", "result": "success"},
            {"type": "usage", "usage_details": {"input_token_count": 1}},
            {"type": "hosted_file", "file_id": "file123"},
            {"type": "hosted_vector_store", "vector_store_id": "vs123"},
            {
                "type": "function_approval_request",
                "id": "req1",
                "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}},
            },
            {
                "type": "function_approval_response",
                "id": "resp1",
                "approved": True,
                "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}},
            },
            {"type": "text_reasoning", "text": "reasoning"},
        ]
    }

    update = ChatResponseUpdate.from_dict(update_data)
    assert len(update.contents) == 12  # unknown_type is skipped with warning
    assert update.contents[0].type == "text"
    assert update.contents[1].type == "data"
    assert update.contents[2].type == "uri"
    assert update.contents[3].type == "error"
    assert update.contents[4].type == "function_call"
    assert update.contents[5].type == "function_result"
    assert update.contents[6].type == "usage"
    assert update.contents[7].type == "hosted_file"
    assert update.contents[8].type == "hosted_vector_store"
    assert update.contents[9].type == "function_approval_request"
    assert update.contents[10].type == "function_approval_response"
    assert update.contents[11].type == "text_reasoning"


def test_agent_run_response_complex_serialization():
    """Test AgentResponse from_dict and to_dict with messages and usage_details."""

    response_data = {
        "messages": [
            {"role": "user", "contents": [{"type": "text", "text": "Hello"}]},
            {"role": "assistant", "contents": [{"type": "text", "text": "Hi"}]},
        ],
        "usage_details": {
            "type": "usage_details",
            "input_token_count": 3,
            "output_token_count": 2,
            "total_token_count": 5,
        },
    }

    response = AgentResponse.from_dict(response_data)
    assert len(response.messages) == 2
    assert isinstance(response.messages[0], Message)
    assert isinstance(response.usage_details, dict)

    # Test to_dict
    response_dict = response.to_dict()
    assert len(response_dict["messages"]) == 2
    assert isinstance(response_dict["messages"][0], dict)
    assert isinstance(response_dict["usage_details"], dict)


def test_agent_run_response_update_all_content_types():
    """Test AgentResponseUpdate from_dict with all content types and role handling."""

    update_data = {
        "contents": [
            {"type": "text", "text": "Hello"},
            {"type": "data", "data": b"base64data", "media_type": "text/plain"},
            {"type": "uri", "uri": "http://example.com", "media_type": "text/html"},
            {"type": "error", "message": "An error occurred"},
            {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}},
            {"type": "function_result", "call_id": "call1", "result": "success"},
            {"type": "usage", "usage_details": {"input_token_count": 1}},
            {"type": "hosted_file", "file_id": "file123"},
            {"type": "hosted_vector_store", "vector_store_id": "vs123"},
            {
                "type": "function_approval_request",
                "id": "req1",
                "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}},
            },
            {
                "type": "function_approval_response",
                "id": "resp1",
                "approved": True,
                "function_call": {"type": "function_call", "call_id": "call1", "name": "func", "arguments": {}},
            },
            {"type": "text_reasoning", "text": "reasoning"},
        ],
        "role": "assistant",  # Test role as dict
    }

    update = AgentResponseUpdate.from_dict(update_data)
    assert len(update.contents) == 12  # unknown_type is logged and ignored
    assert isinstance(update.role, str)  # Role is now a NewType of str
    assert update.role == "assistant"

    # Test to_dict with role conversion
    update_dict = update.to_dict()
    assert len(update_dict["contents"]) == 12  # unknown_type was ignored during from_dict
    assert isinstance(update_dict["role"], str)  # Role serializes to string

    # Test role as string conversion
    update_data_str_role = update_data.copy()
    update_data_str_role["role"] = "user"
    update_str = AgentResponseUpdate.from_dict(update_data_str_role)
    assert isinstance(update_str.role, str)  # Role is now a NewType of str
    assert update_str.role == "user"


# region DeepCopy


class _NonCopyableRaw:
    """Simulates an LLM SDK response object that cannot be deep-copied (e.g., proto/gRPC)."""

    def __deepcopy__(self, memo: dict) -> Any:
        raise TypeError("Cannot deepcopy this object")


def test_content_deepcopy_preserves_raw_representation():
    """Test that deepcopy of Content keeps raw_representation by reference."""
    import copy

    raw = _NonCopyableRaw()
    content = Content.from_text("hello", raw_representation=raw)

    cloned = copy.deepcopy(content)

    assert cloned.text == "hello"
    assert cloned.raw_representation is raw
    assert cloned.additional_properties is not content.additional_properties


def test_message_deepcopy_preserves_raw_representation():
    """Test that deepcopy of Message keeps raw_representation by reference."""
    import copy

    raw = _NonCopyableRaw()
    msg = Message("assistant", ["hello"], raw_representation=raw)

    cloned = copy.deepcopy(msg)

    assert cloned.text == "hello"
    assert cloned.raw_representation is raw
    assert cloned.contents is not msg.contents


def test_agent_response_deepcopy_preserves_raw_representation():
    """Test that deepcopy of AgentResponse keeps raw_representation by reference."""
    import copy

    raw = _NonCopyableRaw()
    response = AgentResponse(
        messages=[Message("assistant", ["test"])],
        raw_representation=raw,
    )

    cloned = copy.deepcopy(response)

    assert cloned.text == "test"
    assert cloned.raw_representation is raw
    assert cloned.messages is not response.messages


def test_chat_response_deepcopy_preserves_raw_representation():
    """Test that deepcopy of ChatResponse keeps raw_representation by reference."""
    import copy

    raw = _NonCopyableRaw()
    response = ChatResponse(
        messages=[Message("assistant", ["test"])],
        raw_representation=raw,
    )

    cloned = copy.deepcopy(response)

    assert cloned.text == "test"
    assert cloned.raw_representation is raw
    assert cloned.messages is not response.messages


def test_chat_response_update_deepcopy_preserves_raw_representation():
    """Test that deepcopy of ChatResponseUpdate keeps raw_representation by reference."""
    import copy

    raw = _NonCopyableRaw()
    update = ChatResponseUpdate(
        contents=[Content.from_text("hello")],
        role="assistant",
        raw_representation=raw,
    )

    cloned = copy.deepcopy(update)

    assert cloned.text == "hello"
    assert cloned.raw_representation is raw
    assert cloned.contents is not update.contents


def test_agent_response_update_deepcopy_preserves_raw_representation():
    """Test that deepcopy of AgentResponseUpdate keeps raw_representation by reference."""
    import copy

    raw = _NonCopyableRaw()
    update = AgentResponseUpdate(
        contents=[Content.from_text("hello")],
        role="assistant",
        raw_representation=raw,
    )

    cloned = copy.deepcopy(update)

    assert cloned.text == "hello"
    assert cloned.raw_representation is raw
    assert cloned.contents is not update.contents


def test_nested_deepcopy_preserves_raw_representation():
    """Test that deepcopy of an AgentResponse with nested Message raw_representations works."""
    import copy

    raw_msg = _NonCopyableRaw()
    raw_response = _NonCopyableRaw()
    response = AgentResponse(
        messages=[Message("assistant", ["hello"], raw_representation=raw_msg)],
        raw_representation=raw_response,
    )

    cloned = copy.deepcopy(response)

    assert cloned.raw_representation is raw_response
    assert cloned.messages[0].raw_representation is raw_msg
    assert cloned.messages is not response.messages
    assert cloned.text == "hello"


def test_content_deepcopy_shallow_copy_fields_identity():
    """Test that Content._SHALLOW_COPY_FIELDS fields are identity-preserved while others are deep-copied."""
    import copy

    raw = _NonCopyableRaw()
    content = Content.from_text("hello", raw_representation=raw)
    content.additional_properties["key"] = "value"

    cloned = copy.deepcopy(content)

    # _SHALLOW_COPY_FIELDS (raw_representation) should be same object
    assert cloned.raw_representation is raw
    # Non-shallow fields should be independent deep copies
    assert cloned.additional_properties is not content.additional_properties
    assert cloned.additional_properties == {"key": "value"}


def test_chat_response_deepcopy_deep_copies_additional_properties():
    """Test that ChatResponse deepcopy deep-copies additional_properties despite it being in DEFAULT_EXCLUDE."""
    import copy

    response = ChatResponse(
        messages=[Message("assistant", ["test"])],
        additional_properties={"key": [1, 2, 3]},
    )

    cloned = copy.deepcopy(response)

    # additional_properties is in DEFAULT_EXCLUDE for serialization but not in _SHALLOW_COPY_FIELDS,
    # so it should be deep-copied (independent copy)
    assert cloned.additional_properties is not response.additional_properties
    assert cloned.additional_properties == {"key": [1, 2, 3]}


# endregion


# region Serialization


@mark.parametrize(
    "content_class,init_kwargs",
    [
        pytest.param(
            Content,
            {
                "type": "text",
                "text": "Hello world",
                "raw_representation": "raw",
            },
            id="text_content",
        ),
        pytest.param(
            Content,
            {
                "type": "text_reasoning",
                "text": "Reasoning text",
                "raw_representation": "raw",
            },
            id="text_reasoning_content",
        ),
        pytest.param(
            Content,
            {
                "type": "data",
                "uri": "data:text/plain;base64,dGVzdCBkYXRh",
            },
            id="data_content_with_uri",
        ),
        pytest.param(
            Content,
            {
                "type": "data",
                "data": b"test data",
                "media_type": "text/plain",
            },
            id="data_content_with_bytes",
        ),
        pytest.param(
            Content,
            {
                "type": "uri",
                "uri": "http://example.com",
                "media_type": "text/html",
            },
            id="uri_content",
        ),
        pytest.param(
            Content,
            {"type": "hosted_file", "file_id": "file-123"},
            id="hosted_file_content",
        ),
        pytest.param(
            Content,
            {
                "type": "hosted_vector_store",
                "vector_store_id": "vs-789",
            },
            id="hosted_vector_store_content",
        ),
        pytest.param(
            Content,
            {
                "type": "function_call",
                "call_id": "call-1",
                "name": "test_func",
                "arguments": {"arg": "val"},
            },
            id="function_call_content",
        ),
        pytest.param(
            Content,
            {
                "type": "function_result",
                "call_id": "call-1",
                "result": "success",
            },
            id="function_result_content",
        ),
        pytest.param(
            Content,
            {
                "type": "error",
                "message": "Error occurred",
                "error_code": "E001",
            },
            id="error_content",
        ),
        pytest.param(
            Content,
            {
                "type": "usage",
                "usage_details": {
                    "type": "usage_details",
                    "input_token_count": 10,
                    "output_token_count": 20,
                    "reasoning_tokens": 5,
                },
            },
            id="usage_content",
        ),
        pytest.param(
            Content,
            {
                "type": "function_approval_request",
                "id": "req-1",
                "function_call": {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}},
            },
            id="function_approval_request",
        ),
        pytest.param(
            Content,
            {
                "type": "function_approval_response",
                "id": "resp-1",
                "approved": True,
                "function_call": {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}},
            },
            id="function_approval_response",
        ),
        pytest.param(
            Message,
            {
                "role": "\1",
                "contents": [
                    {"type": "text", "text": "Hello"},
                    {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}},
                ],
                "message_id": "msg-123",
                "author_name": "User",
            },
            id="chat_message",
        ),
        pytest.param(
            ChatResponse,
            {
                "type": "chat_response",
                "messages": [
                    {
                        "type": "message",
                        "role": "\1",
                        "contents": [{"type": "text", "text": "Hello"}],
                    },
                    {
                        "type": "message",
                        "role": "\1",
                        "contents": [{"type": "text", "text": "Hi there"}],
                    },
                ],
                "finish_reason": "\1",
                "usage_details": {
                    "type": "usage_details",
                    "input_token_count": 10,
                    "output_token_count": 20,
                    "total_token_count": 30,
                },
                "response_id": "resp-123",
                "model": "gpt-4",
            },
            id="chat_response",
        ),
        pytest.param(
            ChatResponseUpdate,
            {
                "contents": [
                    {"type": "text", "text": "Hello"},
                    {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}},
                ],
                "role": "\1",
                "finish_reason": "\1",
                "message_id": "msg-123",
                "response_id": "resp-123",
            },
            id="chat_response_update",
        ),
        pytest.param(
            AgentResponse,
            {
                "messages": [
                    {
                        "role": "\1",
                        "contents": [{"type": "text", "text": "Question"}],
                    },
                    {
                        "role": "\1",
                        "contents": [{"type": "text", "text": "Answer"}],
                    },
                ],
                "response_id": "run-123",
                "usage_details": {
                    "type": "usage_details",
                    "input_token_count": 5,
                    "output_token_count": 3,
                    "total_token_count": 8,
                },
            },
            id="agent_response",
        ),
        pytest.param(
            AgentResponseUpdate,
            {
                "contents": [
                    {"type": "text", "text": "Streaming"},
                    {"type": "function_call", "call_id": "call-1", "name": "test_func", "arguments": {}},
                ],
                "role": "\1",
                "message_id": "msg-123",
                "response_id": "run-123",
                "author_name": "Agent",
            },
            id="agent_response_update",
        ),
    ],
)
def test_content_roundtrip_serialization(content_class: type[Content], init_kwargs: dict[str, Any]):
    """Test to_dict/from_dict roundtrip for all content types."""
    # Create instance using from_dict to handle nested dict-to-object conversions
    content = content_class.from_dict(init_kwargs)

    # Serialize to dict
    content_dict = content.to_dict()

    # Verify type key is in serialized dict
    assert "type" in content_dict
    if hasattr(content, "type"):
        assert content_dict["type"] == content.type  # type: ignore[attr-defined]

    # Deserialize from dict
    reconstructed = content_class.from_dict(content_dict)

    # Verify type
    assert isinstance(reconstructed, content_class)
    # Check type attribute dynamically
    if hasattr(content, "type"):
        assert reconstructed.type == content.type  # type: ignore[attr-defined]

    # Verify key attributes (excluding raw_representation which is not serialized)
    for key, value in init_kwargs.items():
        if key == "type":
            continue
        if key == "raw_representation":
            # raw_representation is intentionally excluded from serialization
            continue

        # Special handling for DataContent created with 'data' parameter
        if hasattr(content, "type") and content.type == "data" and key == "data":
            # DataContent converts 'data' to 'uri', so we skip checking 'data' attribute
            # Instead we verify that uri and media_type are set correctly
            assert hasattr(reconstructed, "uri")
            assert hasattr(reconstructed, "media_type")
            assert reconstructed.media_type == init_kwargs.get("media_type")
            # Verify the uri contains the encoded data
            assert reconstructed.uri.startswith(f"data:{init_kwargs.get('media_type')};base64,")  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
            continue

        reconstructed_value = getattr(reconstructed, key)

        # Special handling for nested SerializationMixin objects
        if hasattr(value, "to_dict"):
            # Compare the serialized forms
            assert reconstructed_value.to_dict() == value.to_dict()
        # Special handling for lists that may contain dicts converted to objects
        elif isinstance(value, list) and value and isinstance(reconstructed_value, list):
            # Check if this is a list of objects that were created from dicts
            if isinstance(value[0], dict) and hasattr(reconstructed_value[0], "to_dict"):
                # Compare each item by serializing the reconstructed object
                assert len(reconstructed_value) == len(value)
                for orig_dict, recon_obj in zip(value, reconstructed_value):
                    recon_dict = recon_obj.to_dict()  # ty: ignore[unresolved-attribute]
                    # Compare all keys from original dict (reconstructed may have extra default fields)
                    for k, v in orig_dict.items():  # ty: ignore[unresolved-attribute]
                        assert k in recon_dict, f"Key '{k}' missing from reconstructed dict"
                        # For nested lists, recursively compare
                        if isinstance(v, list) and v and isinstance(v[0], dict):
                            assert len(recon_dict[k]) == len(v)
                            for orig_item, recon_item in zip(v, recon_dict[k]):
                                # Compare essential keys, ignoring fields like additional_properties
                                for item_key, item_val in orig_item.items():  # ty: ignore[unresolved-attribute]
                                    assert item_key in recon_item
                                    assert recon_item[item_key] == item_val
                        else:
                            assert recon_dict[k] == v, f"Value mismatch for key '{k}'"
            else:
                assert reconstructed_value == value
        # Special handling for dicts that get converted to objects (like UsageDetails, FunctionCallContent)
        elif isinstance(value, dict) and hasattr(reconstructed_value, "to_dict"):
            # Compare the dict with the serialized form of the object
            reconstructed_dict = reconstructed_value.to_dict()
            # Verify all keys from the original dict are in the reconstructed dict
            for k, v in value.items():
                assert k in reconstructed_dict, f"Key '{k}' missing from reconstructed dict"
                assert reconstructed_dict[k] == v, f"Value mismatch for key '{k}'"
        else:
            assert reconstructed_value == value


def test_text_content_with_annotations_serialization():
    """Test TextContent with multiple annotations roundtrip serialization."""
    # Create multiple regions
    region1 = TextSpanRegion(type="text_span", start_index=0, end_index=5)
    region2 = TextSpanRegion(type="text_span", start_index=6, end_index=11)

    # Create multiple citations
    citation1 = Annotation(type="citation", title="Citation 1", url="http://example.com/1", annotated_regions=[region1])

    citation2 = Annotation(type="citation", title="Citation 2", url="http://example.com/2", annotated_regions=[region2])

    # Create TextContent with multiple annotations
    content = Content.from_text(text="Hello world", annotations=[citation1, citation2])

    # Serialize
    content_dict = content.to_dict()

    # Verify we have 2 annotations
    assert len(content_dict["annotations"]) == 2
    assert content_dict["annotations"][0]["title"] == "Citation 1"
    assert content_dict["annotations"][1]["title"] == "Citation 2"

    # Deserialize
    reconstructed = Content.from_dict(content_dict)

    # Verify reconstruction
    assert len(reconstructed.annotations) == 2  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
    assert reconstructed.annotations is not None
    # Annotation are TypedDicts (dicts at runtime)
    assert all(isinstance(ann, dict) for ann in reconstructed.annotations)  # type: ignore[union-attr]  # pyrefly: ignore[not-iterable]
    assert reconstructed.annotations[0]["title"] == "Citation 1"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]
    assert reconstructed.annotations[1]["title"] == "Citation 2"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]
    assert all(isinstance(ann["annotated_regions"][0], dict) for ann in reconstructed.annotations)  # type: ignore[union-attr]  # pyrefly: ignore[not-iterable]


# region FunctionTool.parse_result with Pydantic models


class WeatherResult(BaseModel):
    """A Pydantic model for testing."""

    temperature: float
    condition: str


class NestedModel(BaseModel):
    """A Pydantic model with nested structure."""

    name: str
    weather: WeatherResult


def test_parse_result_pydantic_model():
    """Test that Pydantic BaseModel subclasses are properly serialized using model_dump()."""
    result = WeatherResult(temperature=22.5, condition="sunny")
    parsed = FunctionTool.parse_result(result)

    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"
    assert '"temperature": 22.5' in parsed[0].text or '"temperature":22.5' in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
    assert '"condition": "sunny"' in parsed[0].text or '"condition":"sunny"' in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]


def test_parse_result_pydantic_model_in_list():
    """Test that lists containing Pydantic models are properly serialized."""
    results = [
        WeatherResult(temperature=20.0, condition="cloudy"),
        WeatherResult(temperature=25.0, condition="sunny"),
    ]
    parsed = FunctionTool.parse_result(results)

    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"
    assert parsed[0].text.startswith("[")  # type: ignore[union-attr]  # ty: ignore[unresolved-attribute]
    assert "cloudy" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
    assert "sunny" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]


def test_parse_result_pydantic_model_in_dict():
    """Test that dicts containing Pydantic models are properly serialized."""
    results = {
        "current": WeatherResult(temperature=22.0, condition="partly cloudy"),
        "forecast": WeatherResult(temperature=24.0, condition="sunny"),
    }
    parsed = FunctionTool.parse_result(results)

    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"
    assert "current" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
    assert "forecast" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
    assert "partly cloudy" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
    assert "sunny" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]


def test_parse_result_nested_pydantic_model():
    """Test that nested Pydantic models are properly serialized."""
    result = NestedModel(name="Seattle", weather=WeatherResult(temperature=18.0, condition="rainy"))
    parsed = FunctionTool.parse_result(result)

    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"
    assert "Seattle" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
    assert "rainy" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
    assert "18.0" in parsed[0].text or "18" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]


# region FunctionTool.parse_result with MCP TextContent-like objects


def test_parse_result_text_content_single():
    """Test that objects with text attribute (like MCP TextContent) are properly handled."""

    @dataclass
    class MockTextContent:
        text: str

    result = [MockTextContent("Hello from MCP tool!")]
    parsed = FunctionTool.parse_result(result)

    # Non-Content list items are serialized via _make_dumpable
    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"


def test_parse_result_text_content_multiple():
    """Test that multiple TextContent-like objects are serialized correctly."""

    @dataclass
    class MockTextContent:
        text: str

    result = [MockTextContent("First result"), MockTextContent("Second result")]
    parsed = FunctionTool.parse_result(result)

    # Non-Content list items are serialized via _make_dumpable
    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"


def test_parse_result_text_content_with_non_string_text():
    """Test that objects with non-string text attribute are not treated as TextContent."""

    class BadTextContent:
        def __init__(self):
            self.text = 12345  # Not a string!

    result = [BadTextContent()]
    parsed = FunctionTool.parse_result(result)

    # Should not extract text since it's not a string, will serialize the object
    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"


def test_parse_result_none_returns_empty_string():
    """Test that None returns a list with empty text Content."""
    parsed = FunctionTool.parse_result(None)
    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].type == "text"
    assert parsed[0].text == ""


def test_parse_result_string_passthrough():
    """Test that strings are wrapped in Content."""
    parsed = FunctionTool.parse_result("hello world")
    assert isinstance(parsed, list)
    assert len(parsed) == 1
    assert parsed[0].text == "hello world"

    parsed2 = FunctionTool.parse_result('{"key": "value"}')
    assert isinstance(parsed2, list)
    assert len(parsed2) == 1
    assert parsed2[0].text == '{"key": "value"}'


def test_parse_result_content_object():
    """Test that text Content objects are wrapped in a list."""
    content = Content.from_text("hello")
    result = FunctionTool.parse_result(content)
    assert isinstance(result, list)
    assert len(result) == 1
    assert result[0].type == "text"
    assert result[0].text == "hello"


def test_parse_result_list_of_content():
    """Test that list[Content] with text-only items is returned as list[Content]."""
    contents = [Content.from_text("hello"), Content.from_text("world")]
    result = FunctionTool.parse_result(contents)
    assert isinstance(result, list)
    assert len(result) == 2
    assert result[0].text == "hello"
    assert result[1].text == "world"


def test_parse_result_single_image_content():
    """Test that a single image Content is preserved as list[Content]."""
    image_content = Content.from_data(data=b"fake_png_bytes", media_type="image/png")
    result = FunctionTool.parse_result(image_content)
    assert isinstance(result, list)
    assert len(result) == 1
    assert result[0].type == "data"
    assert result[0].media_type == "image/png"


def test_parse_result_single_text_content():
    """Test that a single text Content returns a list with one text Content."""
    text_content = Content.from_text("just text")
    result = FunctionTool.parse_result(text_content)
    assert isinstance(result, list)
    assert len(result) == 1
    assert result[0].type == "text"
    assert result[0].text == "just text"


def test_parse_result_mixed_content_list():
    """Test that list with text and image Content is preserved."""
    contents = [
        Content.from_text("Chart rendered."),
        Content.from_data(data=b"image_bytes", media_type="image/png"),
    ]
    result = FunctionTool.parse_result(contents)
    assert isinstance(result, list)
    assert len(result) == 2
    assert result[0].type == "text"
    assert result[1].type == "data"


def test_from_function_result_with_content_list():
    """Test Content.from_function_result stores all items uniformly."""
    content_list = [
        Content.from_text("Chart rendered."),
        Content.from_data(data=b"image_bytes", media_type="image/png"),
    ]
    result = Content.from_function_result(call_id="test-123", result=content_list)
    assert result.type == "function_result"
    assert result.call_id == "test-123"
    assert result.result == "Chart rendered."
    assert result.items is not None
    assert len(result.items) == 2
    assert result.items[0].type == "text"
    assert result.items[0].text == "Chart rendered."
    assert result.items[1].type == "data"
    assert result.items[1].media_type == "image/png"


def test_from_function_result_with_string():
    """Test Content.from_function_result with plain string result."""
    result = Content.from_function_result(call_id="test-123", result="just text")
    assert result.type == "function_result"
    assert result.call_id == "test-123"
    assert result.result == "just text"
    assert result.items is not None
    assert len(result.items) == 1
    assert result.items[0].type == "text"
    assert result.items[0].text == "just text"


def test_content_from_function_result_items_in_to_dict():
    """Test that items are included in to_dict serialization."""
    content_list = [
        Content.from_text("done"),
        Content.from_data(data=b"png_data", media_type="image/png"),
    ]
    result = Content.from_function_result(
        call_id="call-1",
        result=content_list,
    )
    d = result.to_dict()
    assert "items" in d
    assert len(d["items"]) == 2
    assert d["items"][0]["type"] == "text"
    assert d["items"][1]["type"] == "data"


def test_from_function_result_with_only_rich_content_list():
    """Test Content.from_function_result with only image items and no text."""
    content_list = [
        Content.from_data(data=b"image_bytes", media_type="image/png"),
    ]
    result = Content.from_function_result(call_id="test-456", result=content_list)
    assert result.type == "function_result"
    assert result.result == ""
    assert result.items is not None
    assert len(result.items) == 1
    assert result.items[0].type == "data"


def test_function_result_items_roundtrip_via_dict():
    """Test that items survive a to_dict/from_dict round-trip as Content objects."""
    content_list = [
        Content.from_text("done"),
        Content.from_data(data=b"png_data", media_type="image/png"),
    ]
    original = Content.from_function_result(call_id="call-rt", result=content_list)
    restored = Content.from_dict(original.to_dict())
    assert restored.items is not None
    assert len(restored.items) == 2
    assert isinstance(restored.items[0], Content)
    assert restored.items[0].type == "text"
    assert restored.items[0].text == "done"
    assert isinstance(restored.items[1], Content)
    assert restored.items[1].type == "data"


def test_from_function_result_with_non_content_list():
    """Test Content.from_function_result with a list of non-Content objects falls back to str."""
    result = Content.from_function_result(call_id="test-789", result=["hello", "world"])
    assert result.type == "function_result"
    assert result.result == "['hello', 'world']"
    assert result.items is not None
    assert len(result.items) == 1
    assert result.items[0].type == "text"


# endregion


# region Test Content._add_usage_content


def test_content_add_usage_content():
    """Test adding two usage content instances combines their usage details."""
    usage1 = Content(
        type="usage",
        usage_details={"input_token_count": 100, "output_token_count": 50},
        raw_representation="raw1",
    )
    usage2 = Content(
        type="usage",
        usage_details={"input_token_count": 200, "output_token_count": 100},
        raw_representation="raw2",
    )

    result = usage1 + usage2

    assert result.type == "usage"
    assert result.usage_details["input_token_count"] == 300  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
    assert result.usage_details["output_token_count"] == 150  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
    # Raw representations should be combined
    assert isinstance(result.raw_representation, list)
    assert "raw1" in result.raw_representation
    assert "raw2" in result.raw_representation


def test_content_add_usage_content_with_none_raw_representation():
    """Test adding usage content when one has None raw_representation."""
    usage1 = Content(
        type="usage",
        usage_details={"input_token_count": 100},
        raw_representation=None,
    )
    usage2 = Content(
        type="usage",
        usage_details={"output_token_count": 50},
        raw_representation="raw2",
    )

    result = usage1 + usage2

    assert result.raw_representation == "raw2"


def test_content_add_usage_content_non_integer_values():
    """Test adding usage content with non-integer values."""
    usage1 = Content(
        type="usage",
        usage_details=cast(UsageDetails, {"model": "gpt-4", "count": 10}),
    )
    usage2 = Content(
        type="usage",
        usage_details=cast(UsageDetails, {"model": "gpt-3.5", "count": 20}),
    )

    result = usage1 + usage2

    # Non-integer "model" should take first non-None value
    assert result.usage_details is not None
    assert "model" not in result.usage_details  # type: ignore[operator]  # pyrefly: ignore[not-iterable]
    # Integer "count" should be summed
    assert result.usage_details["count"] == 30  # type: ignore[index, typeddict-item]  # pyrefly: ignore[unsupported-operation]


# endregion


# region Test Content.has_top_level_media_type


def test_content_has_top_level_media_type():
    """Test has_top_level_media_type returns correct boolean."""
    image = Content(type="uri", uri="https://example.com/image.png", media_type="image/png")

    assert image.has_top_level_media_type("image") is True
    assert image.has_top_level_media_type("IMAGE") is True  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]  # Case insensitive
    assert image.has_top_level_media_type("audio") is False


def test_content_has_top_level_media_type_no_slash():
    """Test has_top_level_media_type when media_type has no slash."""
    content = Content(type="data", media_type="text")

    assert content.has_top_level_media_type("text") is True


def test_content_has_top_level_media_type_raises_without_media_type():
    """Test has_top_level_media_type raises ContentError when no media_type."""
    content = Content(type="text", text="hello")

    with raises(ContentError, match="no media_type found"):
        content.has_top_level_media_type("text")


# endregion


# region Test Content.parse_arguments


def test_content_parse_arguments_none():
    """Test parse_arguments returns None when arguments is None."""
    content = Content(type="function_call", call_id="1", name="test", arguments=None)

    assert content.parse_arguments() is None


def test_content_parse_arguments_empty_string():
    """Test parse_arguments returns empty dict for empty string."""
    content = Content(type="function_call", call_id="1", name="test", arguments="")

    assert content.parse_arguments() == {}


def test_content_parse_arguments_valid_json():
    """Test parse_arguments parses valid JSON string."""
    content = Content(type="function_call", call_id="1", name="test", arguments='{"key": "value"}')

    result = content.parse_arguments()
    assert result == {"key": "value"}


def test_content_parse_arguments_non_dict_json():
    """Test parse_arguments wraps non-dict JSON in 'raw' key."""
    content = Content(type="function_call", call_id="1", name="test", arguments='"just a string"')

    result = content.parse_arguments()
    # The JSON is parsed, and if it's not a dict, wrapped in 'raw'
    assert result == {"raw": "just a string"}


def test_content_parse_arguments_invalid_json():
    """Test parse_arguments wraps invalid JSON in 'raw' key."""
    content = Content(type="function_call", call_id="1", name="test", arguments="not json at all")

    result = content.parse_arguments()
    assert result == {"raw": "not json at all"}


def test_content_parse_arguments_dict_passthrough():
    """Test parse_arguments passes through dict arguments."""
    args = {"key": "value", "num": 42}
    content = Content(type="function_call", call_id="1", name="test", arguments=args)

    result = content.parse_arguments()
    assert result == args


# endregion


# region Test _get_data_bytes_as_str


def test_get_data_bytes_as_str_non_data_uri():
    """Test _get_data_bytes_as_str returns None for non-data URIs."""
    content = Content(type="uri", uri="https://example.com/image.png")
    assert _get_data_bytes_as_str(content) is None


def test_get_data_bytes_as_str_no_base64():
    """Test _get_data_bytes_as_str raises for non-base64 data URI."""
    content = Content(type="uri", uri="data:text/plain,hello")
    with raises(ContentError, match="base64 encoding"):
        _get_data_bytes_as_str(content)


def test_get_data_bytes_as_str_valid():
    """Test _get_data_bytes_as_str extracts base64 data."""
    data = base64.b64encode(b"hello").decode()
    content = Content(type="uri", uri=f"data:text/plain;base64,{data}")
    result = _get_data_bytes_as_str(content)
    assert result == data


# endregion


# region Test _get_data_bytes


def test_get_data_bytes_decodes_base64():
    """Test _get_data_bytes decodes base64 data correctly."""
    original = b"hello world"
    data = base64.b64encode(original).decode()
    content = Content(type="uri", uri=f"data:text/plain;base64,{data}")

    result = _get_data_bytes(content)
    assert result == original


def test_get_data_bytes_invalid_base64():
    """Test _get_data_bytes raises for invalid base64."""
    content = Content(type="uri", uri="data:text/plain;base64,!!invalid!!")
    with raises(ContentError, match="Failed to decode"):
        _get_data_bytes(content)


# endregion


# region Test _parse_content_list


def test_parse_content_list_with_content_objects():
    """Test _parse_content_list passes through Content objects."""
    content = Content(type="text", text="hello")
    result = _parse_content_list([content])

    assert len(result) == 1
    assert result[0] is content


def test_parse_content_list_with_dicts():
    """Test _parse_content_list converts dicts to Content."""
    result = _parse_content_list([{"type": "text", "text": "hello"}])

    assert len(result) == 1
    assert result[0].type == "text"
    assert result[0].text == "hello"


def test_parse_content_list_with_mixed_content_and_dict():
    """Test _parse_content_list handles a mix of Content objects and dicts."""
    content = Content(type="text", text="hello")
    # Pass a mix of Content object and dict
    result = _parse_content_list([content, {"type": "text", "text": "world"}])

    assert len(result) == 2
    assert result[0].text == "hello"
    assert result[1].text == "world"


# endregion


# region Test _validate_uri


def test_validate_uri_known_scheme():
    """Test _validate_uri accepts known URI schemes."""
    result = _validate_uri("https://example.com/file.txt", "text/plain")
    assert result.get("uri") == "https://example.com/file.txt"


def test_validate_uri_data_uri():
    """Test _validate_uri handles data URIs."""
    data = base64.b64encode(b"test").decode()
    uri = f"data:text/plain;base64,{data}"
    result = _validate_uri(uri, None)
    assert "uri" in result


# endregion


# region ResponseStream


async def _generate_updates(count: int = 5) -> AsyncIterable[ChatResponseUpdate]:
    """Helper to generate test updates."""
    for i in range(count):
        yield ChatResponseUpdate(contents=[Content.from_text(f"update_{i}")], role="assistant")


def _combine_updates(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
    """Helper finalizer that combines updates into a response."""
    return ChatResponse.from_updates(updates)


class TestResponseStreamBasicIteration:
    """Tests for basic ResponseStream iteration."""

    async def test_iterate_collects_updates(self) -> None:
        """Iterating through stream collects all updates."""
        stream = ResponseStream(_generate_updates(3), finalizer=_combine_updates)

        collected: list[str] = []
        async for update in stream:
            collected.append(update.text or "")

        assert collected == ["update_0", "update_1", "update_2"]
        assert len(stream.updates) == 3

    async def test_stream_consumed_after_iteration(self) -> None:
        """Stream is marked consumed after full iteration."""
        stream = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        async for _ in stream:
            pass

        assert stream._consumed is True

    async def test_get_final_response_after_iteration(self) -> None:
        """Can get final response after iterating."""
        stream = ResponseStream(_generate_updates(3), finalizer=_combine_updates)

        async for _ in stream:
            pass

        final = await stream.get_final_response()
        assert final.text == "update_0update_1update_2"

    async def test_get_final_response_without_iteration(self) -> None:
        """get_final_response auto-iterates if not consumed."""
        stream = ResponseStream(_generate_updates(3), finalizer=_combine_updates)

        final = await stream.get_final_response()

        assert final.text == "update_0update_1update_2"
        assert stream._consumed is True

    async def test_updates_property_returns_collected(self) -> None:
        """updates property returns collected updates."""
        stream = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        async for _ in stream:
            pass

        assert len(stream.updates) == 2
        assert stream.updates[0].text == "update_0"
        assert stream.updates[1].text == "update_1"

    async def test_auto_finalize_on_iteration_completion(self) -> None:
        """Stream auto-finalizes when async iteration completes."""
        stream = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        async for _ in stream:
            pass

        assert stream._finalized is True
        assert stream._final_result is not None
        assert stream._final_result.text == "update_0update_1"

    async def test_auto_finalize_runs_result_hooks(self) -> None:
        """Result hooks run automatically when iteration completes."""
        hook_called = {"value": False}

        def tracking_hook(response: ChatResponse) -> ChatResponse:
            hook_called["value"] = True
            response.additional_properties["auto_finalized"] = True
            return response

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            result_hooks=[tracking_hook],  # ty: ignore[invalid-argument-type]
        )

        async for _ in stream:
            pass

        assert hook_called["value"] is True
        final = await stream.get_final_response()
        assert final.additional_properties["auto_finalized"] is True

    async def test_get_final_response_idempotent_after_auto_finalize(self) -> None:
        """get_final_response returns cached result after auto-finalization."""
        call_count = {"value": 0}

        def counting_finalizer(updates: list[ChatResponseUpdate]) -> ChatResponse:
            call_count["value"] += 1
            return _combine_updates(updates)

        stream = ResponseStream(_generate_updates(2), finalizer=counting_finalizer)  # type: ignore[arg-type, var-annotated]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

        async for _ in stream:
            pass

        final1 = await stream.get_final_response()
        final2 = await stream.get_final_response()

        assert call_count["value"] == 1
        assert final1.text == final2.text


class TestResponseStreamTransformHooks:
    """Tests for transform hooks (per-update processing)."""

    async def test_transform_hook_called_for_each_update(self) -> None:
        """Transform hook is called for each update during iteration."""
        call_count = {"value": 0}

        def counting_hook(update: ChatResponseUpdate) -> None:
            call_count["value"] += 1

        stream = ResponseStream(
            _generate_updates(3),
            finalizer=_combine_updates,
            transform_hooks=[counting_hook],  # ty: ignore[invalid-argument-type]
        )

        await stream.get_final_response()

        assert call_count["value"] == 3

    async def test_transform_hook_can_modify_update(self) -> None:
        """Transform hook can modify the update."""

        def uppercase_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
            return ChatResponseUpdate(
                contents=[Content.from_text((update.text or "").upper())],
                role=cast(Any, update.role),
            )

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            transform_hooks=[uppercase_hook],  # ty: ignore[invalid-argument-type]
        )

        collected: list[str] = []
        async for update in stream:
            collected.append(update.text or "")  # ty: ignore[unresolved-attribute]

        assert collected == ["UPDATE_0", "UPDATE_1"]

    async def test_multiple_transform_hooks_chained(self) -> None:
        """Multiple transform hooks are called in order."""
        order: list[str] = []

        def hook_a(update: ChatResponseUpdate) -> ChatResponseUpdate:
            order.append("a")
            return update

        def hook_b(update: ChatResponseUpdate) -> ChatResponseUpdate:
            order.append("b")
            return update

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            transform_hooks=[hook_a, hook_b],  # ty: ignore[invalid-argument-type]
        )

        async for _ in stream:
            pass

        assert order == ["a", "b", "a", "b"]

    async def test_transform_hook_returning_none_keeps_previous(self) -> None:
        """Transform hook returning None keeps the previous value."""

        def none_hook(update: ChatResponseUpdate) -> None:
            return None

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            transform_hooks=[none_hook],  # ty: ignore[invalid-argument-type]
        )

        collected: list[str] = []
        async for update in stream:
            collected.append(update.text or "")  # ty: ignore[unresolved-attribute]

        assert collected == ["update_0", "update_1"]

    async def test_with_transform_hook_fluent_api(self) -> None:
        """with_transform_hook adds hook via fluent API."""
        call_count = {"value": 0}

        def counting_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
            call_count["value"] += 1
            return update

        stream = ResponseStream(_generate_updates(3), finalizer=_combine_updates).with_transform_hook(counting_hook)

        async for _ in stream:
            pass

        assert call_count["value"] == 3

    async def test_async_transform_hook(self) -> None:
        """Async transform hooks are awaited."""

        async def async_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
            return ChatResponseUpdate(
                contents=[Content.from_text(f"async_{update.text}")],
                role=cast(Any, update.role),
            )

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            transform_hooks=[async_hook],  # ty: ignore[invalid-argument-type]
        )

        collected: list[str] = []
        async for update in stream:
            collected.append(update.text or "")  # ty: ignore[unresolved-attribute]

        assert collected == ["async_update_0", "async_update_1"]


class TestResponseStreamCleanupHooks:
    """Tests for cleanup hooks (after stream consumption, before finalizer)."""

    async def test_cleanup_hook_called_after_iteration(self) -> None:
        """Cleanup hook is called after iteration completes."""
        cleanup_called = {"value": False}

        def cleanup_hook() -> None:
            cleanup_called["value"] = True

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            cleanup_hooks=[cleanup_hook],
        )

        async for _ in stream:
            pass

        assert cleanup_called["value"] is True

    async def test_cleanup_hook_called_only_once(self) -> None:
        """Cleanup hook is called only once even if get_final_response called."""
        call_count = {"value": 0}

        def cleanup_hook() -> None:
            call_count["value"] += 1

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            cleanup_hooks=[cleanup_hook],
        )

        async for _ in stream:
            pass
        await stream.get_final_response()

        assert call_count["value"] == 1

    async def test_multiple_cleanup_hooks(self) -> None:
        """Multiple cleanup hooks are called in order."""
        order: list[str] = []

        def hook_a() -> None:
            order.append("a")

        def hook_b() -> None:
            order.append("b")

        stream = ResponseStream(
            _generate_updates(1),
            finalizer=_combine_updates,
            cleanup_hooks=[hook_a, hook_b],
        )

        async for _ in stream:
            pass

        assert order == ["a", "b"]

    async def test_with_cleanup_hook_fluent_api(self) -> None:
        """with_cleanup_hook adds hook via fluent API."""
        cleanup_called = {"value": False}

        def cleanup_hook() -> None:
            cleanup_called["value"] = True

        stream = ResponseStream(_generate_updates(2), finalizer=_combine_updates).with_cleanup_hook(cleanup_hook)

        async for _ in stream:
            pass

        assert cleanup_called["value"] is True

    async def test_async_cleanup_hook(self) -> None:
        """Async cleanup hooks are awaited."""
        cleanup_called = {"value": False}

        async def async_cleanup() -> None:
            cleanup_called["value"] = True

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            cleanup_hooks=[async_cleanup],
        )

        async for _ in stream:
            pass

        assert cleanup_called["value"] is True


class TestResponseStreamResultHooks:
    """Tests for result hooks (after finalizer)."""

    async def test_result_hook_called_after_finalizer(self) -> None:
        """Result hook is called after finalizer produces result."""

        def add_metadata(response: ChatResponse) -> ChatResponse:
            response.additional_properties["processed"] = True
            return response

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            result_hooks=[add_metadata],  # ty: ignore[invalid-argument-type]
        )

        final = await stream.get_final_response()

        assert final.additional_properties["processed"] is True  # ty: ignore[unresolved-attribute]

    async def test_result_hook_can_transform_result(self) -> None:
        """Result hook can transform the final result."""

        def wrap_text(response: ChatResponse) -> ChatResponse:
            return ChatResponse(messages=Message("assistant", [f"[{response.text}]"]))

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            result_hooks=[wrap_text],  # ty: ignore[invalid-argument-type]
        )

        final = await stream.get_final_response()

        assert final.text == "[update_0update_1]"  # ty: ignore[unresolved-attribute]

    async def test_multiple_result_hooks_chained(self) -> None:
        """Multiple result hooks are called in order."""

        def add_prefix(response: ChatResponse) -> ChatResponse:
            return ChatResponse(messages=Message("assistant", [f"prefix_{response.text}"]))

        def add_suffix(response: ChatResponse) -> ChatResponse:
            return ChatResponse(messages=Message("assistant", [f"{response.text}_suffix"]))

        stream = ResponseStream(
            _generate_updates(1),
            finalizer=_combine_updates,
            result_hooks=[add_prefix, add_suffix],  # ty: ignore[invalid-argument-type]
        )

        final = await stream.get_final_response()

        assert final.text == "prefix_update_0_suffix"  # ty: ignore[unresolved-attribute]

    async def test_result_hook_returning_none_keeps_previous(self) -> None:
        """Result hook returning None keeps the previous value."""
        hook_called = {"value": False}

        def none_hook(response: ChatResponse) -> None:
            hook_called["value"] = True
            return

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            result_hooks=[none_hook],  # ty: ignore[invalid-argument-type]
        )

        final = await stream.get_final_response()

        assert hook_called["value"] is True
        assert final.text == "update_0update_1"

    async def test_with_result_hook_fluent_api(self) -> None:
        """with_result_hook adds hook via fluent API."""

        def add_metadata(response: ChatResponse) -> ChatResponse:
            response.additional_properties["via_fluent"] = True
            return response

        stream = ResponseStream(_generate_updates(2), finalizer=_combine_updates).with_result_hook(add_metadata)

        final = await stream.get_final_response()

        assert final.additional_properties["via_fluent"] is True

    async def test_async_result_hook(self) -> None:
        """Async result hooks are awaited."""

        async def async_hook(response: ChatResponse) -> ChatResponse:
            return ChatResponse(messages=Message("assistant", [f"async_{response.text}"]))

        stream = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            result_hooks=[async_hook],  # ty: ignore[invalid-argument-type]
        )

        final = await stream.get_final_response()

        assert final.text == "async_update_0update_1"  # ty: ignore[unresolved-attribute]


class TestResponseStreamFinalizer:
    """Tests for the finalizer."""

    async def test_finalizer_receives_all_updates(self) -> None:
        """Finalizer receives all collected updates."""
        received_updates: list[ChatResponseUpdate] = []

        def capturing_finalizer(updates: list[ChatResponseUpdate]) -> ChatResponse:
            received_updates.extend(updates)
            return ChatResponse(messages=Message("assistant", ["done"]))

        stream = ResponseStream(_generate_updates(3), finalizer=capturing_finalizer)  # type: ignore[arg-type, var-annotated]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

        await stream.get_final_response()

        assert len(received_updates) == 3
        assert received_updates[0].text == "update_0"
        assert received_updates[2].text == "update_2"

    async def test_no_finalizer_returns_updates(self) -> None:
        """get_final_response returns collected updates if no finalizer configured."""
        stream: ResponseStream[ChatResponseUpdate, Sequence[ChatResponseUpdate]] = ResponseStream(_generate_updates(2))

        final = await stream.get_final_response()

        assert len(final) == 2
        assert final[0].text == "update_0"
        assert final[1].text == "update_1"

    async def test_async_finalizer(self) -> None:
        """Async finalizer is awaited."""

        async def async_finalizer(updates: list[ChatResponseUpdate]) -> ChatResponse:
            text = "".join(u.text or "" for u in updates)
            return ChatResponse(messages=Message("assistant", [f"async_{text}"]))

        stream = ResponseStream(_generate_updates(2), finalizer=async_finalizer)  # type: ignore[arg-type, var-annotated]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

        final = await stream.get_final_response()

        assert final.text == "async_update_0update_1"  # ty: ignore[unresolved-attribute]

    async def test_finalized_only_once(self) -> None:
        """Finalizer is only called once even with multiple get_final_response calls."""
        call_count = {"value": 0}

        def counting_finalizer(updates: list[ChatResponseUpdate]) -> ChatResponse:
            call_count["value"] += 1
            return ChatResponse(messages=Message("assistant", ["done"]))

        stream = ResponseStream(_generate_updates(2), finalizer=counting_finalizer)  # type: ignore[arg-type, var-annotated]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]

        await stream.get_final_response()
        await stream.get_final_response()

        assert call_count["value"] == 1


class TestResponseStreamMapAndWithFinalizer:
    """Tests for ResponseStream.map() and .with_finalizer() functionality."""

    async def test_map_delegates_iteration(self) -> None:
        """Mapped stream delegates iteration to inner stream."""
        inner = ResponseStream(_generate_updates(3), finalizer=_combine_updates)

        outer = inner.map(lambda u: u, _combine_updates)

        collected: list[str] = []
        async for update in outer:
            collected.append(update.text or "")

        assert collected == ["update_0", "update_1", "update_2"]
        assert inner._consumed is True

    async def test_map_transforms_updates(self) -> None:
        """map() transforms each update."""
        inner = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        def add_prefix(update: ChatResponseUpdate) -> ChatResponseUpdate:
            return ChatResponseUpdate(
                contents=[Content.from_text(f"mapped_{update.text}")],
                role=cast(Any, update.role),
            )

        outer = inner.map(add_prefix, _combine_updates)

        collected: list[str] = []
        async for update in outer:
            collected.append(update.text or "")

        assert collected == ["mapped_update_0", "mapped_update_1"]

    async def test_map_requires_finalizer(self) -> None:
        """map() requires a finalizer since inner's won't work with new type."""
        inner = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        # map() now requires a finalizer parameter
        outer = inner.map(lambda u: u, _combine_updates)

        final = await outer.get_final_response()
        assert final.text == "update_0update_1"

    async def test_map_calls_inner_result_hooks(self) -> None:
        """map() calls inner's result hooks when get_final_response() is called."""
        inner_result_hook_called = {"value": False}

        def inner_result_hook(response: ChatResponse) -> ChatResponse:
            inner_result_hook_called["value"] = True
            return ChatResponse(messages=Message("assistant", [f"hooked_{response.text}"]))

        inner = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            result_hooks=[inner_result_hook],  # ty: ignore[invalid-argument-type]
        )
        outer = inner.map(lambda u: u, _combine_updates)

        await outer.get_final_response()

        # Inner's result_hooks ARE called when get_final_response() is invoked
        assert inner_result_hook_called["value"] is True

    async def test_with_finalizer_calls_inner_finalizer(self) -> None:
        """with_finalizer() still calls inner's finalizer first."""
        inner_finalizer_called = {"value": False}

        def inner_finalizer(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
            inner_finalizer_called["value"] = True
            return ChatResponse(messages=Message("assistant", ["inner_result"]))

        inner = ResponseStream(
            _generate_updates(2),
            finalizer=inner_finalizer,
        )
        outer = inner.with_finalizer(_combine_updates)

        final = await outer.get_final_response()

        # Inner's finalizer IS called first
        assert inner_finalizer_called["value"] is True
        # But the outer result is from outer's finalizer (working on outer's updates)
        assert final.text == "update_0update_1"

    async def test_with_finalizer_plus_result_hooks(self) -> None:
        """with_finalizer() works with result hooks."""
        inner = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        def outer_hook(response: ChatResponse) -> ChatResponse:
            return ChatResponse(messages=Message("assistant", [f"outer_{response.text}"]))

        outer = inner.with_finalizer(_combine_updates).with_result_hook(outer_hook)

        final = await outer.get_final_response()

        assert final.text == "outer_update_0update_1"

    async def test_map_with_finalizer(self) -> None:
        """map() takes a finalizer and transforms updates."""
        inner = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        def add_prefix(update: ChatResponseUpdate) -> ChatResponseUpdate:
            return ChatResponseUpdate(
                contents=[Content.from_text(f"mapped_{update.text}")],
                role=cast(Any, update.role),
            )

        outer = inner.map(add_prefix, _combine_updates)

        collected: list[str] = []
        async for update in outer:
            collected.append(update.text or "")

        assert collected == ["mapped_update_0", "mapped_update_1"]

        final = await outer.get_final_response()
        assert final.text == "mapped_update_0mapped_update_1"

    async def test_flat_map_expands_updates(self) -> None:
        """flat_map() can transform one update into many updates."""
        inner = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        def expand(update: ChatResponseUpdate) -> list[ChatResponseUpdate]:
            return [
                ChatResponseUpdate(contents=[Content.from_text(update.text)], role=cast(Any, update.role)),
                ChatResponseUpdate(contents=[Content.from_text(f"{update.text}_extra")], role=cast(Any, update.role)),
            ]

        outer = inner.flat_map(expand, _combine_updates)

        collected: list[str] = []
        async for update in outer:
            collected.append(update.text or "")

        assert collected == ["update_0", "update_0_extra", "update_1", "update_1_extra"]

        final = await outer.get_final_response()
        assert final.text == "update_0update_0_extraupdate_1update_1_extra"

    async def test_flat_map_skips_empty_mappings(self) -> None:
        """flat_map() supports zero-output transforms."""
        inner = ResponseStream(_generate_updates(3), finalizer=_combine_updates)

        def keep_odd(update: ChatResponseUpdate) -> list[ChatResponseUpdate]:
            return [update] if update.text == "update_1" else []

        outer = inner.flat_map(keep_odd, _combine_updates)

        collected = [update.text async for update in outer]
        assert collected == ["update_1"]

        final = await outer.get_final_response()
        assert final.text == "update_1"

    async def test_flat_map_calls_inner_result_hooks(self) -> None:
        """flat_map() preserves inner result hooks."""
        inner_result_hook_called = {"value": False}

        def inner_result_hook(response: ChatResponse) -> ChatResponse:
            inner_result_hook_called["value"] = True
            return response

        inner = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            result_hooks=[inner_result_hook],  # ty: ignore[invalid-argument-type]
        )
        outer = inner.flat_map(lambda u: [u], _combine_updates)

        await outer.get_final_response()

        assert inner_result_hook_called["value"] is True

    async def test_outer_transform_hooks_independent(self) -> None:
        """Outer stream has its own independent transform hooks."""
        inner_hook_calls = {"value": 0}
        outer_hook_calls = {"value": 0}

        def inner_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
            inner_hook_calls["value"] += 1
            return update

        def outer_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
            outer_hook_calls["value"] += 1
            return update

        inner = ResponseStream(
            _generate_updates(2),
            finalizer=_combine_updates,
            transform_hooks=[inner_hook],  # ty: ignore[invalid-argument-type]
        )
        outer = inner.map(lambda u: u, _combine_updates).with_transform_hook(outer_hook)

        async for _ in outer:
            pass

        assert inner_hook_calls["value"] == 2
        assert outer_hook_calls["value"] == 2

    async def test_preserves_single_consumption(self) -> None:
        """Inner stream is only consumed once."""
        consumption_count = {"value": 0}

        async def counting_generator() -> AsyncIterable[ChatResponseUpdate]:
            consumption_count["value"] += 1
            for i in range(2):
                yield ChatResponseUpdate(contents=[Content.from_text(f"u{i}")], role="assistant")

        inner = ResponseStream(counting_generator(), finalizer=_combine_updates)
        outer = inner.map(lambda u: u, _combine_updates)

        async for _ in outer:
            pass
        await outer.get_final_response()

        assert consumption_count["value"] == 1

    async def test_async_map_transform(self) -> None:
        """map() supports async transform function."""
        inner = ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        async def async_map(update: ChatResponseUpdate) -> ChatResponseUpdate:
            return ChatResponseUpdate(
                contents=[Content.from_text(f"async_{update.text}")],
                role=cast(Any, update.role),
            )

        outer = inner.map(async_map, _combine_updates)

        collected: list[str] = []
        async for update in outer:
            collected.append(update.text or "")

        assert collected == ["async_update_0", "async_update_1"]

    async def test_from_awaitable(self) -> None:
        """from_awaitable() wraps an awaitable ResponseStream."""

        async def get_stream() -> ResponseStream[ChatResponseUpdate, ChatResponse]:
            return ResponseStream(_generate_updates(2), finalizer=_combine_updates)

        outer = ResponseStream.from_awaitable(get_stream())

        collected: list[str] = []
        async for update in outer:
            collected.append(update.text or "")

        assert collected == ["update_0", "update_1"]

        final = await outer.get_final_response()
        assert final.text == "update_0update_1"


class TestResponseStreamExecutionOrder:
    """Tests verifying the correct execution order of hooks."""

    async def test_execution_order_iteration_then_finalize(self) -> None:
        """Verify execution order: transform -> cleanup -> finalizer -> result."""
        order: list[str] = []

        def transform_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
            order.append(f"transform_{update.text}")
            return update

        def cleanup_hook() -> None:
            order.append("cleanup")

        def finalizer(updates: list[ChatResponseUpdate]) -> ChatResponse:
            order.append("finalizer")
            return ChatResponse(messages=Message("assistant", ["done"]))

        def result_hook(response: ChatResponse) -> ChatResponse:
            order.append("result")
            return response

        stream = ResponseStream(  # type: ignore[var-annotated]
            _generate_updates(2),
            finalizer=finalizer,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
            transform_hooks=[transform_hook],  # ty: ignore[invalid-argument-type]
            cleanup_hooks=[cleanup_hook],
            result_hooks=[result_hook],  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
        )

        async for _ in stream:
            pass
        await stream.get_final_response()

        assert order == [
            "transform_update_0",
            "transform_update_1",
            "cleanup",
            "finalizer",
            "result",
        ]

    async def test_cleanup_runs_before_finalizer_on_direct_finalize(self) -> None:
        """Cleanup hooks run before finalizer even when not iterating manually."""
        order: list[str] = []

        def cleanup_hook() -> None:
            order.append("cleanup")

        def finalizer(updates: list[ChatResponseUpdate]) -> ChatResponse:
            order.append("finalizer")
            return ChatResponse(messages=Message("assistant", ["done"]))

        stream = ResponseStream(  # type: ignore[var-annotated]
            _generate_updates(2),
            finalizer=finalizer,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
            cleanup_hooks=[cleanup_hook],
        )

        await stream.get_final_response()

        assert order == ["cleanup", "finalizer"]


class TestResponseStreamAwaitableSource:
    """Tests for ResponseStream with awaitable stream sources."""

    async def test_awaitable_stream_source(self) -> None:
        """ResponseStream can accept an awaitable that resolves to an async iterable."""

        async def get_stream() -> AsyncIterable[ChatResponseUpdate]:
            return _generate_updates(2)

        stream = ResponseStream(get_stream(), finalizer=_combine_updates)

        collected: list[str] = []
        async for update in stream:
            collected.append(update.text or "")

        assert collected == ["update_0", "update_1"]

    async def test_await_stream(self) -> None:
        """ResponseStream can be awaited to resolve stream source."""

        async def get_stream() -> AsyncIterable[ChatResponseUpdate]:
            return _generate_updates(2)

        stream = await ResponseStream(get_stream(), finalizer=_combine_updates)

        collected: list[str] = []
        async for update in stream:
            collected.append(update.text or "")

        assert collected == ["update_0", "update_1"]


class TestResponseStreamEdgeCases:
    """Tests for edge cases and error handling."""

    async def test_empty_stream(self) -> None:
        """Empty stream produces empty result."""

        async def empty_gen() -> AsyncIterable[ChatResponseUpdate]:
            return
            yield  # type: ignore[misc]  # Make it a generator

        stream = ResponseStream(empty_gen(), finalizer=_combine_updates)

        final = await stream.get_final_response()

        assert final.text == ""
        assert len(stream.updates) == 0

    async def test_hooks_not_called_on_empty_stream_iteration(self) -> None:
        """Transform hooks not called when stream is empty."""
        hook_calls = {"value": 0}

        def transform_hook(update: ChatResponseUpdate) -> ChatResponseUpdate:
            hook_calls["value"] += 1
            return update

        async def empty_gen() -> AsyncIterable[ChatResponseUpdate]:
            return
            yield  # type: ignore[misc]

        stream = ResponseStream(
            empty_gen(),
            finalizer=_combine_updates,
            transform_hooks=[transform_hook],  # ty: ignore[invalid-argument-type]
        )

        async for _ in stream:
            pass

        assert hook_calls["value"] == 0

    async def test_cleanup_called_even_on_empty_stream(self) -> None:
        """Cleanup hooks are called even when stream is empty."""
        cleanup_called = {"value": False}

        def cleanup_hook() -> None:
            cleanup_called["value"] = True

        async def empty_gen() -> AsyncIterable[ChatResponseUpdate]:
            return
            yield  # type: ignore[misc]

        stream = ResponseStream(
            empty_gen(),
            finalizer=_combine_updates,
            cleanup_hooks=[cleanup_hook],
        )

        async for _ in stream:
            pass

        assert cleanup_called["value"] is True

    async def test_all_constructor_parameters(self) -> None:
        """All constructor parameters work together."""
        events: list[str] = []

        def transform(u: ChatResponseUpdate) -> ChatResponseUpdate:
            events.append("transform")
            return u

        def cleanup() -> None:
            events.append("cleanup")

        def finalizer(updates: list[ChatResponseUpdate]) -> ChatResponse:
            events.append("finalizer")
            return ChatResponse(messages=Message("assistant", ["done"]))

        def result(r: ChatResponse) -> ChatResponse:
            events.append("result")
            return r

        stream = ResponseStream(  # type: ignore[var-annotated]
            _generate_updates(1),
            finalizer=finalizer,  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
            transform_hooks=[transform],  # ty: ignore[invalid-argument-type]
            cleanup_hooks=[cleanup],
            result_hooks=[result],  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
        )

        await stream.get_final_response()

        assert events == ["transform", "cleanup", "finalizer", "result"]


# endregion


# region OAuth Consent Content


def test_oauth_consent_request_creation():
    """Test Content.from_oauth_consent_request creates the correct content."""
    content = Content.from_oauth_consent_request(
        consent_link="https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc",
    )
    assert content.type == "oauth_consent_request"
    assert content.consent_link == "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc"
    assert content.user_input_request is True


def test_oauth_consent_request_serialization_roundtrip():
    """Test that oauth_consent_request content serializes and includes consent_link."""
    content = Content.from_oauth_consent_request(
        consent_link="https://login.microsoftonline.com/consent",
    )
    d = content.to_dict()
    assert d["type"] == "oauth_consent_request"
    assert d["consent_link"] == "https://login.microsoftonline.com/consent"
    assert d["user_input_request"] is True


# endregion


# region prepend_instructions_to_messages tests


def test_prepend_instructions_basic():
    """Test that instructions are prepended as system message."""
    from agent_framework._types import prepend_instructions_to_messages

    messages = [Message("user", ["Hello"])]
    result = prepend_instructions_to_messages(messages, "You are helpful.")
    assert len(result) == 2
    assert result[0].role == "system"
    assert result[0].text == "You are helpful."
    assert result[1].role == "user"


def test_prepend_instructions_none():
    """Test that None instructions returns messages unchanged."""
    from agent_framework._types import prepend_instructions_to_messages

    messages = [Message("user", ["Hello"])]
    result = prepend_instructions_to_messages(messages, None)
    assert result is messages


def test_prepend_instructions_skips_duplicate():
    """Test that duplicate system instructions are not prepended again."""
    from agent_framework._types import prepend_instructions_to_messages

    messages = [
        Message("system", ["You are helpful."]),
        Message("user", ["Hello"]),
    ]
    result = prepend_instructions_to_messages(messages, "You are helpful.")
    assert len(result) == 2
    assert result[0].role == "system"
    assert result[0].text == "You are helpful."
    assert result[1].role == "user"


def test_prepend_instructions_skips_duplicate_list():
    """Test deduplication with a list of instructions."""
    from agent_framework._types import prepend_instructions_to_messages

    messages = [
        Message("system", ["First instruction"]),
        Message("system", ["Second instruction"]),
        Message("user", ["Hello"]),
    ]
    result = prepend_instructions_to_messages(messages, ["First instruction", "Second instruction"])
    assert len(result) == 3
    assert result[0].text == "First instruction"
    assert result[1].text == "Second instruction"
    assert result[2].text == "Hello"


def test_prepend_instructions_adds_when_different():
    """Test that different instructions are still prepended."""
    from agent_framework._types import prepend_instructions_to_messages

    messages = [
        Message("system", ["Old instruction"]),
        Message("user", ["Hello"]),
    ]
    result = prepend_instructions_to_messages(messages, "New instruction")
    assert len(result) == 3
    assert result[0].role == "system"
    assert result[0].text == "New instruction"
    assert result[1].text == "Old instruction"
    assert result[2].text == "Hello"


def test_prepend_instructions_custom_role():
    """Test prepending with a custom role."""
    from agent_framework._types import prepend_instructions_to_messages

    messages = [Message("user", ["Hello"])]
    result = prepend_instructions_to_messages(messages, "Be concise.", role="developer")
    assert len(result) == 2
    assert result[0].role == "developer"


# endregion


# region finish_reason


def test_agent_response_init_with_finish_reason() -> None:
    """Test that AgentResponse correctly initializes and stores finish_reason."""
    response = AgentResponse(
        messages=[Message("assistant", [Content.from_text("test")])],
        finish_reason="stop",
    )
    assert response.finish_reason == "stop"


def test_agent_response_update_init_with_finish_reason() -> None:
    """Test that AgentResponseUpdate correctly initializes and stores finish_reason."""
    update = AgentResponseUpdate(
        contents=[Content.from_text("test")],
        role="assistant",
        finish_reason="stop",
    )
    assert update.finish_reason == "stop"


def test_map_chat_to_agent_update_forwards_finish_reason() -> None:
    """Test that mapping a ChatResponseUpdate with finish_reason forwards it."""
    chat_update = ChatResponseUpdate(
        contents=[Content.from_text("test")],
        finish_reason="length",
    )
    agent_update = map_chat_to_agent_update(chat_update, agent_name="test_agent")

    assert agent_update.finish_reason == "length"
    assert agent_update.author_name == "test_agent"


def test_process_update_propagates_finish_reason_to_agent_response() -> None:
    """Test that _process_update correctly updates an AgentResponse from an AgentResponseUpdate."""
    response = AgentResponse(messages=[Message("assistant", [Content.from_text("test")])])
    update = AgentResponseUpdate(
        contents=[Content.from_text("more text")],
        role="assistant",
        finish_reason="stop",
    )

    # Process the update
    _process_update(response, update)

    assert response.finish_reason == "stop"


def test_process_update_does_not_overwrite_with_none() -> None:
    """Test that _process_update does not overwrite an existing finish_reason with None."""
    response = AgentResponse(
        messages=[Message("assistant", [Content.from_text("test")])],
        finish_reason="length",
    )
    update = AgentResponseUpdate(
        contents=[Content.from_text("more text")],
        role="assistant",
        finish_reason=None,
    )

    # Process the update
    _process_update(response, update)

    assert response.finish_reason == "length"


def test_agent_response_serialization_includes_finish_reason() -> None:
    """Test that AgentResponse serializes correctly, including finish_reason."""
    response = AgentResponse(
        messages=[Message("assistant", [Content.from_text("test")])],
        response_id="test_123",
        finish_reason="stop",
    )

    # Serialize using the framework's API and verify finish_reason is included.
    data = response.to_dict()
    assert "finish_reason" in data
    assert data["finish_reason"] == "stop"


def test_agent_response_update_serialization_includes_finish_reason() -> None:
    """Test that AgentResponseUpdate serializes correctly, including finish_reason."""
    update = AgentResponseUpdate(
        contents=[Content.from_text("test")],
        role="assistant",
        response_id="test_456",
        finish_reason="tool_calls",
    )

    data = update.to_dict()
    assert "finish_reason" in data
    assert data["finish_reason"] == "tool_calls"


# endregion
