"""OpenLLMetry (Traceloop) compatibility span processor."""

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

from plyra_trace.semconv import SpanAttributes, SpanKind


class OpenLLMetrySpanProcessor(SpanProcessor):
    """
    Span processor that normalizes OpenLLMetry (Traceloop) spans.

    OpenLLMetry (from Traceloop) uses different attribute names
    than plyra/OpenInference.
    This processor maps:
    - traceloop.span.kind → plyra.span.kind
    - traceloop.entity.name → agent.name or tool.name
    - gen_ai.* → llm.*
    - llm.request.model → llm.model_name
    - llm.usage.* → llm.token_count.*

    Traceloop semantic conventions:
    https://www.traceloop.com/docs/openllmetry/tracing/semantic-conventions
    """

    # Mapping from Traceloop span kinds to plyra span kinds
    SPAN_KIND_MAP = {
        "workflow": SpanKind.CHAIN,
        "task": SpanKind.TOOL,
        "agent": SpanKind.AGENT,
        "tool": SpanKind.TOOL,
    }

    def on_start(self, span: Span, parent_context: otel_context.Context | None = None) -> None:
        """Called when a span is started."""
        pass

    def on_end(self, span: ReadableSpan) -> None:
        """Called when a span ends. Normalize Traceloop attributes to plyra format."""
        if not span.attributes:
            return

        attrs = dict(span.attributes)

        # Map traceloop.span.kind to plyra.span.kind
        traceloop_kind = attrs.get("traceloop.span.kind")
        if traceloop_kind:
            plyra_kind = self.SPAN_KIND_MAP.get(traceloop_kind)
            if plyra_kind and SpanAttributes.SPAN_KIND not in attrs:
                # Note: We can't modify the span after it's ended
                # This processor documents the mapping for future implementations
                # that might use mutable span wrappers or exporters
                pass

        # Map traceloop.entity.name to appropriate plyra attribute
        entity_name = attrs.get("traceloop.entity.name")
        if entity_name:
            if traceloop_kind == "agent" and SpanAttributes.AGENT_NAME not in attrs:
                pass  # Would map to agent.name
            elif traceloop_kind in ("tool", "task") and SpanAttributes.TOOL_NAME not in attrs:
                pass  # Would map to tool.name

        # Map gen_ai.prompt.* to llm.input_messages.*
        # Map gen_ai.completion.* to llm.output_messages.*

        # Map llm.request.model to llm.model_name
        if "llm.request.model" in attrs and SpanAttributes.LLM_MODEL_NAME not in attrs:
            pass  # Would map llm.request.model → llm.model_name

        # Map llm.usage.prompt_tokens to llm.token_count.prompt
        if "llm.usage.prompt_tokens" in attrs and SpanAttributes.LLM_TOKEN_COUNT_PROMPT not in attrs:
            pass  # Would map llm.usage.prompt_tokens → llm.token_count.prompt

        # Map llm.usage.completion_tokens to llm.token_count.completion
        if "llm.usage.completion_tokens" in attrs and SpanAttributes.LLM_TOKEN_COUNT_COMPLETION not in attrs:
            pass  # Would map llm.usage.completion_tokens → llm.token_count.completion

        # Map llm.usage.total_tokens to llm.token_count.total
        if "llm.usage.total_tokens" in attrs and SpanAttributes.LLM_TOKEN_COUNT_TOTAL not in attrs:
            pass  # Would map llm.usage.total_tokens → llm.token_count.total

    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
