from __future__ import annotations

from datetime import datetime, timezone
from uuid import uuid4

from v1.agent.utils import convert_message_to_history


class _FakeMessage:
    def __init__(self, *, role: str, metadata: dict[str, object] | None) -> None:
        self.id = uuid4()
        self.seq = 1
        self.role = role
        self.content = "content"
        self.metadata = metadata
        self.timestamp = datetime.now(timezone.utc)


def test_convert_message_to_history_uses_ui_schema_key_for_tool_message() -> None:
    message = _FakeMessage(
        role="tool",
        metadata={
            "tool_agent_output": {
                "ui_schema": {"version": "2.0", "root": {"type": "stack"}}
            }
        },
    )

    result = convert_message_to_history(message)  # type: ignore[arg-type]

    assert "ui_schema" in result
    assert "uiSchema" not in result
    assert result["ui_schema"] == {"version": "2.0", "root": {"type": "stack"}}


def test_convert_message_to_history_uses_ui_schema_key_for_assistant_message() -> None:
    message = _FakeMessage(
        role="assistant",
        metadata={
            "worker_agent_output": {
                "ui_schema": {"version": "2.0", "root": {"type": "stack"}}
            }
        },
    )

    result = convert_message_to_history(message)  # type: ignore[arg-type]

    assert "ui_schema" in result
    assert "uiSchema" not in result
    assert result["ui_schema"] == {"version": "2.0", "root": {"type": "stack"}}