"""Coalex execution context — sets coalex.* attributes on a parent span."""

from __future__ import annotations

import contextlib
from collections.abc import Generator
from contextvars import ContextVar

from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode

_coalex_agent_id: ContextVar[str | None] = ContextVar("coalex_agent_id", default=None)
_coalex_request_id: ContextVar[str | None] = ContextVar("coalex_request_id", default=None)
_coalex_agent_version: ContextVar[str | None] = ContextVar("coalex_agent_version", default=None)


@contextlib.contextmanager
def coalex_context(
    *,
    agent_id: str,
    request_id: str | None = None,
    version: str | None = None,
) -> Generator[None, None, None]:
    """Create a parent span with Coalex attributes for all child spans.

    All spans created inside this context are descendants of a
    ``coalex.invocation`` span that carries the agent metadata.
    This makes every child span queryable by agent_id, request_id,
    and version in the Bronze layer.

    Context vars are set before the span is created so that
    ``CoalexAttributePropagator`` can propagate attributes even
    across trace boundaries (e.g. when LangChain auto-instrumentors
    create new root spans with separate trace IDs).

    Args:
        agent_id: Identifier for the AI agent.
        request_id: Unique request/invocation ID.
        version: Agent version string.
    """
    token_a = _coalex_agent_id.set(agent_id)
    token_r = _coalex_request_id.set(request_id)
    token_v = _coalex_agent_version.set(version)

    try:
        tracer = trace.get_tracer("coalex")

        attrs: dict[str, str] = {"coalex.agent_id": agent_id}
        if request_id is not None:
            attrs["coalex.request_id"] = request_id
        if version is not None:
            attrs["coalex.agent_version"] = version

        with tracer.start_as_current_span("coalex.invocation", attributes=attrs) as span:
            yield
            span.set_status(Status(StatusCode.OK))
    finally:
        _coalex_agent_id.reset(token_a)
        _coalex_request_id.reset(token_r)
        _coalex_agent_version.reset(token_v)
