"""Span processors with masking support for Brizz SDK."""

import contextlib
import logging

from opentelemetry import context
from opentelemetry.context import Context
from opentelemetry.sdk.trace import ReadableSpan, Span
from opentelemetry.sdk.trace.export import (
    BatchSpanProcessor,
    SimpleSpanProcessor,
    SpanExporter,
)

from brizz._internal.models import AttributesMaskingRule, SpanMaskingConfig

from ...config import BrizzConfig
from ...masking.patterns import DEFAULT_PII_PATTERN_ENTRIES
from ...masking.utils import mask_attributes
from ...semantic_conventions import BRIZZ, PROPERTIES_CONTEXT_KEY

logger = logging.getLogger("brizz.masking")

# Default span masking rules
DEFAULT_MASKING_RULES = [
    AttributesMaskingRule(
        attribute_pattern="gen_ai.prompt",
        mode="partial",
        patterns=DEFAULT_PII_PATTERN_ENTRIES,
    ),
    AttributesMaskingRule(
        attribute_pattern="gen_ai.completion",
        mode="partial",
        patterns=DEFAULT_PII_PATTERN_ENTRIES,
    ),
    AttributesMaskingRule(
        attribute_pattern="traceloop.entity.input",
        mode="partial",
        patterns=DEFAULT_PII_PATTERN_ENTRIES,
    ),
    AttributesMaskingRule(
        attribute_pattern="traceloop.entity.output",
        mode="partial",
        patterns=DEFAULT_PII_PATTERN_ENTRIES,
    ),
]


def _apply_masking(span: Span, config: BrizzConfig) -> None:
    """Apply masking to a span if configured."""
    masking_config = getattr(config.masking, "span_masking", None) if config.masking else None
    if masking_config:
        try:
            _mask_span(span, masking_config)
        except Exception as error:
            logger.error("Error in span masking during on_start: %s", error)


def _read_association_properties(span: Span) -> dict[str, str] | None:
    """Read association properties from context and set them on the span."""
    association_properties = context.get_value(PROPERTIES_CONTEXT_KEY)
    if association_properties and hasattr(association_properties, "items"):
        properties = dict(association_properties.items())
        for key, value in properties.items():
            span.set_attribute(f"{BRIZZ}.{key}", value)
        return properties
    return None


def _reapply_properties(span: ReadableSpan, properties: dict[str, str]) -> None:
    """Re-apply brizz properties if they were evicted by BoundedAttributes.

    See docs/bugs/BRZ-759-conversation-view-missing-final-response.md

    OTel's Span.end() creates a new ReadableSpan but passes the same _attributes
    (BoundedAttributes) reference, so we write to it directly. span.attributes
    is a read-only mappingproxy and cannot be used for writes.
    """
    attrs = getattr(span, "_attributes", None)
    if attrs is None:
        return

    for key, value in properties.items():
        attr_key = f"{BRIZZ}.{key}"
        if attrs.get(attr_key) is None:
            with contextlib.suppress(TypeError, KeyError):
                attrs[attr_key] = value


class BrizzSimpleSpanProcessor(SimpleSpanProcessor):
    """Simple span processor with masking and context support."""

    def __init__(self, span_exporter: SpanExporter, config: BrizzConfig) -> None:
        super().__init__(span_exporter)
        self.config = config
        # Cache brizz properties per span keyed by (trace_id, span_id).
        # on_start receives a _Span, on_end receives a new ReadableSpan — different
        # objects, so we use a dict to pass properties between them.
        # Python's GIL makes dict get/set/pop atomic; no lock needed.
        # See docs/bugs/BRZ-759-conversation-view-missing-final-response.md
        self._props: dict[tuple[int, int], dict[str, str]] = {}

    def on_start(self, span: Span, parent_context: Context | None = None) -> None:
        # Apply PII masking rules to span attributes if configured
        _apply_masking(span, self.config)
        # Set brizz.* attributes from context (e.g. session.id) and cache them
        # so on_end can restore them if OTel's BoundedAttributes evicts them
        properties = _read_association_properties(span)
        if properties and hasattr(span, "context") and span.context:
            self._props[(span.context.trace_id, span.context.span_id)] = properties
        super().on_start(span, parent_context)

    def on_end(self, span: ReadableSpan) -> None:
        # Re-apply brizz.* attributes that may have been evicted by BoundedAttributes
        if hasattr(span, "context") and span.context:
            properties = self._props.pop((span.context.trace_id, span.context.span_id), None)
            if properties:
                _reapply_properties(span, properties)
        super().on_end(span)


class BrizzBatchSpanProcessor(BatchSpanProcessor):
    """Batch span processor with masking and context support."""

    def __init__(self, span_exporter: SpanExporter, config: BrizzConfig) -> None:
        super().__init__(span_exporter)
        self.config = config
        self._props: dict[tuple[int, int], dict[str, str]] = {}

    def on_start(self, span: Span, parent_context: Context | None = None) -> None:
        _apply_masking(span, self.config)
        properties = _read_association_properties(span)
        if properties and hasattr(span, "context") and span.context:
            self._props[(span.context.trace_id, span.context.span_id)] = properties
        super().on_start(span, parent_context)

    def on_end(self, span: ReadableSpan) -> None:
        if hasattr(span, "context") and span.context:
            properties = self._props.pop((span.context.trace_id, span.context.span_id), None)
            if properties:
                _reapply_properties(span, properties)
        super().on_end(span)


def _mask_span(span: Span, config: SpanMaskingConfig) -> Span:
    """Apply masking to a span based on the provided configuration."""
    if not span.attributes:
        return span

    rules = config.rules if config.rules else []
    if not getattr(config, "disable_default_rules", False):
        rules = rules + DEFAULT_MASKING_RULES

    try:
        masked_attributes = mask_attributes(span.attributes, rules, getattr(config, "_output_original_value", False))
        span.set_attributes(masked_attributes)
        return span
    except Exception as error:
        logger.error("Error masking span: %s", error)
        return span
