MaverickMCP

"""
OpenTelemetry distributed tracing integration for MaverickMCP.

This module provides comprehensive distributed tracing capabilities including:
- Automatic span creation for database queries, external API calls, and tool executions
- Integration with FastMCP and FastAPI
- Support for multiple tracing backends (Jaeger, Zipkin, OTLP)
- Correlation with structured logging
"""

import functools
import os
import time
from collections.abc import Callable
from contextlib import contextmanager
from typing import Any

from maverick_mcp.config.settings import settings
from maverick_mcp.utils.logging import get_logger

# OpenTelemetry imports with graceful fallback
try:
    from opentelemetry import trace  # type: ignore[import-untyped]
    from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
        OTLPSpanExporter,  # type: ignore[import-untyped]
    )
    from opentelemetry.exporter.zipkin.json import (
        ZipkinExporter,  # type: ignore[import-untyped]
    )
    from opentelemetry.instrumentation.asyncio import (
        AsyncioInstrumentor,  # type: ignore[import-untyped]
    )
    from opentelemetry.instrumentation.asyncpg import (
        AsyncPGInstrumentor,  # type: ignore[import-untyped]
    )
    from opentelemetry.instrumentation.fastapi import (
        FastAPIInstrumentor,  # type: ignore[import-untyped]
    )
    from opentelemetry.instrumentation.httpx import (
        HTTPXInstrumentor,  # type: ignore[import-untyped]
    )
    from opentelemetry.instrumentation.redis import (
        RedisInstrumentor,  # type: ignore[import-untyped]
    )
    from opentelemetry.instrumentation.requests import (
        RequestsInstrumentor,  # type: ignore[import-untyped]
    )
    from opentelemetry.instrumentation.sqlalchemy import (
        SQLAlchemyInstrumentor,  # type: ignore[import-untyped]
    )
    from opentelemetry.propagate import (
        set_global_textmap,  # type: ignore[import-untyped]
    )
    from opentelemetry.propagators.b3 import (
        B3MultiFormat,  # type: ignore[import-untyped]
    )
    from opentelemetry.sdk.resources import Resource  # type: ignore[import-untyped]
    from opentelemetry.sdk.trace import TracerProvider  # type: ignore[import-untyped]
    from opentelemetry.sdk.trace.export import (  # type: ignore[import-untyped]
        BatchSpanProcessor,
        ConsoleSpanExporter,
    )
    from opentelemetry.semconv.resource import (
        ResourceAttributes,  # type: ignore[import-untyped]
    )
    from opentelemetry.trace import Status, StatusCode  # type: ignore[import-untyped]

    OTEL_AVAILABLE = True
except ImportError:
    # Create stub classes for when OpenTelemetry is not available
    class _TracerStub:
        def start_span(self, name: str, **kwargs):
            return _SpanStub()

        def start_as_current_span(self, name: str, **kwargs):
            return _SpanStub()

    class _SpanStub:
        def __enter__(self):
            return self

        def __exit__(self, *args):
            pass

        def set_attribute(self, key: str, value: Any):
            pass

        def set_status(self, status):
            pass

        def record_exception(self, exception: Exception):
            pass

        def add_event(self, name: str, attributes: dict[str, Any] | None = None):
            pass

    # Create stub types for type annotations
    class TracerProvider:
        pass

    trace = type("trace", (), {"get_tracer": lambda name: _TracerStub()})()
    OTEL_AVAILABLE = False


logger = get_logger(__name__)


class TracingService:
    """Service for distributed tracing configuration and management."""

    def __init__(self):
        self.tracer = None
        self.enabled = False
        self._initialize_tracing()

    def _initialize_tracing(self):
        """Initialize OpenTelemetry tracing."""
        if not OTEL_AVAILABLE:
            return

        # Check if tracing is enabled
        tracing_enabled = os.getenv("OTEL_TRACING_ENABLED", "false").lower() == "true"
        if not tracing_enabled and settings.environment != "development":
            logger.info("OpenTelemetry tracing disabled")
            return

        try:
            # Create resource
            resource = Resource.create(
                {
                    ResourceAttributes.SERVICE_NAME: settings.app_name,
                    ResourceAttributes.SERVICE_VERSION: os.getenv(
                        "RELEASE_VERSION", "unknown"
                    ),
                    ResourceAttributes.SERVICE_NAMESPACE: "maverick-mcp",
                    ResourceAttributes.DEPLOYMENT_ENVIRONMENT: settings.environment,
                }
            )

            # Configure tracer provider
            tracer_provider = TracerProvider(resource=resource)
            trace.set_tracer_provider(tracer_provider)

            # Configure exporters
            self._configure_exporters(tracer_provider)

            # Configure propagators
            self._configure_propagators()

            # Instrument libraries
            self._instrument_libraries()

            # Create tracer
            self.tracer = trace.get_tracer(__name__)
            self.enabled = True

            logger.info("OpenTelemetry tracing initialized successfully")

        except Exception as e:
            logger.error(f"Failed to initialize OpenTelemetry tracing: {e}")

    def _configure_exporters(self, tracer_provider: TracerProvider):
        """Configure trace exporters based on environment variables."""
        # Console exporter (for development)
        if settings.environment == "development":
            console_exporter = ConsoleSpanExporter()
            tracer_provider.add_span_processor(BatchSpanProcessor(console_exporter))  # type: ignore[attr-defined]

        # Jaeger exporter via OTLP (modern approach)
        jaeger_endpoint = os.getenv("JAEGER_ENDPOINT")
        if jaeger_endpoint:
            # Modern Jaeger deployments accept OTLP on port 4317 (gRPC) or 4318 (HTTP)
            # Convert legacy Jaeger collector endpoint to OTLP format if needed
            if "14268" in jaeger_endpoint:  # Legacy Jaeger HTTP port
                otlp_endpoint = jaeger_endpoint.replace(":14268", ":4318").replace(
                    "/api/traces", ""
                )
                logger.info(
                    f"Converting legacy Jaeger endpoint {jaeger_endpoint} to OTLP: {otlp_endpoint}"
                )
            else:
                otlp_endpoint = jaeger_endpoint

            jaeger_otlp_exporter = OTLPSpanExporter(
                endpoint=otlp_endpoint,
                # Add Jaeger-specific headers if needed
                headers={},
            )
            tracer_provider.add_span_processor(BatchSpanProcessor(jaeger_otlp_exporter))  # type: ignore[attr-defined]
            logger.info(f"Jaeger OTLP exporter configured: {otlp_endpoint}")

        # Zipkin exporter
        zipkin_endpoint = os.getenv("ZIPKIN_ENDPOINT")
        if zipkin_endpoint:
            zipkin_exporter = ZipkinExporter(endpoint=zipkin_endpoint)
            tracer_provider.add_span_processor(BatchSpanProcessor(zipkin_exporter))  # type: ignore[attr-defined]
            logger.info(f"Zipkin exporter configured: {zipkin_endpoint}")

        # OTLP exporter (for services like Honeycomb, New Relic, etc.)
        otlp_endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")
        if otlp_endpoint:
            otlp_exporter = OTLPSpanExporter(
                endpoint=otlp_endpoint,
                headers={"x-honeycomb-team": os.getenv("HONEYCOMB_API_KEY", "")},
            )
            tracer_provider.add_span_processor(BatchSpanProcessor(otlp_exporter))  # type: ignore[attr-defined]
            logger.info(f"OTLP exporter configured: {otlp_endpoint}")

    def _configure_propagators(self):
        """Configure trace propagators for cross-service communication."""
        # Use B3 propagator for maximum compatibility
        set_global_textmap(B3MultiFormat())
        logger.info("B3 trace propagator configured")

    def _instrument_libraries(self):
        """Automatically instrument common libraries."""
        try:
            # FastAPI instrumentation
            FastAPIInstrumentor().instrument()

            # Database instrumentation
            SQLAlchemyInstrumentor().instrument()
            AsyncPGInstrumentor().instrument()

            # HTTP client instrumentation
            RequestsInstrumentor().instrument()
            HTTPXInstrumentor().instrument()

            # Redis instrumentation
            RedisInstrumentor().instrument()

            # Asyncio instrumentation
            AsyncioInstrumentor().instrument()

            logger.info("Auto-instrumentation completed successfully")

        except Exception as e:
            logger.warning(f"Some auto-instrumentation failed: {e}")

    @contextmanager
    def trace_operation(
        self,
        operation_name: str,
        attributes: dict[str, Any] | None = None,
        record_exception: bool = True,
    ):
        """
        Context manager for tracing operations.

        Args:
            operation_name: Name of the operation being traced
            attributes: Additional attributes to add to the span
            record_exception: Whether to record exceptions in the span
        """
        if not self.enabled:
            yield None
            return

        with self.tracer.start_as_current_span(operation_name) as span:
            # Add attributes
            if attributes:
                for key, value in attributes.items():
                    span.set_attribute(key, str(value))

            try:
                yield span
                span.set_status(Status(StatusCode.OK))
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR, str(e)))
                if record_exception:
                    span.record_exception(e)
                raise

    def trace_tool_execution(self, func: Callable) -> Callable:
        """
        Decorator to trace tool execution.

        Args:
            func: The tool function to trace

        Returns:
            Decorated function with tracing
        """

        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            if not self.enabled:
                return await func(*args, **kwargs)

            tool_name = getattr(func, "__name__", "unknown_tool")
            with self.trace_operation(
                f"tool.{tool_name}",
                attributes={
                    "tool.name": tool_name,
                    "tool.args_count": len(args),
                    "tool.kwargs_count": len(kwargs),
                },
            ) as span:
                # Add user context if available
                for arg in args:
                    if hasattr(arg, "user_id"):
                        span.set_attribute("user.id", str(arg.user_id))
                        break

                start_time = time.time()
                result = await func(*args, **kwargs)
                duration = time.time() - start_time

                span.set_attribute("tool.duration_seconds", duration)
                span.set_attribute("tool.success", True)

                return result

        return wrapper

    def trace_database_query(
        self, query_type: str, table: str | None = None, query: str | None = None
    ):
        """
        Context manager for tracing database queries.

        Args:
            query_type: Type of query (SELECT, INSERT, UPDATE, DELETE)
            table: Table name being queried
            query: The actual SQL query (will be truncated for security)
        """
        attributes = {
            "db.operation": query_type,
            "db.system": "postgresql",
        }

        if table:
            attributes["db.table"] = table

        if query:
            # Truncate query for security and performance
            attributes["db.statement"] = (
                query[:200] + "..." if len(query) > 200 else query
            )

        return self.trace_operation(f"db.{query_type.lower()}", attributes)

    def trace_external_api_call(self, service: str, endpoint: str, method: str = "GET"):
        """
        Context manager for tracing external API calls.

        Args:
            service: Name of the external service
            endpoint: API endpoint being called
            method: HTTP method
        """
        attributes = {
            "http.method": method,
            "http.url": endpoint,
            "service.name": service,
        }

        return self.trace_operation(f"http.{method.lower()}", attributes)

    def trace_cache_operation(self, operation: str, cache_type: str = "redis"):
        """
        Context manager for tracing cache operations.

        Args:
            operation: Cache operation (get, set, delete, etc.)
            cache_type: Type of cache (redis, memory, etc.)
        """
        attributes = {
            "cache.operation": operation,
            "cache.type": cache_type,
        }

        return self.trace_operation(f"cache.{operation}", attributes)

    def add_event(self, name: str, attributes: dict[str, Any] | None = None):
        """Add an event to the current span."""
        if not self.enabled:
            return

        current_span = trace.get_current_span()
        if current_span:
            current_span.add_event(name, attributes or {})

    def set_user_context(self, user_id: str, email: str | None = None):
        """Set user context on the current span."""
        if not self.enabled:
            return

        current_span = trace.get_current_span()
        if current_span:
            current_span.set_attribute("user.id", user_id)
            if email:
                current_span.set_attribute("user.email", email)


# Global tracing service instance
_tracing_service: TracingService | None = None


def get_tracing_service() -> TracingService:
    """Get or create the global tracing service."""
    global _tracing_service
    if _tracing_service is None:
        _tracing_service = TracingService()
    return _tracing_service


def trace_tool(func: Callable) -> Callable:
    """Decorator for tracing tool execution."""
    tracing = get_tracing_service()
    return tracing.trace_tool_execution(func)


@contextmanager
def trace_operation(
    operation_name: str,
    attributes: dict[str, Any] | None = None,
    record_exception: bool = True,
):
    """Context manager for tracing operations."""
    tracing = get_tracing_service()
    with tracing.trace_operation(operation_name, attributes, record_exception) as span:
        yield span


@contextmanager
def trace_database_query(
    query_type: str, table: str | None = None, query: str | None = None
):
    """Context manager for tracing database queries."""
    tracing = get_tracing_service()
    with tracing.trace_database_query(query_type, table, query) as span:
        yield span


@contextmanager
def trace_external_api_call(service: str, endpoint: str, method: str = "GET"):
    """Context manager for tracing external API calls."""
    tracing = get_tracing_service()
    with tracing.trace_external_api_call(service, endpoint, method) as span:
        yield span


@contextmanager
def trace_cache_operation(operation: str, cache_type: str = "redis"):
    """Context manager for tracing cache operations."""
    tracing = get_tracing_service()
    with tracing.trace_cache_operation(operation, cache_type) as span:
        yield span


def initialize_tracing():
    """Initialize the global tracing service."""
    logger.info("Initializing distributed tracing...")
    tracing = get_tracing_service()

    if tracing.enabled:
        logger.info("Distributed tracing initialized successfully")
    else:
        logger.info("Distributed tracing disabled or unavailable")