"""
Monitoring and Observability Integration Example

Demonstrates how to set up OpenTelemetry instrumentation
for MCP tools and Django endpoints.
"""

from opentelemetry import trace, metrics
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.exporter.jaeger.thrift import JaegerExporter as JaegerMetricsExporter
from opentelemetry.instrumentation.django import DjangoInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor
import functools
from drf_mcp.provider import MCPProvider


# Configure Jaeger exporter
jaeger_exporter = JaegerExporter(
    agent_host_name="localhost",
    agent_port=6831,
)

# Set up tracing
trace_provider = TracerProvider()
trace_provider.add_span_processor(BatchSpanProcessor(jaeger_exporter))
trace.set_tracer_provider(trace_provider)

# Set up metrics
metrics_reader = PeriodicExportingMetricReader(
    JaegerMetricsExporter(
        agent_host_name="localhost",
        agent_port=6831,
    )
)
metrics_provider = MeterProvider(metric_readers=[metrics_reader])
metrics.set_meter_provider(metrics_provider)

# Instrument Django
DjangoInstrumentor().instrument()
RequestsInstrumentor().instrument()

# Get tracer
tracer = trace.get_tracer(__name__)
meter = metrics.get_meter(__name__)


# Decorator to trace MCP tool execution
def trace_mcp_tool(func):
    """Decorator to automatically trace MCP tool execution."""
    
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        with tracer.start_as_current_span(f"mcp_tool.{func.__name__}") as span:
            span.set_attribute("mcp.tool.name", func.__name__)
            span.set_attribute("mcp.tool.args", str(args))
            
            try:
                result = await func(*args, **kwargs)
                span.set_attribute("mcp.tool.status", "success")
                return result
            except Exception as e:
                span.set_attribute("mcp.tool.status", "error")
                span.set_attribute("mcp.tool.error", str(e))
                raise
    
    return wrapper


# Example traced tool
@trace_mcp_tool
async def get_user_profile(user_id: int):
    """Get user profile with distributed tracing."""
    from django.contrib.auth.models import User
    
    with tracer.start_as_current_span("db.query.user") as span:
        span.set_attribute("db.user_id", user_id)
        user = await User.objects.aget(id=user_id)
        return {
            'id': user.id,
            'username': user.username,
            'email': user.email
        }


# Initialize MCP Provider with monitoring
provider = MCPProvider()

# Counter for tool invocations
tool_counter = meter.create_counter(
    name="mcp.tools.invocations",
    description="Number of MCP tool invocations",
    unit="1"
)

# Histogram for tool execution time
tool_duration = meter.create_histogram(
    name="mcp.tools.duration",
    description="MCP tool execution duration in milliseconds",
    unit="ms"
)


# Usage configuration
if __name__ == "__main__":
    import logging
    
    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    logger = logging.getLogger(__name__)
    logger.info("Monitoring and observability configured")
    logger.info(f"Jaeger exporter configured for localhost:6831")
    logger.info("Distributed tracing enabled for MCP tools")
    
    print("""
    Monitoring Setup Complete
    ========================
    
    To view traces:
    1. Start Jaeger locally: docker run -d -p 6831:6831/udp jaeger
    2. Open http://localhost:16686 in your browser
    3. Search for service: your-service-name
    
    Metrics are collected for:
    - Tool invocations count
    - Tool execution duration
    - Database query performance
    - HTTP request handling
    """)
