"""Span processors and exporters for filtering and standardizing spans before export.

Includes:
- NoneAttributeFilterProcessor: Filters out None attribute values
- SpanSuppressionProcessor: Suppresses noisy framework spans (PROD-382)
- FilteringSpanExporter: Wraps any exporter to filter suppressed spans at export time
"""

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

logger = logging.getLogger(__name__)


class NoneAttributeFilterProcessor(SpanProcessor):
    """Span processor that filters out None values from span attributes.

    This prevents OpenTelemetry encoding errors when attributes are set to None,
    particularly for gen_ai.response.model and other attributes that might be
    extracted from response objects with None values.
    """

    def __init__(self, exporter: SpanExporter) -> None:
        """Initialize the filter processor with an exporter.

        Args:
            exporter: The span exporter to use after filtering
        """
        self._exporter = exporter

    def on_start(
        self, span: ReadableSpan, parent_context: Optional[object] = None
    ) -> None:
        """Called when a span is started."""
        # No filtering needed on start
        pass

    def on_end(self, span: ReadableSpan) -> None:
        """Called when a span ends. Filters None values before exporting."""
        # Filter out None values from attributes
        if hasattr(span, "attributes") and span.attributes:
            filtered_attributes = {}
            for key, value in span.attributes.items():
                if value is not None:
                    filtered_attributes[key] = value

            # Replace attributes with filtered version
            # Note: ReadableSpan.attributes is immutable (BoundedAttributes)
            # We can only modify internal _attributes on SDK Span implementation
            if hasattr(span, "_attributes"):
                span._attributes = filtered_attributes

        # Export the span
        self._exporter.export([span])

    def shutdown(self) -> None:
        """Shutdown the processor and exporter."""
        self._exporter.shutdown()

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

        Args:
            timeout_millis: Timeout in milliseconds (default: 30000)

        Returns:
            True if flush succeeded, False otherwise
        """
        if hasattr(self._exporter, "force_flush"):
            return self._exporter.force_flush(timeout_millis)
        return True


class SpanSuppressionProcessor(SpanProcessor):
    """Span processor that suppresses noisy framework spans at export time (PROD-382).

    Filters spans whose names match SUPPRESSED_SPAN_NAMES or start with
    suppressed prefixes. This removes auto-instrumented noise (AutoGen messaging,
    LangChain chat model spans, HTTP internals) without affecting Klira-owned spans.
    """

    def __init__(self, suppressed_names: Optional[Set[str]] = None) -> None:
        self._suppressed_names = suppressed_names or set()

    def on_start(
        self, span: ReadableSpan, parent_context: Optional[object] = None
    ) -> None:
        pass

    def on_end(self, span: ReadableSpan) -> None:
        span_name = span.name if hasattr(span, "name") else ""
        if self._should_suppress(span_name):
            # Mark as non-recording by setting internal flag
            if hasattr(span, "_attributes") and isinstance(span._attributes, dict):
                span._attributes["_klira_suppressed"] = True

    def _should_suppress(self, span_name: str) -> bool:
        """Check if a span should be suppressed."""
        # Exact match
        if span_name in self._suppressed_names:
            return True
        # Prefix match for spans like "autogen create RoundRobinGroupChat..."
        for suppressed in self._suppressed_names:
            if span_name.startswith(suppressed):
                return True
        return False

    def shutdown(self) -> None:
        pass

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


class FilteringSpanExporter(SpanExporter):
    """Wraps any SpanExporter to filter out suppressed spans before export (PROD-382).

    This is the primary suppression mechanism. It wraps Traceloop's OTLPSpanExporter
    so that noisy framework spans (AutoGen messaging, LangChain chat model spans,
    HTTP internals) are dropped before they reach the Klira API.
    """

    def __init__(self, inner: SpanExporter, suppressed_names: Set[str]) -> None:
        self._inner = inner
        self._suppressed_names = suppressed_names

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        filtered = [s for s in spans if not self._should_suppress(s)]
        if not filtered:
            return SpanExportResult.SUCCESS
        return self._inner.export(filtered)

    def _should_suppress(self, span: ReadableSpan) -> bool:
        name = span.name if hasattr(span, "name") else ""
        if name in self._suppressed_names:
            return True
        for suppressed in self._suppressed_names:
            if name.startswith(suppressed):
                return True
        return False

    def shutdown(self) -> None:
        self._inner.shutdown()

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        if hasattr(self._inner, "force_flush"):
            return self._inner.force_flush(timeout_millis)
        return True
