"""AgentBill Span Processor (P3-D)

Client-side span processing pipeline providing:
- Identity attribute injection (on_start)
- PII redaction (on_end)
- Deduplication (on_end)
- Head-based sampling with business event bypass (on_end)

[SPAN_PROCESSOR:P3-D] All code in this module is part of the P3-D implementation.
To find all P3-D code: grep -rn "SPAN_PROCESSOR:P3-D" sdks/python/

Configuration:
    pii_redaction: list[str]  - Rules: "STRIP_PROMPTS", "MASK_EMAILS", "MASK_KEYS"
    dedup_window: int         - LRU cache size for dedup (default: 1000)
    sampling_rate: float      - 0.0-1.0, 1.0 = keep all (default: 1.0)
"""
import re
import random
from collections import OrderedDict
from typing import Any, Dict, List, Optional, Set

# PII redaction rule constants
STRIP_PROMPTS = "STRIP_PROMPTS"
MASK_EMAILS = "MASK_EMAILS"
MASK_KEYS = "MASK_KEYS"

# Email regex pattern
_EMAIL_PATTERN = re.compile(
    r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z]{2,}\b',
    re.IGNORECASE,
)

# Key patterns to mask
_SENSITIVE_KEY_SUFFIXES = ("_key", "_secret", "_token", "_password")

# Prompt/completion attribute prefixes to strip
_PROMPT_PREFIXES = (
    "gen_ai.prompt.",
    "gen_ai.completion.",
    "agentbill.prompt_sample",
    "agentbill.response_sample",
)

# Protected attributes — never redacted by any rule
_PROTECTED_PREFIXES = (
    "service.",
    "telemetry.",
    "gen_ai.request.model",
    "gen_ai.system",
    "gen_ai.operation.name",
    "gen_ai.usage.",
    "agentbill.event_name",
    "agentbill.is_business_event",
)


def _is_protected(key: str) -> bool:
    """Check if an attribute key is protected from redaction."""
    for prefix in _PROTECTED_PREFIXES:
        if key == prefix or key.startswith(prefix):
            return True
    return False


class LRUDedup:
    """Bounded LRU cache for span deduplication by (trace_id, span_id)."""

    def __init__(self, maxsize: int = 1000):
        self._cache: OrderedDict = OrderedDict()
        self._maxsize = maxsize

    def seen(self, trace_id: str, span_id: str) -> bool:
        """Returns True if this span was already seen (duplicate)."""
        key = (trace_id, span_id)
        if key in self._cache:
            self._cache.move_to_end(key)
            return True
        self._cache[key] = True
        if len(self._cache) > self._maxsize:
            self._cache.popitem(last=False)
        return False


class AgentBillSpanProcessor:
    """[SPAN_PROCESSOR:P3-D] Client-side span processing pipeline.
    
    Integrates with both legacy and native OTEL modes.
    With default config (no redaction, rate=1.0), this is a zero-cost pass-through.
    """

    def __init__(
        self,
        customer_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        service_name: str = "agentbill-python-sdk",
        pii_redaction: Optional[List[str]] = None,
        dedup_window: int = 1000,
        sampling_rate: float = 1.0,
        debug: bool = False,
    ):
        self._customer_id = customer_id
        self._agent_id = agent_id
        self._service_name = service_name
        self._pii_rules: Set[str] = set(pii_redaction or [])
        self._dedup = LRUDedup(maxsize=dedup_window)
        self._sampling_rate = max(0.0, min(1.0, sampling_rate))
        self._debug = debug

    # ------------------------------------------------------------------
    # on_start: called when a span is created
    # ------------------------------------------------------------------
    def on_start(self, attributes: Dict[str, Any]) -> Dict[str, Any]:
        """[SPAN_PROCESSOR:P3-D] Inject identity attributes into span at creation.
        
        Called by tracer.start_span(). Mutates and returns the attributes dict.
        Context attributes (P3-C) are already merged by the tracer before this call.
        """
        # Identity injection
        if self._customer_id:
            attributes.setdefault("agentbill.customer_id", self._customer_id)
        if self._agent_id:
            attributes.setdefault("agentbill.agent_id", self._agent_id)
        attributes.setdefault("service.name", self._service_name)

        # Sampling decision (head-based)
        is_business = attributes.get("agentbill.is_business_event", False)
        if is_business or self._sampling_rate >= 1.0:
            attributes["agentbill.sampled"] = True
        else:
            sampled = random.random() < self._sampling_rate
            attributes["agentbill.sampled"] = sampled

        return attributes

    # ------------------------------------------------------------------
    # on_end: called when a span is finished, before export
    # ------------------------------------------------------------------
    def on_end(self, span_dict: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """[SPAN_PROCESSOR:P3-D] Process span before export.
        
        Args:
            span_dict: A dict with at minimum {trace_id, span_id, attributes}.
                       For legacy mode this is the Span object's __dict__; 
                       for native mode it's the ReadableSpan export record.
        
        Returns:
            The processed span_dict, or None if the span should be dropped.
        """
        trace_id = span_dict.get("trace_id", "")
        span_id = span_dict.get("span_id", "")
        attributes = span_dict.get("attributes", {})

        # 1. Sampling check — drop if not sampled
        if not attributes.get("agentbill.sampled", True):
            if self._debug:
                print(f"[AgentBill P3-D] Span dropped by sampling: {span_id[:8]}...")
            return None

        # 2. Deduplication
        if self._dedup.seen(trace_id, span_id):
            if self._debug:
                print(f"[AgentBill P3-D] Duplicate span dropped: {span_id[:8]}...")
            return None

        # 3. PII Redaction
        if self._pii_rules:
            span_dict["attributes"] = self._redact(attributes)

        return span_dict

    # ------------------------------------------------------------------
    # PII redaction pipeline
    # ------------------------------------------------------------------
    def _redact(self, attributes: Dict[str, Any]) -> Dict[str, Any]:
        """[SPAN_PROCESSOR:P3-D] Apply configured PII redaction rules."""
        result: Dict[str, Any] = {}

        for key, value in attributes.items():
            if _is_protected(key):
                result[key] = value
                continue

            # STRIP_PROMPTS: remove prompt/completion content
            if STRIP_PROMPTS in self._pii_rules:
                if any(key.startswith(p) or key == p for p in _PROMPT_PREFIXES):
                    continue  # drop this attribute

            # MASK_KEYS: remove sensitive key/secret/token/password attrs
            if MASK_KEYS in self._pii_rules:
                key_lower = key.lower()
                if any(key_lower.endswith(s) for s in _SENSITIVE_KEY_SUFFIXES):
                    continue  # drop this attribute

            # MASK_EMAILS: redact email patterns in string values
            if MASK_EMAILS in self._pii_rules and isinstance(value, str):
                value = _EMAIL_PATTERN.sub("[REDACTED_EMAIL]", value)

            result[key] = value

        return result
