from __future__ import annotations

import pytest

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


class _FailingRawEvents:
    async def publish_raw_event(self, **_kwargs):
        raise RuntimeError("publish failed")


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


class _DummyMessagesService:
    def __init__(self):
        self._terminaluse_client = _FailingTerminalUseClient()


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


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


@pytest.mark.asyncio
async def test_publish_claude_event_raises_when_publish_fails():
    module = MessagesModule(messages_service=_DummyMessagesService())  # type: ignore[arg-type]

    with pytest.raises(RuntimeError, match="publish failed"):
        await module._publish_claude_event(
            task_id="task_1",
            raw_event={"type": "UserMessage", "data": {}},
        )


@pytest.mark.asyncio
async def test_publish_claude_event_sets_idempotency_key_argument():
    raw_events = _CapturingRawEvents()
    module = MessagesModule(
        messages_service=_CapturingMessagesService(raw_events)  # type: ignore[arg-type]
    )

    await module._publish_claude_event(
        task_id="task_1",
        raw_event={"type": "UserMessage", "data": {}},
    )

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


@pytest.mark.asyncio
async def test_publish_claude_event_forwards_request_headers_and_delivery_event_id():
    raw_events = _CapturingRawEvents()
    module = MessagesModule(
        messages_service=_CapturingMessagesService(raw_events)  # type: ignore[arg-type]
    )

    await module._publish_claude_event(
        task_id="task_1",
        raw_event={"type": "UserMessage", "data": {}},
        request_headers={
            "x-request-id": "req-abc",
            "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
            "tracestate": "rojo=00f067aa0ba902b7",
            "baggage": "team=platform",
            "cookie": "blocked",
        },
        delivery_event_id="evt-abc",
    )

    assert len(raw_events.calls) == 1
    call = raw_events.calls[0]
    request_options = call.get("request_options")
    assert request_options is not None
    additional_headers = request_options.get("additional_headers")
    assert additional_headers["x-request-id"] == "req-abc"
    assert additional_headers["traceparent"] == "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
    assert additional_headers["tracestate"] == "rojo=00f067aa0ba902b7"
    assert additional_headers["baggage"] == "team=platform"
    assert additional_headers["x-terminaluse-delivery-event-id"] == "evt-abc"
    assert "cookie" not in additional_headers
