"""Tests for the distributed tracer."""

from __future__ import annotations

import pytest

from a2a_mesh.tracer import MeshTracer


class TestMeshTracer:
    """Tests for MeshTracer."""

    @pytest.mark.asyncio
    async def test_trace_task(self, tracer: MeshTracer) -> None:
        async with tracer.trace_task("dispatch", agent_name="research") as span:
            span.cost = 0.05

        assert len(tracer.spans) == 1
        recorded = tracer.spans[0]
        assert recorded.operation == "dispatch"
        assert recorded.agent_name == "research"
        assert recorded.cost == 0.05
        assert recorded.duration_ms > 0
        assert recorded.status == "ok"

    @pytest.mark.asyncio
    async def test_trace_task_error(self, tracer: MeshTracer) -> None:
        with pytest.raises(ValueError, match="boom"):
            async with tracer.trace_task("failing") as span:
                raise ValueError("boom")

        assert len(tracer.spans) == 1
        assert tracer.spans[0].status == "error"

    @pytest.mark.asyncio
    async def test_get_traces(self, tracer: MeshTracer) -> None:
        for i in range(5):
            async with tracer.trace_task(f"op-{i}", agent_name=f"agent-{i}"):
                pass

        traces = tracer.get_traces(limit=3)
        assert len(traces) == 3
        # Most recent first
        assert traces[0].operation == "op-4"

    @pytest.mark.asyncio
    async def test_total_cost(self, tracer: MeshTracer) -> None:
        async with tracer.trace_task("a") as span:
            span.cost = 0.10
        async with tracer.trace_task("b") as span:
            span.cost = 0.05

        assert abs(tracer.total_cost() - 0.15) < 1e-10

    @pytest.mark.asyncio
    async def test_trace_ids(self, tracer: MeshTracer) -> None:
        async with tracer.trace_task("op") as span:
            pass

        assert span.trace_id
        assert span.span_id
        assert len(span.trace_id) == 32

    @pytest.mark.asyncio
    async def test_get_trace_by_id(self, tracer: MeshTracer) -> None:
        async with tracer.trace_task("first") as span1:
            pass
        trace_id = span1.trace_id

        async with tracer.trace_task("other"):
            pass

        matching = tracer.get_trace_by_id(trace_id)
        assert len(matching) == 1
        assert matching[0].operation == "first"

    @pytest.mark.asyncio
    async def test_span_attributes(self, tracer: MeshTracer) -> None:
        async with tracer.trace_task(
            "custom",
            attributes={"model": "gpt-4", "tokens": 150},
        ) as span:
            pass

        assert tracer.spans[0].attributes["model"] == "gpt-4"

    def test_shutdown(self, tracer: MeshTracer) -> None:
        tracer.shutdown()
        # Should not raise
