# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/services/observability_service.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Observability Service Implementation.
This module provides OpenTelemetry-style observability for MCP Gateway,
capturing traces, spans, events, and metrics for all operations.

It includes:
- Trace creation and management
- Span tracking with hierarchical nesting
- Event logging within spans
- Metrics collection and storage
- Query and filtering capabilities
- Integration with FastAPI middleware

Examples:
    >>> from mcpgateway.services.observability_service import ObservabilityService  # doctest: +SKIP
    >>> service = ObservabilityService()  # doctest: +SKIP
    >>> trace_id = service.start_trace(db, "GET /tools", http_method="GET", http_url="/tools")  # doctest: +SKIP
    >>> span_id = service.start_span(db, trace_id, "database_query", resource_type="database")  # doctest: +SKIP
    >>> service.end_span(db, span_id, status="ok")  # doctest: +SKIP
    >>> service.end_trace(db, trace_id, status="ok", http_status_code=200)  # doctest: +SKIP
"""

# Standard
from contextlib import contextmanager
from contextvars import ContextVar
from datetime import datetime, timezone
import logging
import re
import traceback
from typing import Any, Dict, List, Optional, Pattern, Tuple
import uuid

# Third-Party
from sqlalchemy import desc
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import joinedload, Session

# First-Party
from mcpgateway.db import ObservabilityEvent, ObservabilityMetric, ObservabilitySpan, ObservabilityTrace

logger = logging.getLogger(__name__)

# Precompiled regex for W3C Trace Context traceparent header parsing
# Format: version-trace_id-parent_id-trace_flags
_TRACEPARENT_RE: Pattern[str] = re.compile(r"^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$")

# Context variable for tracking the current trace_id across async calls
current_trace_id: ContextVar[Optional[str]] = ContextVar("current_trace_id", default=None)


def utc_now() -> datetime:
    """Return current UTC time with timezone.

    Returns:
        datetime: Current time in UTC with timezone info
    """
    return datetime.now(timezone.utc)


def ensure_timezone_aware(dt: datetime) -> datetime:
    """Ensure datetime is timezone-aware (UTC).

    SQLite returns naive datetimes even when stored with timezone info.
    This helper ensures consistency for datetime arithmetic.

    Args:
        dt: Datetime that may be naive or aware

    Returns:
        Timezone-aware datetime in UTC
    """
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt


def parse_traceparent(traceparent: str) -> Optional[Tuple[str, str, str]]:
    """Parse W3C Trace Context traceparent header.

    Format: version-trace_id-parent_id-trace_flags
    Example: 00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01

    Args:
        traceparent: W3C traceparent header value

    Returns:
        Tuple of (trace_id, parent_span_id, trace_flags) or None if invalid

    Examples:
        >>> parse_traceparent("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")  # doctest: +SKIP
        ('0af7651916cd43dd8448eb211c80319c', 'b7ad6b7169203331', '01')
    """
    # W3C Trace Context format: 00-trace_id(32hex)-parent_id(16hex)-flags(2hex)
    # Uses precompiled regex for performance
    match = _TRACEPARENT_RE.match(traceparent.lower())

    if not match:
        logger.warning(f"Invalid traceparent format: {traceparent}")
        return None

    version, trace_id, parent_id, flags = match.groups()

    # Only support version 00 for now
    if version != "00":
        logger.warning(f"Unsupported traceparent version: {version}")
        return None

    # Validate trace_id and parent_id are not all zeros
    if trace_id == "0" * 32 or parent_id == "0" * 16:
        logger.warning("Invalid traceparent with zero trace_id or parent_id")
        return None

    return (trace_id, parent_id, flags)


def generate_w3c_trace_id() -> str:
    """Generate a W3C compliant trace ID (32 hex characters).

    Returns:
        32-character lowercase hex string

    Examples:
        >>> trace_id = generate_w3c_trace_id()  # doctest: +SKIP
        >>> len(trace_id)  # doctest: +SKIP
        32
    """
    return uuid.uuid4().hex + uuid.uuid4().hex[:16]


def generate_w3c_span_id() -> str:
    """Generate a W3C compliant span ID (16 hex characters).

    Returns:
        16-character lowercase hex string

    Examples:
        >>> span_id = generate_w3c_span_id()  # doctest: +SKIP
        >>> len(span_id)  # doctest: +SKIP
        16
    """
    return uuid.uuid4().hex[:16]


def format_traceparent(trace_id: str, span_id: str, sampled: bool = True) -> str:
    """Format a W3C traceparent header value.

    Args:
        trace_id: 32-character hex trace ID
        span_id: 16-character hex span ID
        sampled: Whether the trace is sampled (affects trace-flags)

    Returns:
        W3C traceparent header value

    Examples:
        >>> format_traceparent("0af7651916cd43dd8448eb211c80319c", "b7ad6b7169203331")  # doctest: +SKIP
        '00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01'
    """
    flags = "01" if sampled else "00"
    return f"00-{trace_id}-{span_id}-{flags}"


class ObservabilityService:
    """Service for managing observability traces, spans, events, and metrics.

    This service provides comprehensive observability capabilities similar to
    OpenTelemetry, allowing tracking of request flows through the system.

    Examples:
        >>> service = ObservabilityService()  # doctest: +SKIP
        >>> trace_id = service.start_trace(db, "POST /tools/invoke")  # doctest: +SKIP
        >>> span_id = service.start_span(db, trace_id, "tool_execution")  # doctest: +SKIP
        >>> service.end_span(db, span_id, status="ok")  # doctest: +SKIP
        >>> service.end_trace(db, trace_id, status="ok")  # doctest: +SKIP
    """

    def _safe_commit(self, db: Session, context: str) -> bool:
        """Commit and rollback on failure without raising.

        Args:
            db: SQLAlchemy session for the current operation.
            context: Short label for the commit context (used in logs).

        Returns:
            True when commit succeeds, False when a rollback was performed.
        """
        try:
            db.commit()
            return True
        except SQLAlchemyError as exc:
            logger.warning(f"Observability commit failed ({context}): {exc}")
            try:
                db.rollback()
            except SQLAlchemyError as rollback_exc:
                logger.debug(f"Observability rollback failed ({context}): {rollback_exc}")
            return False

    # ==============================
    # Trace Management
    # ==============================

    def start_trace(
        self,
        db: Session,
        name: str,
        trace_id: Optional[str] = None,
        parent_span_id: Optional[str] = None,
        http_method: Optional[str] = None,
        http_url: Optional[str] = None,
        user_email: Optional[str] = None,
        user_agent: Optional[str] = None,
        ip_address: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        resource_attributes: Optional[Dict[str, Any]] = None,
    ) -> str:
        """Start a new trace.

        Args:
            db: Database session
            name: Trace name (e.g., "POST /tools/invoke")
            trace_id: External trace ID (for distributed tracing, W3C format)
            parent_span_id: Parent span ID from upstream service
            http_method: HTTP method (GET, POST, etc.)
            http_url: Full request URL
            user_email: Authenticated user email
            user_agent: Client user agent string
            ip_address: Client IP address
            attributes: Additional trace attributes
            resource_attributes: Resource attributes (service name, version, etc.)

        Returns:
            Trace ID (UUID string or W3C format)

        Examples:
            >>> trace_id = service.start_trace(  # doctest: +SKIP
            ...     db,
            ...     "POST /tools/invoke",
            ...     http_method="POST",
            ...     http_url="https://api.example.com/tools/invoke",
            ...     user_email="user@example.com"
            ... )
        """
        # Use provided trace_id or generate new UUID
        if not trace_id:
            trace_id = str(uuid.uuid4())

        # Add parent context to attributes if provided
        attrs = attributes or {}
        if parent_span_id:
            attrs["parent_span_id"] = parent_span_id

        trace = ObservabilityTrace(
            trace_id=trace_id,
            name=name,
            start_time=utc_now(),
            status="unset",
            http_method=http_method,
            http_url=http_url,
            user_email=user_email,
            user_agent=user_agent,
            ip_address=ip_address,
            attributes=attrs,
            resource_attributes=resource_attributes or {},
            created_at=utc_now(),
        )
        db.add(trace)
        self._safe_commit(db, "start_trace")
        logger.debug(f"Started trace {trace_id}: {name}")
        return trace_id

    def end_trace(
        self,
        db: Session,
        trace_id: str,
        status: str = "ok",
        status_message: Optional[str] = None,
        http_status_code: Optional[int] = None,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> None:
        """End a trace.

        Args:
            db: Database session
            trace_id: Trace ID to end
            status: Trace status (ok, error)
            status_message: Optional status message
            http_status_code: HTTP response status code
            attributes: Additional attributes to merge

        Examples:
            >>> service.end_trace(  # doctest: +SKIP
            ...     db,
            ...     trace_id,
            ...     status="ok",
            ...     http_status_code=200
            ... )
        """
        trace = db.query(ObservabilityTrace).filter_by(trace_id=trace_id).first()
        if not trace:
            logger.warning(f"Trace {trace_id} not found")
            return

        end_time = utc_now()
        duration_ms = (end_time - ensure_timezone_aware(trace.start_time)).total_seconds() * 1000

        trace.end_time = end_time
        trace.duration_ms = duration_ms
        trace.status = status
        trace.status_message = status_message
        if http_status_code is not None:
            trace.http_status_code = http_status_code
        if attributes:
            trace.attributes = {**(trace.attributes or {}), **attributes}

        self._safe_commit(db, "end_trace")
        logger.debug(f"Ended trace {trace_id}: {status} ({duration_ms:.2f}ms)")

    def get_trace(self, db: Session, trace_id: str, include_spans: bool = False) -> Optional[ObservabilityTrace]:
        """Get a trace by ID.

        Args:
            db: Database session
            trace_id: Trace ID
            include_spans: Whether to load spans eagerly

        Returns:
            Trace object or None if not found

        Examples:
            >>> trace = service.get_trace(db, trace_id, include_spans=True)  # doctest: +SKIP
            >>> if trace:  # doctest: +SKIP
            ...     print(f"Trace: {trace.name}, Spans: {len(trace.spans)}")  # doctest: +SKIP
        """
        query = db.query(ObservabilityTrace).filter_by(trace_id=trace_id)
        if include_spans:
            query = query.options(joinedload(ObservabilityTrace.spans))
        return query.first()

    # ==============================
    # Span Management
    # ==============================

    def start_span(
        self,
        db: Session,
        trace_id: str,
        name: str,
        parent_span_id: Optional[str] = None,
        kind: str = "internal",
        resource_name: Optional[str] = None,
        resource_type: Optional[str] = None,
        resource_id: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        commit: bool = True,
    ) -> str:
        """Start a new span within a trace.

        Args:
            db: Database session
            trace_id: Parent trace ID
            name: Span name (e.g., "database_query", "tool_invocation")
            parent_span_id: Parent span ID (for nested spans)
            kind: Span kind (internal, server, client, producer, consumer)
            resource_name: Resource name being operated on
            resource_type: Resource type (tool, resource, prompt, etc.)
            resource_id: Resource ID
            attributes: Additional span attributes
            commit: Whether to commit the transaction (default True).
                Set to False when using fresh_db_session() which handles commits.

        Returns:
            Span ID (UUID string)

        Examples:
            >>> span_id = service.start_span(  # doctest: +SKIP
            ...     db,
            ...     trace_id,
            ...     "tool_invocation",
            ...     resource_type="tool",
            ...     resource_name="get_weather"
            ... )
        """
        span_id = str(uuid.uuid4())
        span = ObservabilitySpan(
            span_id=span_id,
            trace_id=trace_id,
            parent_span_id=parent_span_id,
            name=name,
            kind=kind,
            start_time=utc_now(),
            status="unset",
            resource_name=resource_name,
            resource_type=resource_type,
            resource_id=resource_id,
            attributes=attributes or {},
            created_at=utc_now(),
        )
        db.add(span)
        if commit:
            self._safe_commit(db, "start_span")
        logger.debug(f"Started span {span_id}: {name} (trace={trace_id})")
        return span_id

    def end_span(
        self,
        db: Session,
        span_id: str,
        status: str = "ok",
        status_message: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        commit: bool = True,
    ) -> None:
        """End a span.

        Args:
            db: Database session
            span_id: Span ID to end
            status: Span status (ok, error)
            status_message: Optional status message
            attributes: Additional attributes to merge
            commit: Whether to commit the transaction (default True).
                Set to False when using fresh_db_session() which handles commits.

        Examples:
            >>> service.end_span(db, span_id, status="ok")  # doctest: +SKIP
        """
        span = db.query(ObservabilitySpan).filter_by(span_id=span_id).first()
        if not span:
            logger.warning(f"Span {span_id} not found")
            return

        end_time = utc_now()
        duration_ms = (end_time - ensure_timezone_aware(span.start_time)).total_seconds() * 1000

        span.end_time = end_time
        span.duration_ms = duration_ms
        span.status = status
        span.status_message = status_message
        if attributes:
            span.attributes = {**(span.attributes or {}), **attributes}

        if commit:
            self._safe_commit(db, "end_span")
        logger.debug(f"Ended span {span_id}: {status} ({duration_ms:.2f}ms)")

    @contextmanager
    def trace_span(
        self,
        db: Session,
        trace_id: str,
        name: str,
        parent_span_id: Optional[str] = None,
        resource_type: Optional[str] = None,
        resource_name: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
    ):
        """Context manager for automatic span lifecycle management.

        Args:
            db: Database session
            trace_id: Parent trace ID
            name: Span name
            parent_span_id: Parent span ID (optional)
            resource_type: Resource type
            resource_name: Resource name
            attributes: Additional attributes

        Yields:
            Span ID

        Raises:
            Exception: Re-raises any exception after logging it in the span

        Examples:
            >>> with service.trace_span(db, trace_id, "database_query") as span_id:  # doctest: +SKIP
            ...     results = db.query(Tool).all()  # doctest: +SKIP
        """
        span_id = self.start_span(db, trace_id, name, parent_span_id, resource_type=resource_type, resource_name=resource_name, attributes=attributes)
        try:
            yield span_id
            self.end_span(db, span_id, status="ok")
        except Exception as e:
            self.end_span(db, span_id, status="error", status_message=str(e))
            self.add_event(db, span_id, "exception", severity="error", message=str(e), exception_type=type(e).__name__, exception_message=str(e), exception_stacktrace=traceback.format_exc())
            raise

    @contextmanager
    def trace_tool_invocation(
        self,
        db: Session,
        tool_name: str,
        arguments: Dict[str, Any],
        integration_type: Optional[str] = None,
    ):
        """Context manager for tracing MCP tool invocations.

        This automatically creates a span for tool execution, capturing timing,
        arguments, results, and errors.

        Args:
            db: Database session
            tool_name: Name of the tool being invoked
            arguments: Tool arguments (will be sanitized)
            integration_type: Integration type (MCP, REST, A2A, etc.)

        Yields:
            Tuple of (span_id, result_dict) - update result_dict with tool results

        Raises:
            Exception: Re-raises any exception from tool invocation after logging

        Examples:
            >>> with service.trace_tool_invocation(db, "weather", {"city": "NYC"}) as (span_id, result):  # doctest: +SKIP
            ...     response = await http_client.post(...)  # doctest: +SKIP
            ...     result["status_code"] = response.status_code  # doctest: +SKIP
            ...     result["response_size"] = len(response.content)  # doctest: +SKIP
        """
        trace_id = current_trace_id.get()
        if not trace_id:
            # No active trace, yield a no-op
            result_dict: Dict[str, Any] = {}
            yield (None, result_dict)
            return

        # Sanitize arguments (remove sensitive data)
        safe_args = {k: ("***REDACTED***" if any(sensitive in k.lower() for sensitive in ["password", "token", "key", "secret"]) else v) for k, v in arguments.items()}

        # Start tool invocation span
        span_id = self.start_span(
            db=db,
            trace_id=trace_id,
            name=f"tool.invoke.{tool_name}",
            kind="client",
            resource_type="tool",
            resource_name=tool_name,
            attributes={
                "tool.name": tool_name,
                "tool.integration_type": integration_type,
                "tool.argument_count": len(arguments),
                "tool.arguments": safe_args,
            },
        )

        result_dict = {}
        try:
            yield (span_id, result_dict)

            # End span with results
            self.end_span(
                db=db,
                span_id=span_id,
                status="ok",
                attributes={
                    "tool.result": result_dict,
                },
            )
        except Exception as e:
            # Log error in span
            self.end_span(db=db, span_id=span_id, status="error", status_message=str(e))

            self.add_event(
                db=db,
                span_id=span_id,
                name="tool.error",
                severity="error",
                message=str(e),
                exception_type=type(e).__name__,
                exception_message=str(e),
                exception_stacktrace=traceback.format_exc(),
            )
            raise

    # ==============================
    # Event Management
    # ==============================

    def add_event(
        self,
        db: Session,
        span_id: str,
        name: str,
        severity: Optional[str] = None,
        message: Optional[str] = None,
        exception_type: Optional[str] = None,
        exception_message: Optional[str] = None,
        exception_stacktrace: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> int:
        """Add an event to a span.

        Args:
            db: Database session
            span_id: Parent span ID
            name: Event name
            severity: Log severity (debug, info, warning, error, critical)
            message: Event message
            exception_type: Exception class name
            exception_message: Exception message
            exception_stacktrace: Exception stacktrace
            attributes: Additional event attributes

        Returns:
            Event ID

        Examples:
            >>> event_id = service.add_event(  # doctest: +SKIP
            ...     db,  # doctest: +SKIP
            ...     span_id,  # doctest: +SKIP
            ...     "database_connection_error",  # doctest: +SKIP
            ...     severity="error",  # doctest: +SKIP
            ...     message="Failed to connect to database"  # doctest: +SKIP
            ... )  # doctest: +SKIP
        """
        event = ObservabilityEvent(
            span_id=span_id,
            name=name,
            timestamp=utc_now(),
            severity=severity,
            message=message,
            exception_type=exception_type,
            exception_message=exception_message,
            exception_stacktrace=exception_stacktrace,
            attributes=attributes or {},
            created_at=utc_now(),
        )
        db.add(event)
        if not self._safe_commit(db, "add_event"):
            return 0
        db.refresh(event)
        logger.debug(f"Added event to span {span_id}: {name}")
        return event.id

    # ==============================
    # Token Usage Tracking
    # ==============================

    def record_token_usage(
        self,
        db: Session,
        span_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        model: Optional[str] = None,
        input_tokens: int = 0,
        output_tokens: int = 0,
        total_tokens: Optional[int] = None,
        estimated_cost_usd: Optional[float] = None,
        provider: Optional[str] = None,
    ) -> None:
        """Record token usage for LLM calls.

        Args:
            db: Database session
            span_id: Span ID to attach token usage to
            trace_id: Trace ID (will use current context if not provided)
            model: Model name (e.g., "gpt-4", "claude-3-opus")
            input_tokens: Number of input/prompt tokens
            output_tokens: Number of output/completion tokens
            total_tokens: Total tokens (calculated if not provided)
            estimated_cost_usd: Estimated cost in USD
            provider: LLM provider (openai, anthropic, etc.)

        Examples:
            >>> service.record_token_usage(  # doctest: +SKIP
            ...     db, span_id="abc123",
            ...     model="gpt-4",
            ...     input_tokens=100,
            ...     output_tokens=50,
            ...     estimated_cost_usd=0.015
            ... )
        """
        if not trace_id:
            trace_id = current_trace_id.get()

        if not trace_id:
            logger.warning("Cannot record token usage: no active trace")
            return

        # Calculate total if not provided
        if total_tokens is None:
            total_tokens = input_tokens + output_tokens

        # Estimate cost if not provided and we have model info
        if estimated_cost_usd is None and model:
            estimated_cost_usd = self._estimate_token_cost(model, input_tokens, output_tokens)

        # Store in span attributes if span_id provided
        if span_id:
            span = db.query(ObservabilitySpan).filter_by(span_id=span_id).first()
            if span:
                attrs = span.attributes or {}
                attrs.update(
                    {
                        "llm.model": model,
                        "llm.provider": provider,
                        "llm.input_tokens": input_tokens,
                        "llm.output_tokens": output_tokens,
                        "llm.total_tokens": total_tokens,
                        "llm.estimated_cost_usd": estimated_cost_usd,
                    }
                )
                span.attributes = attrs
                self._safe_commit(db, "record_token_usage")

        # Also record as metrics for aggregation
        if input_tokens > 0:
            self.record_metric(
                db=db,
                name="llm.tokens.input",
                value=float(input_tokens),
                metric_type="counter",
                unit="tokens",
                trace_id=trace_id,
                attributes={"model": model, "provider": provider},
            )

        if output_tokens > 0:
            self.record_metric(
                db=db,
                name="llm.tokens.output",
                value=float(output_tokens),
                metric_type="counter",
                unit="tokens",
                trace_id=trace_id,
                attributes={"model": model, "provider": provider},
            )

        if estimated_cost_usd:
            self.record_metric(
                db=db,
                name="llm.cost",
                value=estimated_cost_usd,
                metric_type="counter",
                unit="usd",
                trace_id=trace_id,
                attributes={"model": model, "provider": provider},
            )

        logger.debug(f"Recorded token usage: {input_tokens} in, {output_tokens} out, ${estimated_cost_usd:.6f}")

    def _estimate_token_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Estimate cost based on model and token counts.

        Pricing as of January 2025 (prices may change).

        Args:
            model: Model name
            input_tokens: Input token count
            output_tokens: Output token count

        Returns:
            Estimated cost in USD
        """
        # Pricing per 1M tokens (input, output)
        pricing = {
            # OpenAI
            "gpt-4": (30.0, 60.0),
            "gpt-4-turbo": (10.0, 30.0),
            "gpt-4o": (2.5, 10.0),
            "gpt-4o-mini": (0.15, 0.60),
            "gpt-3.5-turbo": (0.50, 1.50),
            # Anthropic
            "claude-3-opus": (15.0, 75.0),
            "claude-3-sonnet": (3.0, 15.0),
            "claude-3-haiku": (0.25, 1.25),
            "claude-3.5-sonnet": (3.0, 15.0),
            "claude-3.5-haiku": (0.80, 4.0),
            # Fallback for unknown models
            "default": (1.0, 3.0),
        }

        # Find matching pricing (case-insensitive, partial match)
        model_lower = model.lower()
        input_price, output_price = pricing.get("default")

        for model_key, prices in pricing.items():
            if model_key in model_lower:
                input_price, output_price = prices
                break

        # Calculate cost (pricing is per 1M tokens)
        input_cost = (input_tokens / 1_000_000) * input_price
        output_cost = (output_tokens / 1_000_000) * output_price

        return input_cost + output_cost

    # ==============================
    # Agent-to-Agent (A2A) Tracing
    # ==============================

    @contextmanager
    def trace_a2a_request(
        self,
        db: Session,
        agent_id: str,
        agent_name: Optional[str] = None,
        operation: Optional[str] = None,
        request_data: Optional[Dict[str, Any]] = None,
    ):
        """Context manager for tracing Agent-to-Agent requests.

        This automatically creates a span for A2A communication, capturing timing,
        request/response data, and errors.

        Args:
            db: Database session
            agent_id: Target agent ID
            agent_name: Human-readable agent name
            operation: Operation being performed (e.g., "query", "execute", "status")
            request_data: Request payload (will be sanitized)

        Yields:
            Tuple of (span_id, result_dict) - update result_dict with A2A results

        Raises:
            Exception: Re-raises any exception from A2A call after logging

        Examples:
            >>> with service.trace_a2a_request(db, "agent-123", "WeatherAgent", "query") as (span_id, result):  # doctest: +SKIP
            ...     response = await http_client.post(...)  # doctest: +SKIP
            ...     result["status_code"] = response.status_code  # doctest: +SKIP
            ...     result["response_time_ms"] = 45.2  # doctest: +SKIP
        """
        trace_id = current_trace_id.get()
        if not trace_id:
            # No active trace, yield a no-op
            result_dict: Dict[str, Any] = {}
            yield (None, result_dict)
            return

        # Sanitize request data
        safe_data = {}
        if request_data:
            safe_data = {k: ("***REDACTED***" if any(sensitive in k.lower() for sensitive in ["password", "token", "key", "secret", "auth"]) else v) for k, v in request_data.items()}

        # Start A2A span
        span_id = self.start_span(
            db=db,
            trace_id=trace_id,
            name=f"a2a.call.{agent_name or agent_id}",
            kind="client",
            resource_type="agent",
            resource_name=agent_name or agent_id,
            attributes={
                "a2a.agent_id": agent_id,
                "a2a.agent_name": agent_name,
                "a2a.operation": operation,
                "a2a.request_data": safe_data,
            },
        )

        result_dict = {}
        try:
            yield (span_id, result_dict)

            # End span with results
            self.end_span(
                db=db,
                span_id=span_id,
                status="ok",
                attributes={
                    "a2a.result": result_dict,
                },
            )
        except Exception as e:
            # Log error in span
            self.end_span(db=db, span_id=span_id, status="error", status_message=str(e))

            self.add_event(
                db=db,
                span_id=span_id,
                name="a2a.error",
                severity="error",
                message=str(e),
                exception_type=type(e).__name__,
                exception_message=str(e),
                exception_stacktrace=traceback.format_exc(),
            )
            raise

    # ==============================
    # Transport Metrics
    # ==============================

    def record_transport_activity(
        self,
        db: Session,
        transport_type: str,
        operation: str,
        message_count: int = 1,
        bytes_sent: Optional[int] = None,
        bytes_received: Optional[int] = None,
        connection_id: Optional[str] = None,
        error: Optional[str] = None,
    ) -> None:
        """Record transport-specific activity metrics.

        Args:
            db: Database session
            transport_type: Transport type (sse, websocket, stdio, http)
            operation: Operation type (connect, disconnect, send, receive, error)
            message_count: Number of messages processed
            bytes_sent: Bytes sent (if applicable)
            bytes_received: Bytes received (if applicable)
            connection_id: Connection/session identifier
            error: Error message if operation failed

        Examples:
            >>> service.record_transport_activity(  # doctest: +SKIP
            ...     db, transport_type="sse",
            ...     operation="send",
            ...     message_count=1,
            ...     bytes_sent=1024
            ... )
        """
        trace_id = current_trace_id.get()

        # Record message count
        if message_count > 0:
            self.record_metric(
                db=db,
                name=f"transport.{transport_type}.messages",
                value=float(message_count),
                metric_type="counter",
                unit="messages",
                trace_id=trace_id,
                attributes={
                    "transport": transport_type,
                    "operation": operation,
                    "connection_id": connection_id,
                },
            )

        # Record bytes sent
        if bytes_sent:
            self.record_metric(
                db=db,
                name=f"transport.{transport_type}.bytes_sent",
                value=float(bytes_sent),
                metric_type="counter",
                unit="bytes",
                trace_id=trace_id,
                attributes={
                    "transport": transport_type,
                    "operation": operation,
                    "connection_id": connection_id,
                },
            )

        # Record bytes received
        if bytes_received:
            self.record_metric(
                db=db,
                name=f"transport.{transport_type}.bytes_received",
                value=float(bytes_received),
                metric_type="counter",
                unit="bytes",
                trace_id=trace_id,
                attributes={
                    "transport": transport_type,
                    "operation": operation,
                    "connection_id": connection_id,
                },
            )

        # Record errors
        if error:
            self.record_metric(
                db=db,
                name=f"transport.{transport_type}.errors",
                value=1.0,
                metric_type="counter",
                unit="errors",
                trace_id=trace_id,
                attributes={
                    "transport": transport_type,
                    "operation": operation,
                    "connection_id": connection_id,
                    "error": error,
                },
            )

        logger.debug(f"Recorded {transport_type} transport activity: {operation} ({message_count} messages)")

    # ==============================
    # Metric Management
    # ==============================

    def record_metric(
        self,
        db: Session,
        name: str,
        value: float,
        metric_type: str = "gauge",
        unit: Optional[str] = None,
        resource_type: Optional[str] = None,
        resource_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> int:
        """Record a metric.

        Args:
            db: Database session
            name: Metric name (e.g., "http.request.duration")
            value: Metric value
            metric_type: Metric type (counter, gauge, histogram)
            unit: Metric unit (ms, count, bytes, etc.)
            resource_type: Resource type
            resource_id: Resource ID
            trace_id: Associated trace ID
            attributes: Additional metric attributes/labels

        Returns:
            Metric ID

        Examples:
            >>> metric_id = service.record_metric(  # doctest: +SKIP
            ...     db,  # doctest: +SKIP
            ...     "http.request.duration",  # doctest: +SKIP
            ...     123.45,  # doctest: +SKIP
            ...     metric_type="histogram",  # doctest: +SKIP
            ...     unit="ms",  # doctest: +SKIP
            ...     trace_id=trace_id  # doctest: +SKIP
            ... )  # doctest: +SKIP
        """
        metric = ObservabilityMetric(
            name=name,
            value=value,
            metric_type=metric_type,
            timestamp=utc_now(),
            unit=unit,
            resource_type=resource_type,
            resource_id=resource_id,
            trace_id=trace_id,
            attributes=attributes or {},
            created_at=utc_now(),
        )
        db.add(metric)
        if not self._safe_commit(db, "record_metric"):
            return 0
        db.refresh(metric)
        logger.debug(f"Recorded metric: {name} = {value} {unit or ''}")
        return metric.id

    # ==============================
    # Query Methods
    # ==============================

    # pylint: disable=too-many-positional-arguments,too-many-arguments,too-many-locals
    def query_traces(
        self,
        db: Session,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        min_duration_ms: Optional[float] = None,
        max_duration_ms: Optional[float] = None,
        status: Optional[str] = None,
        status_in: Optional[List[str]] = None,
        status_not_in: Optional[List[str]] = None,
        http_status_code: Optional[int] = None,
        http_status_code_in: Optional[List[int]] = None,
        http_method: Optional[str] = None,
        http_method_in: Optional[List[str]] = None,
        user_email: Optional[str] = None,
        user_email_in: Optional[List[str]] = None,
        attribute_filters: Optional[Dict[str, Any]] = None,
        attribute_filters_or: Optional[Dict[str, Any]] = None,
        attribute_search: Optional[str] = None,
        name_contains: Optional[str] = None,
        order_by: str = "start_time_desc",
        limit: int = 100,
        offset: int = 0,
    ) -> List[ObservabilityTrace]:
        """Query traces with advanced filters.

        Supports both simple filters (single value) and list filters (multiple values with OR logic).
        All top-level filters are combined with AND logic unless using _or suffix.

        Args:
            db: Database session
            start_time: Filter traces after this time
            end_time: Filter traces before this time
            min_duration_ms: Filter traces with duration >= this value (milliseconds)
            max_duration_ms: Filter traces with duration <= this value (milliseconds)
            status: Filter by single status (ok, error)
            status_in: Filter by multiple statuses (OR logic)
            status_not_in: Exclude these statuses (NOT logic)
            http_status_code: Filter by single HTTP status code
            http_status_code_in: Filter by multiple HTTP status codes (OR logic)
            http_method: Filter by single HTTP method (GET, POST, etc.)
            http_method_in: Filter by multiple HTTP methods (OR logic)
            user_email: Filter by single user email
            user_email_in: Filter by multiple user emails (OR logic)
            attribute_filters: JSON attribute filters (AND logic - all must match)
            attribute_filters_or: JSON attribute filters (OR logic - any must match)
            attribute_search: Free-text search within JSON attributes (partial match)
            name_contains: Filter traces where name contains this substring
            order_by: Sort order (start_time_desc, start_time_asc, duration_desc, duration_asc)
            limit: Maximum results (1-1000)
            offset: Result offset

        Returns:
            List of traces

        Raises:
            ValueError: If invalid parameters are provided

        Examples:
            >>> # Find slow errors from multiple endpoints
            >>> traces = service.query_traces(  # doctest: +SKIP
            ...     db,
            ...     status="error",
            ...     min_duration_ms=100.0,
            ...     http_method_in=["POST", "PUT"],
            ...     attribute_filters={"http.route": "/api/tools"},
            ...     limit=50
            ... )
            >>> # Exclude health checks and find slow requests
            >>> traces = service.query_traces(  # doctest: +SKIP
            ...     db,
            ...     min_duration_ms=1000.0,
            ...     name_contains="api",
            ...     status_not_in=["ok"],
            ...     order_by="duration_desc"
            ... )
        """
        # Third-Party
        # pylint: disable=import-outside-toplevel
        from sqlalchemy import cast, or_, String

        # pylint: enable=import-outside-toplevel
        # Validate limit
        if limit < 1 or limit > 1000:
            raise ValueError("limit must be between 1 and 1000")

        # Validate order_by
        valid_orders = ["start_time_desc", "start_time_asc", "duration_desc", "duration_asc"]
        if order_by not in valid_orders:
            raise ValueError(f"order_by must be one of: {', '.join(valid_orders)}")

        query = db.query(ObservabilityTrace)

        # Time range filters
        if start_time:
            query = query.filter(ObservabilityTrace.start_time >= start_time)
        if end_time:
            query = query.filter(ObservabilityTrace.start_time <= end_time)

        # Duration filters
        if min_duration_ms is not None:
            query = query.filter(ObservabilityTrace.duration_ms >= min_duration_ms)
        if max_duration_ms is not None:
            query = query.filter(ObservabilityTrace.duration_ms <= max_duration_ms)

        # Status filters (with OR and NOT support)
        if status:
            query = query.filter(ObservabilityTrace.status == status)
        if status_in:
            query = query.filter(ObservabilityTrace.status.in_(status_in))
        if status_not_in:
            query = query.filter(~ObservabilityTrace.status.in_(status_not_in))

        # HTTP status code filters (with OR support)
        if http_status_code:
            query = query.filter(ObservabilityTrace.http_status_code == http_status_code)
        if http_status_code_in:
            query = query.filter(ObservabilityTrace.http_status_code.in_(http_status_code_in))

        # HTTP method filters (with OR support)
        if http_method:
            query = query.filter(ObservabilityTrace.http_method == http_method)
        if http_method_in:
            query = query.filter(ObservabilityTrace.http_method.in_(http_method_in))

        # User email filters (with OR support)
        if user_email:
            query = query.filter(ObservabilityTrace.user_email == user_email)
        if user_email_in:
            query = query.filter(ObservabilityTrace.user_email.in_(user_email_in))

        # Name substring filter
        if name_contains:
            query = query.filter(ObservabilityTrace.name.ilike(f"%{name_contains}%"))

        # Attribute-based filtering with AND logic (all filters must match)
        if attribute_filters:
            for key, value in attribute_filters.items():
                # Use JSON path access for filtering
                # Supports both SQLite (via json_extract) and PostgreSQL (via ->>)
                query = query.filter(ObservabilityTrace.attributes[key].astext == str(value))

        # Attribute-based filtering with OR logic (any filter must match)
        if attribute_filters_or:
            or_conditions = []
            for key, value in attribute_filters_or.items():
                or_conditions.append(ObservabilityTrace.attributes[key].astext == str(value))
            if or_conditions:
                query = query.filter(or_(*or_conditions))

        # Free-text search across all attribute values
        if attribute_search:
            # Cast JSON attributes to text and search for substring
            # Works with both SQLite and PostgreSQL
            # Escape special characters to prevent SQL injection
            safe_search = attribute_search.replace("%", "\\%").replace("_", "\\_")
            query = query.filter(cast(ObservabilityTrace.attributes, String).ilike(f"%{safe_search}%"))

        # Apply ordering
        if order_by == "start_time_desc":
            query = query.order_by(desc(ObservabilityTrace.start_time))
        elif order_by == "start_time_asc":
            query = query.order_by(ObservabilityTrace.start_time)
        elif order_by == "duration_desc":
            query = query.order_by(desc(ObservabilityTrace.duration_ms))
        else:  # duration_asc (validated above)
            query = query.order_by(ObservabilityTrace.duration_ms)

        # Apply pagination
        query = query.limit(limit).offset(offset)

        return query.all()

    # pylint: disable=too-many-positional-arguments,too-many-arguments,too-many-locals
    def query_spans(
        self,
        db: Session,
        trace_id: Optional[str] = None,
        trace_id_in: Optional[List[str]] = None,
        resource_type: Optional[str] = None,
        resource_type_in: Optional[List[str]] = None,
        resource_name: Optional[str] = None,
        resource_name_in: Optional[List[str]] = None,
        name_contains: Optional[str] = None,
        kind: Optional[str] = None,
        kind_in: Optional[List[str]] = None,
        status: Optional[str] = None,
        status_in: Optional[List[str]] = None,
        status_not_in: Optional[List[str]] = None,
        start_time: Optional[datetime] = None,
        end_time: Optional[datetime] = None,
        min_duration_ms: Optional[float] = None,
        max_duration_ms: Optional[float] = None,
        attribute_filters: Optional[Dict[str, Any]] = None,
        attribute_search: Optional[str] = None,
        order_by: str = "start_time_desc",
        limit: int = 100,
        offset: int = 0,
    ) -> List[ObservabilitySpan]:
        """Query spans with advanced filters.

        Supports filtering by trace, resource, kind, status, duration, and attributes.
        All top-level filters are combined with AND logic. List filters use OR logic.

        Args:
            db: Database session
            trace_id: Filter by single trace ID
            trace_id_in: Filter by multiple trace IDs (OR logic)
            resource_type: Filter by single resource type (tool, database, plugin, etc.)
            resource_type_in: Filter by multiple resource types (OR logic)
            resource_name: Filter by single resource name
            resource_name_in: Filter by multiple resource names (OR logic)
            name_contains: Filter spans where name contains this substring
            kind: Filter by span kind (client, server, internal)
            kind_in: Filter by multiple kinds (OR logic)
            status: Filter by single status (ok, error)
            status_in: Filter by multiple statuses (OR logic)
            status_not_in: Exclude these statuses (NOT logic)
            start_time: Filter spans after this time
            end_time: Filter spans before this time
            min_duration_ms: Filter spans with duration >= this value (milliseconds)
            max_duration_ms: Filter spans with duration <= this value (milliseconds)
            attribute_filters: JSON attribute filters (AND logic)
            attribute_search: Free-text search within JSON attributes
            order_by: Sort order (start_time_desc, start_time_asc, duration_desc, duration_asc)
            limit: Maximum results (1-1000)
            offset: Result offset

        Returns:
            List of spans

        Raises:
            ValueError: If invalid parameters are provided

        Examples:
            >>> # Find slow database queries
            >>> spans = service.query_spans(  # doctest: +SKIP
            ...     db,
            ...     resource_type="database",
            ...     min_duration_ms=100.0,
            ...     order_by="duration_desc",
            ...     limit=50
            ... )
            >>> # Find tool invocation errors
            >>> spans = service.query_spans(  # doctest: +SKIP
            ...     db,
            ...     resource_type="tool",
            ...     status="error",
            ...     name_contains="invoke"
            ... )
        """
        # Third-Party
        # pylint: disable=import-outside-toplevel
        from sqlalchemy import cast, String

        # pylint: enable=import-outside-toplevel
        # Validate limit
        if limit < 1 or limit > 1000:
            raise ValueError("limit must be between 1 and 1000")

        # Validate order_by
        valid_orders = ["start_time_desc", "start_time_asc", "duration_desc", "duration_asc"]
        if order_by not in valid_orders:
            raise ValueError(f"order_by must be one of: {', '.join(valid_orders)}")

        query = db.query(ObservabilitySpan)

        # Trace ID filters (with OR support)
        if trace_id:
            query = query.filter(ObservabilitySpan.trace_id == trace_id)
        if trace_id_in:
            query = query.filter(ObservabilitySpan.trace_id.in_(trace_id_in))

        # Resource type filters (with OR support)
        if resource_type:
            query = query.filter(ObservabilitySpan.resource_type == resource_type)
        if resource_type_in:
            query = query.filter(ObservabilitySpan.resource_type.in_(resource_type_in))

        # Resource name filters (with OR support)
        if resource_name:
            query = query.filter(ObservabilitySpan.resource_name == resource_name)
        if resource_name_in:
            query = query.filter(ObservabilitySpan.resource_name.in_(resource_name_in))

        # Name substring filter
        if name_contains:
            query = query.filter(ObservabilitySpan.name.ilike(f"%{name_contains}%"))

        # Kind filters (with OR support)
        if kind:
            query = query.filter(ObservabilitySpan.kind == kind)
        if kind_in:
            query = query.filter(ObservabilitySpan.kind.in_(kind_in))

        # Status filters (with OR and NOT support)
        if status:
            query = query.filter(ObservabilitySpan.status == status)
        if status_in:
            query = query.filter(ObservabilitySpan.status.in_(status_in))
        if status_not_in:
            query = query.filter(~ObservabilitySpan.status.in_(status_not_in))

        # Time range filters
        if start_time:
            query = query.filter(ObservabilitySpan.start_time >= start_time)
        if end_time:
            query = query.filter(ObservabilitySpan.start_time <= end_time)

        # Duration filters
        if min_duration_ms is not None:
            query = query.filter(ObservabilitySpan.duration_ms >= min_duration_ms)
        if max_duration_ms is not None:
            query = query.filter(ObservabilitySpan.duration_ms <= max_duration_ms)

        # Attribute-based filtering with AND logic
        if attribute_filters:
            for key, value in attribute_filters.items():
                query = query.filter(ObservabilitySpan.attributes[key].astext == str(value))

        # Free-text search across all attribute values
        if attribute_search:
            safe_search = attribute_search.replace("%", "\\%").replace("_", "\\_")
            query = query.filter(cast(ObservabilitySpan.attributes, String).ilike(f"%{safe_search}%"))

        # Apply ordering
        if order_by == "start_time_desc":
            query = query.order_by(desc(ObservabilitySpan.start_time))
        elif order_by == "start_time_asc":
            query = query.order_by(ObservabilitySpan.start_time)
        elif order_by == "duration_desc":
            query = query.order_by(desc(ObservabilitySpan.duration_ms))
        else:  # duration_asc (validated above)
            query = query.order_by(ObservabilitySpan.duration_ms)

        # Apply pagination
        query = query.limit(limit).offset(offset)

        return query.all()

    def get_trace_with_spans(self, db: Session, trace_id: str) -> Optional[ObservabilityTrace]:
        """Get a complete trace with all spans and events.

        Args:
            db: Database session
            trace_id: Trace ID

        Returns:
            Trace with spans and events loaded

        Examples:
            >>> trace = service.get_trace_with_spans(db, trace_id)  # doctest: +SKIP
            >>> if trace:  # doctest: +SKIP
            ...     for span in trace.spans:  # doctest: +SKIP
            ...         print(f"Span: {span.name}, Events: {len(span.events)}")  # doctest: +SKIP
        """
        return db.query(ObservabilityTrace).filter_by(trace_id=trace_id).options(joinedload(ObservabilityTrace.spans).joinedload(ObservabilitySpan.events)).first()

    def delete_old_traces(self, db: Session, before_time: datetime) -> int:
        """Delete traces older than a given time.

        Args:
            db: Database session
            before_time: Delete traces before this time

        Returns:
            Number of traces deleted

        Examples:
            >>> from datetime import timedelta  # doctest: +SKIP
            >>> cutoff = utc_now() - timedelta(days=30)  # doctest: +SKIP
            >>> deleted = service.delete_old_traces(db, cutoff)  # doctest: +SKIP
            >>> print(f"Deleted {deleted} old traces")  # doctest: +SKIP
        """
        deleted = db.query(ObservabilityTrace).filter(ObservabilityTrace.start_time < before_time).delete()
        if not self._safe_commit(db, "delete_old_traces"):
            return 0
        logger.info(f"Deleted {deleted} traces older than {before_time}")
        return deleted
