"""Tests for coalex.coalex_context()."""

import pytest
from coalex import coalex_context
from opentelemetry import trace
from opentelemetry.trace import StatusCode


def test_context_creates_invocation_span(memory_exporter):
    """coalex_context() should create a coalex.invocation parent span."""
    with coalex_context(agent_id="test-agent"):
        pass

    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 1
    assert spans[0].name == "coalex.invocation"


def test_context_sets_agent_id(memory_exporter):
    """coalex_context() should set coalex.agent_id attribute."""
    with coalex_context(agent_id="claims-processor"):
        pass

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["coalex.agent_id"] == "claims-processor"


def test_context_sets_optional_attrs(memory_exporter):
    """coalex_context() should set request_id and version when provided."""
    with coalex_context(agent_id="my-agent", request_id="req_001", version="1.0.0"):
        pass

    spans = memory_exporter.get_finished_spans()
    attrs = spans[0].attributes
    assert attrs["coalex.agent_id"] == "my-agent"
    assert attrs["coalex.request_id"] == "req_001"
    assert attrs["coalex.agent_version"] == "1.0.0"


def test_context_omits_none_attrs(memory_exporter):
    """coalex_context() should not set attributes for None values."""
    with coalex_context(agent_id="my-agent"):
        pass

    spans = memory_exporter.get_finished_spans()
    attrs = spans[0].attributes
    assert "coalex.request_id" not in attrs
    assert "coalex.agent_version" not in attrs


def test_child_spans_are_descendants(memory_exporter):
    """Spans created inside coalex_context() should be children of the invocation span."""
    tracer = trace.get_tracer("test")

    with coalex_context(agent_id="my-agent"), tracer.start_as_current_span("child-span"):
        pass

    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 2

    child = next(s for s in spans if s.name == "child-span")
    parent = next(s for s in spans if s.name == "coalex.invocation")

    assert child.parent is not None
    assert child.parent.span_id == parent.context.span_id


def test_context_status_ok_on_success(memory_exporter):
    """coalex_context() should set status OK when no exception is raised."""
    with coalex_context(agent_id="test-agent"):
        pass

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.OK


def test_context_status_error_on_exception(memory_exporter):
    """coalex_context() should set status ERROR and record exception on failure."""
    with pytest.raises(ValueError, match="boom"), coalex_context(agent_id="test-agent"):
        raise ValueError("boom")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].status.status_code == StatusCode.ERROR
    assert "boom" in spans[0].status.description
    exception_events = [e for e in spans[0].events if e.name == "exception"]
    assert len(exception_events) == 1


def test_cross_trace_propagation(memory_exporter):
    """Spans with new trace IDs still get coalex attributes via context vars."""
    tracer = trace.get_tracer("test")

    with (
        coalex_context(agent_id="test-agent", request_id="req_123"),
        tracer.start_span("langchain.RunnableSequence") as new_root,
    ):
        assert new_root.attributes.get("coalex.agent_id") == "test-agent"
        assert new_root.attributes.get("coalex.request_id") == "req_123"


def test_context_vars_cleaned_up(memory_exporter):
    """Context vars are reset after coalex_context exits."""
    from coalex.context import _coalex_agent_id

    assert _coalex_agent_id.get(None) is None

    with coalex_context(agent_id="test-agent"):
        assert _coalex_agent_id.get(None) == "test-agent"

    assert _coalex_agent_id.get(None) is None


def test_nested_context_override(memory_exporter):
    """Nested coalex_context overrides outer values, restores on exit."""
    from coalex.context import _coalex_agent_id

    with coalex_context(agent_id="outer-agent"):
        assert _coalex_agent_id.get(None) == "outer-agent"

        with coalex_context(agent_id="inner-agent"):
            assert _coalex_agent_id.get(None) == "inner-agent"

        assert _coalex_agent_id.get(None) == "outer-agent"
