"""Filtering span exporter that only exports AI/LLM-related spans."""

import logging
from collections.abc import Sequence

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult

logger = logging.getLogger(__name__)

# OpenInference span kinds we want to capture
OPENINFERENCE_SPAN_KINDS = {
    "LLM",
    "EMBEDDING",
    "TOOL",
    "CHAIN",
    "AGENT",
    "RETRIEVER",
    "RERANKER",
}

# Attribute prefixes that indicate AI/LLM spans
AI_ATTRIBUTE_PREFIXES = [
    "openinference.",
    "gen_ai.",
    "llm.",
    "coalex.",
]


class FilteringSpanExporter(SpanExporter):
    """Span exporter that filters out non-AI/LLM spans.

    Only exports spans that have OpenInference span kind or
    AI/LLM/Coalex attribute prefixes.
    """

    def __init__(self, wrapped_exporter: SpanExporter, enabled: bool = True):
        self.wrapped_exporter = wrapped_exporter
        self.enabled = enabled
        self._filtered_count = 0
        self._exported_count = 0

    def _should_export(self, span: ReadableSpan) -> bool:
        if not self.enabled:
            return True

        attributes = dict(span.attributes or {})

        # Check 1: OpenInference span kind
        span_kind = attributes.get("openinference.span.kind")
        if span_kind and span_kind in OPENINFERENCE_SPAN_KINDS:
            return True

        # Check 2: AI/Coalex attribute prefixes
        for attr_name in attributes:
            for prefix in AI_ATTRIBUTE_PREFIXES:
                if attr_name.startswith(prefix):
                    return True

        return False

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        if not self.enabled:
            return self.wrapped_exporter.export(spans)

        filtered_spans = [s for s in spans if self._should_export(s)]

        self._filtered_count += len(spans) - len(filtered_spans)
        self._exported_count += len(filtered_spans)

        if filtered_spans:
            logger.debug(
                f"Exporting {len(filtered_spans)}/{len(spans)} spans to Coalex "
                f"(filtered out {len(spans) - len(filtered_spans)} non-AI spans)"
            )
            return self.wrapped_exporter.export(filtered_spans)

        return SpanExportResult.SUCCESS

    def shutdown(self) -> None:
        if self._filtered_count > 0 or self._exported_count > 0:
            logger.info(
                f"FilteringSpanExporter stats: exported={self._exported_count}, filtered={self._filtered_count}"
            )
        self.wrapped_exporter.shutdown()

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