"""Span processor to propagate Coalex attributes from parent to child spans."""

import logging

from opentelemetry import context, trace
from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor

from coalex.context import _coalex_agent_id, _coalex_agent_version, _coalex_request_id

logger = logging.getLogger(__name__)

_CONTEXT_VAR_MAP = {
    "coalex.agent_id": _coalex_agent_id,
    "coalex.request_id": _coalex_request_id,
    "coalex.agent_version": _coalex_agent_version,
}


class CoalexAttributePropagator(SpanProcessor):
    """Propagates Coalex attributes from parent to child spans.

    Ensures that coalex.* attributes set on the coalex.invocation parent span
    are automatically copied to all child spans (like auto-instrumented LLM calls).

    When a parent span doesn't carry coalex attributes (e.g. a new root span
    created by LangChain's auto-instrumentor), falls back to reading from
    ``contextvars`` set by ``coalex_context()``.
    """

    ATTRIBUTES_TO_PROPAGATE = [
        "coalex.agent_id",
        "coalex.request_id",
        "coalex.agent_version",
    ]

    def on_start(self, span: Span, parent_context: context.Context | None = None) -> None:
        if not span.is_recording():
            return

        ctx = parent_context if parent_context is not None else context.get_current()
        parent_span = trace.get_current_span(ctx)

        for attr_name in self.ATTRIBUTES_TO_PROPAGATE:
            # 1) Try parent span first (fast path for same-trace children)
            parent_value = None
            if parent_span and hasattr(parent_span, "attributes"):
                parent_value = getattr(parent_span, "attributes", {}).get(attr_name)

            # 2) Fallback: read from context vars (handles cross-trace boundaries)
            if parent_value is None:
                context_var = _CONTEXT_VAR_MAP.get(attr_name)
                if context_var is not None:
                    parent_value = context_var.get(None)

            # 3) Set attribute if we found a value and span doesn't already have it
            if parent_value is not None:
                span_attrs = getattr(span, "attributes", {})
                if attr_name not in span_attrs:
                    span.set_attribute(attr_name, parent_value)

    def on_end(self, span: ReadableSpan) -> None:
        pass

    def shutdown(self) -> None:
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return True
