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

import logging
from typing import Sequence
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
from opentelemetry.sdk.trace import ReadableSpan

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.",
]

# Specific attributes that indicate Coalex context
COALEX_CONTEXT_ATTRIBUTES = {
    "session.id",
    "prompt.version",
    "account_id",  # SDK sets this
}


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

    Only exports spans that:
    - Have OpenInference span kind (LLM, TOOL, CHAIN, etc.)
    - Have GenAI/LLM semantic convention attributes
    - Have Coalex context attributes (session.id, etc.)

    This prevents non-AI spans (Azure SDK, HTTP clients, etc.) from being
    sent to Coalex when customers use multiple auto-instrumentations.
    """

    def __init__(self, wrapped_exporter: SpanExporter, enabled: bool = True):
        """
        Initialize the filtering exporter.

        Args:
            wrapped_exporter: The actual OTLP exporter to wrap
            enabled: Whether filtering is enabled (default: True)
        """
        self.wrapped_exporter = wrapped_exporter
        self.enabled = enabled
        self._filtered_count = 0
        self._exported_count = 0

    def _should_export(self, span: ReadableSpan) -> bool:
        """
        Check if span should be exported to Coalex.

        Args:
            span: The span to check

        Returns:
            True if span should be exported, False otherwise
        """
        if not self.enabled:
            return True  # Filtering disabled, export all

        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 attribute prefixes
        for attr_name in attributes.keys():
            for prefix in AI_ATTRIBUTE_PREFIXES:
                if attr_name.startswith(prefix):
                    return True

        # Check 3: Coalex context attributes
        for attr_name in COALEX_CONTEXT_ATTRIBUTES:
            if attr_name in attributes and attributes[attr_name]:
                # session.id = "unknown" doesn't count
                if attr_name == "session.id" and attributes[attr_name] == "unknown":
                    continue
                return True

        return False

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        """
        Filter spans and export only AI/LLM-related ones.

        Args:
            spans: Sequence of spans to export

        Returns:
            SpanExportResult indicating success or failure
        """
        if not self.enabled:
            return self.wrapped_exporter.export(spans)

        # Filter 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)

        # No spans to export, but that's OK
        return SpanExportResult.SUCCESS

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

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        """Force flush the wrapped exporter."""
        return self.wrapped_exporter.force_flush(timeout_millis)
