from __future__ import annotations

from collections.abc import AsyncIterator

import pytest

from terminaluse.lib.core.services.adk.messages import MessagesService
from terminaluse.lib.types.message_parts import DataPart, TextPart


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 _FailingRawEvents:
    async def publish_raw_event(self, **_kwargs):
        raise RuntimeError("publish failed")


class _FailingTerminalUseClient:
    def __init__(self):
        self.raw_events = _FailingRawEvents()


class _FakeUiMessage:
    def __init__(self, message_id: str):
        self.id = message_id

    def model_dump(self):
        return {"id": self.id}


class _FakePager:
    def __init__(self, pages: list[list[_FakeUiMessage]]):
        self._pages = pages
        self.items = pages[0] if pages else []

    async def __aiter__(self) -> AsyncIterator[_FakeUiMessage]:
        for page in self._pages:
            for item in page:
                yield item


class _ListMessagesV2Client:
    def __init__(self, pager: _FakePager):
        self._pager = pager

    async def list(self, **_kwargs):
        return self._pager


class _ListTerminalUseClient:
    def __init__(self, pager: _FakePager):
        self.messages_v2 = _ListMessagesV2Client(pager)


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

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


class _CapturingTerminalUseClient:
    def __init__(self, raw_events: _CapturingRawEvents):
        self.raw_events = raw_events


def _build_service() -> MessagesService:
    return MessagesService(
        terminaluse_client=_FailingTerminalUseClient(),  # type: ignore[arg-type]
        tracer=_NoopTracer(),  # type: ignore[arg-type]
    )


def test_content_to_parts_text_part():
    """Test that TextPart is converted correctly."""
    service = _build_service()
    content = TextPart(text="hello world")

    parts = service._content_to_parts(content)

    assert parts == [{"type": "text", "text": "hello world"}]


def test_content_to_parts_data_part():
    """Test that DataPart is converted correctly."""
    service = _build_service()
    content = DataPart(data={"key": "value", "count": 42})

    parts = service._content_to_parts(content)

    assert parts == [{"type": "data", "data": {"key": "value", "count": 42}}]


@pytest.mark.asyncio
async def test_send_raises_when_raw_event_publish_fails():
    service = _build_service()
    content = TextPart(text="hello")

    with pytest.raises(RuntimeError, match="publish failed"):
        await service.send(task_id="task_1", content=content)


@pytest.mark.asyncio
async def test_list_messages_with_limit_returns_current_page_only():
    pager = _FakePager(
        pages=[
            [_FakeUiMessage("m1")],
            [_FakeUiMessage("m2")],
        ]
    )
    service = MessagesService(
        terminaluse_client=_ListTerminalUseClient(pager),  # type: ignore[arg-type]
        tracer=_NoopTracer(),  # type: ignore[arg-type]
    )

    messages = await service.list_messages(task_id="task_1", limit=1)

    assert [msg.id for msg in messages] == ["m1"]


@pytest.mark.asyncio
async def test_list_messages_without_limit_exhausts_all_pages():
    pager = _FakePager(
        pages=[
            [_FakeUiMessage("m1")],
            [_FakeUiMessage("m2"), _FakeUiMessage("m3")],
        ]
    )
    service = MessagesService(
        terminaluse_client=_ListTerminalUseClient(pager),  # type: ignore[arg-type]
        tracer=_NoopTracer(),  # type: ignore[arg-type]
    )

    messages = await service.list_messages(task_id="task_1")

    assert [msg.id for msg in messages] == ["m1", "m2", "m3"]


@pytest.mark.asyncio
async def test_publish_ui_message_sets_idempotency_key_argument():
    raw_events = _CapturingRawEvents()
    service = MessagesService(
        terminaluse_client=_CapturingTerminalUseClient(raw_events),  # type: ignore[arg-type]
        tracer=_NoopTracer(),  # type: ignore[arg-type]
    )

    await service._publish_ui_message(
        task_id="task_1",
        role="user",
        parts=[{"type": "text", "text": "hello"}],
    )

    assert len(raw_events.calls) == 1
    call = raw_events.calls[0]
    assert isinstance(call["idempotency_key"], str)
    assert call["idempotency_key"] != ""
