"""Tests for AgentMessagePart types and ctx.messages.send() integration.

These tests verify:
1. TextPart and DataPart types work correctly
2. ctx.messages.send() accepts TextPart, DataPart, and str
3. Messages are published in the correct UIMessage format
"""

from __future__ import annotations

import pytest

from terminaluse.lib.types.message_parts import DataPart, TextPart


class TestMessagePartTypes:
    """Tests for TextPart and DataPart types."""

    def test_text_part_has_type_discriminator(self):
        """TextPart should have type='text' discriminator."""
        part = TextPart(text="hello")
        assert part.type == "text"
        assert part.text == "hello"

    def test_text_part_serializes_correctly(self):
        """TextPart should serialize to UIMessage format."""
        part = TextPart(text="hello world")
        data = part.model_dump()
        assert data == {"type": "text", "text": "hello world"}

    def test_data_part_has_type_discriminator(self):
        """DataPart should have type='data' discriminator."""
        part = DataPart(data={"key": "value"})
        assert part.type == "data"
        assert part.data == {"key": "value"}

    def test_data_part_serializes_correctly(self):
        """DataPart should serialize to UIMessage format."""
        part = DataPart(data={"status": "complete", "count": 42})
        data = part.model_dump()
        assert data == {"type": "data", "data": {"status": "complete", "count": 42}}

    def test_data_part_accepts_any_data(self):
        """DataPart should accept any JSON-serializable data."""
        # Dict
        assert DataPart(data={"a": 1}).data == {"a": 1}
        # List
        assert DataPart(data=[1, 2, 3]).data == [1, 2, 3]
        # String
        assert DataPart(data="raw string").data == "raw string"
        # Number
        assert DataPart(data=42).data == 42
        # Nested
        assert DataPart(data={"nested": {"deep": [1, 2]}}).data == {"nested": {"deep": [1, 2]}}


class _CapturingRawEvents:
    """Captures raw_events.publish_raw_event calls for testing."""

    def __init__(self):
        self.calls: list[dict] = []

    async def publish_raw_event(self, **kwargs):
        self.calls.append(kwargs)


class _CapturingTerminalUseClient:
    """Fake TerminalUse client that captures raw event calls."""

    def __init__(self):
        self.raw_events = _CapturingRawEvents()


class _NoopSpan:
    async def __aenter__(self):
        return None

    async def __aexit__(self, exc_type, exc, tb):
        return False


class _NoopTrace:
    def span(self, **kwargs):
        return _NoopSpan()


class _NoopTracer:
    def trace(self, *_args, **_kwargs):
        return _NoopTrace()


class TestMessagesServiceWithMessageParts:
    """Tests for MessagesService handling TextPart and DataPart."""

    def setup_method(self):
        """Set up test fixtures."""
        from terminaluse.lib.core.services.adk.messages import MessagesService

        self.client = _CapturingTerminalUseClient()
        self.service = MessagesService(
            terminaluse_client=self.client,  # type: ignore[arg-type]
            tracer=_NoopTracer(),  # type: ignore[arg-type]
        )

    @pytest.mark.asyncio
    async def test_send_text_part_publishes_ui_message(self):
        """MessagesService.send() should accept TextPart and publish UIMessage."""
        part = TextPart(text="hello from agent")

        await self.service.send(task_id="task_123", content=part)

        assert len(self.client.raw_events.calls) == 1
        call = self.client.raw_events.calls[0]
        assert call["task_id"] == "task_123"
        assert call["sdk_type"] == "terminal_use"
        # The raw_event should be a UIMessage
        raw_event = call["raw_event"]
        assert raw_event["role"] == "assistant"
        assert len(raw_event["parts"]) == 1
        assert raw_event["parts"][0] == {"type": "text", "text": "hello from agent"}

    @pytest.mark.asyncio
    async def test_send_data_part_publishes_ui_message(self):
        """MessagesService.send() should accept DataPart and publish UIMessage."""
        part = DataPart(data={"status": "complete", "items": [1, 2, 3]})

        await self.service.send(task_id="task_456", content=part)

        assert len(self.client.raw_events.calls) == 1
        call = self.client.raw_events.calls[0]
        raw_event = call["raw_event"]
        assert raw_event["role"] == "assistant"
        assert raw_event["parts"][0] == {
            "type": "data",
            "data": {"status": "complete", "items": [1, 2, 3]},
        }


class TestMessagesModuleWithAgentMessagePart:
    """Tests for MessagesModule handling TextPart, DataPart, and str.

    These tests verify that MessagesModule.send() correctly routes
    AgentMessagePart types to MessagesService, bypassing CreateMessageParams.
    """

    @pytest.mark.asyncio
    async def test_messages_module_send_text_part_routes_to_service(self):
        """MessagesModule.send() with TextPart should route to MessagesService."""
        from unittest.mock import AsyncMock, MagicMock

        from terminaluse.lib.adk._modules.messages import MessagesModule

        # Create module with mocked service
        mock_service = MagicMock()
        mock_service.send = AsyncMock()

        module = MessagesModule(messages_service=mock_service)

        part = TextPart(text="hello")
        await module.send(task_id="test_task", content=part)

        # Verify it called MessagesService.send() with the TextPart
        mock_service.send.assert_called_once()
        call_kwargs = mock_service.send.call_args.kwargs
        assert call_kwargs["task_id"] == "test_task"
        assert call_kwargs["content"] == part

    @pytest.mark.asyncio
    async def test_messages_module_send_data_part_routes_to_service(self):
        """MessagesModule.send() with DataPart should route to MessagesService."""
        from unittest.mock import AsyncMock, MagicMock

        from terminaluse.lib.adk._modules.messages import MessagesModule

        mock_service = MagicMock()
        mock_service.send = AsyncMock()

        module = MessagesModule(messages_service=mock_service)

        part = DataPart(data={"key": "value"})
        await module.send(task_id="test_task", content=part)

        mock_service.send.assert_called_once()
        call_kwargs = mock_service.send.call_args.kwargs
        assert call_kwargs["task_id"] == "test_task"
        assert call_kwargs["content"] == part

    @pytest.mark.asyncio
    async def test_messages_module_send_string_wraps_to_text_part(self):
        """MessagesModule.send() with string should wrap to TextPart."""
        from unittest.mock import AsyncMock, MagicMock

        from terminaluse.lib.adk._modules.messages import MessagesModule

        mock_service = MagicMock()
        mock_service.send = AsyncMock()

        module = MessagesModule(messages_service=mock_service)

        await module.send(task_id="test_task", content="hello world")

        mock_service.send.assert_called_once()
        call_kwargs = mock_service.send.call_args.kwargs
        assert call_kwargs["task_id"] == "test_task"
        # Should be wrapped in TextPart
        assert isinstance(call_kwargs["content"], TextPart)
        assert call_kwargs["content"].text == "hello world"

    @pytest.mark.asyncio
    async def test_messages_module_send_codex_event_publishes_raw_event(self):
        """MessagesModule.send() with ThreadEvent should publish codex raw event."""
        from unittest.mock import AsyncMock, MagicMock

        codex_events = pytest.importorskip("agents.extensions.experimental.codex.events")
        from terminaluse.lib.adk._modules.messages import MessagesModule

        client = _CapturingTerminalUseClient()
        mock_service = MagicMock()
        mock_service.send = AsyncMock()
        mock_service._terminaluse_client = client

        module = MessagesModule(messages_service=mock_service)
        codex_event = codex_events.ThreadStartedEvent(thread_id="thread_123")

        await module.send(task_id="test_task", content=codex_event, sdk_version="0.9.0")

        assert len(client.raw_events.calls) == 1
        call = client.raw_events.calls[0]
        assert call["task_id"] == "test_task"
        assert call["sdk_type"] == "codex"
        assert call["sdk_version"] == "0.9.0"
        assert call["raw_event"] == {"type": "thread.started", "thread_id": "thread_123"}
        mock_service.send.assert_not_called()

    @pytest.mark.asyncio
    async def test_messages_module_send_codex_event_object_publishes_raw_event(self):
        """MessagesModule.send() should accept non-thread-started ThreadEvent payloads."""
        from unittest.mock import AsyncMock, MagicMock

        codex_events = pytest.importorskip("agents.extensions.experimental.codex.events")
        from terminaluse.lib.adk._modules.messages import MessagesModule

        client = _CapturingTerminalUseClient()
        mock_service = MagicMock()
        mock_service.send = AsyncMock()
        mock_service._terminaluse_client = client

        module = MessagesModule(messages_service=mock_service)
        codex_event = codex_events.TurnFailedEvent(error=codex_events.ThreadError(message="quota exceeded"))

        await module.send(task_id="test_task", content=codex_event, sdk_version="0.9.0")

        assert len(client.raw_events.calls) == 1
        call = client.raw_events.calls[0]
        assert call["task_id"] == "test_task"
        assert call["sdk_type"] == "codex"
        assert call["sdk_version"] == "0.9.0"
        assert call["raw_event"] == {"type": "turn.failed", "error": {"message": "quota exceeded"}}
        mock_service.send.assert_not_called()

    @pytest.mark.asyncio
    async def test_messages_module_send_plain_dict_raises_type_error(self):
        """MessagesModule.send() should reject plain dict content."""
        from unittest.mock import AsyncMock, MagicMock

        from terminaluse.lib.adk._modules.messages import MessagesModule

        mock_service = MagicMock()
        mock_service.send = AsyncMock()
        mock_service._terminaluse_client = _CapturingTerminalUseClient()

        module = MessagesModule(messages_service=mock_service)

        with pytest.raises(TypeError, match="Unsupported content type"):
            await module.send(task_id="test_task", content={"foo": "bar"})

        mock_service.send.assert_not_called()


class TestContextMessagesModuleWithMessageParts:
    """Tests for ContextMessagesModule handling TextPart, DataPart, and str."""

    @pytest.mark.asyncio
    async def test_send_string_wraps_to_text_part(self, monkeypatch):
        """ctx.messages.send('hello') should wrap string in TextPart."""
        from terminaluse.lib.types.task_context import ContextMessagesModule

        captured_content = []

        async def fake_send(task_id: str, content, **kwargs):
            captured_content.append(content)

        # Mock at the terminaluse.lib module level
        import terminaluse.lib

        class FakeMessages:
            send = staticmethod(fake_send)

        class FakeAdk:
            messages = FakeMessages()

        monkeypatch.setattr(terminaluse.lib, "adk", FakeAdk())

        module = ContextMessagesModule(task_id="task_789")
        await module.send("hello world")

        assert len(captured_content) == 1
        content = captured_content[0]
        # After refactor, should be TextPart
        assert isinstance(content, TextPart)
        assert content.text == "hello world"
        assert content.type == "text"

    @pytest.mark.asyncio
    async def test_send_text_part_passes_through(self, monkeypatch):
        """ctx.messages.send(TextPart(...)) should pass through unchanged."""
        from terminaluse.lib.types.task_context import ContextMessagesModule

        captured_content = []

        async def fake_send(task_id: str, content, **kwargs):
            captured_content.append(content)

        import terminaluse.lib

        class FakeMessages:
            send = staticmethod(fake_send)

        class FakeAdk:
            messages = FakeMessages()

        monkeypatch.setattr(terminaluse.lib, "adk", FakeAdk())

        module = ContextMessagesModule(task_id="task_789")
        part = TextPart(text="explicit text part")
        await module.send(part)

        assert len(captured_content) == 1
        assert captured_content[0] is part

    @pytest.mark.asyncio
    async def test_send_data_part_passes_through(self, monkeypatch):
        """ctx.messages.send(DataPart(...)) should pass through unchanged."""
        from terminaluse.lib.types.task_context import ContextMessagesModule

        captured_content = []

        async def fake_send(task_id: str, content, **kwargs):
            captured_content.append(content)

        import terminaluse.lib

        class FakeMessages:
            send = staticmethod(fake_send)

        class FakeAdk:
            messages = FakeMessages()

        monkeypatch.setattr(terminaluse.lib, "adk", FakeAdk())

        module = ContextMessagesModule(task_id="task_789")
        part = DataPart(data={"key": "value"})
        await module.send(part)

        assert len(captured_content) == 1
        assert captured_content[0] is part
