"""OpenInference compatibility span processor."""

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

from plyra_trace.semconv import SpanAttributes


class OpenInferenceSpanProcessor(SpanProcessor):
    """
    Span processor that normalizes OpenInference spans to plyra format.

    OpenInference and plyra-trace use the same semantic conventions
    for most attributes, so this processor mainly:
    - Maps openinference.span.kind → plyra.span.kind
    - Ensures both attributes are set for cross-platform compatibility

    OpenInference spec:
    https://arize-ai.github.io/openinference/spec/semantic_conventions.html
    """

    def on_start(self, span: Span, parent_context: otel_context.Context | None = None) -> None:
        """Called when a span is started."""
        # Check if this span has openinference.span.kind set
        # If so, copy it to plyra.span.kind
        # This allows spans created by OpenInference instrumentation
        # to be recognized by plyra-trace
        pass  # Attributes are set after span creation, handled in on_end

    def on_end(self, span: ReadableSpan) -> None:
        """Called when a span ends.

        Normalize OpenInference attributes to plyra format.
        """
        if not span.attributes:
            return

        attrs = dict(span.attributes)

        # Map openinference.span.kind to plyra.span.kind if present
        oi_span_kind = attrs.get(SpanAttributes.OPENINFERENCE_SPAN_KIND)
        if oi_span_kind and SpanAttributes.SPAN_KIND not in attrs:
            # The span already has openinference.span.kind but not plyra.span.kind
            # This means it was created by OpenInference instrumentation
            # We can't modify the span after it's ended, but we can note this
            # for future implementations that might use a different approach
            pass

        # OpenInference attributes are already compatible with plyra
        # No mapping needed for:
        # - input.value, output.value, input.mime_type, output.mime_type
        # - llm.*, tool.*, retrieval.*, embedding.*
        # - session.id, user.id, metadata, tag.tags

    def shutdown(self) -> None:
        """Called when the span processor is shutdown."""
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush spans."""
        return True
