"""
Distributed tracing middleware using OpenTelemetry.

Provides:
- Automatic tracing for FastAPI requests
- Span creation and context propagation
- Integration with OTLP exporters
- Custom span attributes and events
""" 
import logging
from typing import Optional
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.resources import Resource, SERVICE_NAME, SERVICE_VERSION
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

logger = logging.getLogger(__name__)

# 全局的 tracer instance
_tracer: Optional[trace.Tracer] = None
_tracer_provider: Optional[TracerProvider] = None


def setup_tracing(
    app,
    service_name: str = "pixel-api",
    service_version: str = "0.1.0",
    otlp_endpoint: str = "http://localhost:4317",
    enabled: bool = True
):
    """Set up distributed tracing for the FastAPI application.
    
    Args:
        app: FastAPI application instance
        service_name: Name of the service
        service_version: Version of the service
        otlp_endpoint: OTLP exporter endpoint
        enabled: Whether tracing is enabled
    """
    global _tracer, _tracer_provider
    
    if not enabled:
        logger.info("Distributed tracing is disabled")
        return
    
    try:
        # 创建 resource with service information
        resource = Resource(attributes={
            SERVICE_NAME: service_name,
            SERVICE_VERSION: service_version
        })
        
        # 创建 tracer provider
        _tracer_provider = TracerProvider(resource=resource)
        
        # 创建 OTLP exporter
        otlp_exporter = OTLPSpanExporter(
            endpoint=otlp_endpoint,
            insecure=True  # Use insecure for local development
        )
        
        # 创建 span processor
        span_processor = BatchSpanProcessor(otlp_exporter)
        _tracer_provider.add_span_processor(span_processor)
        
        # Set global tracer provider
        trace.set_tracer_provider(_tracer_provider)
        
        # 获取 tracer instance
        _tracer = trace.get_tracer(__name__)
        
        # Instrument FastAPI
        FastAPIInstrumentor.instrument_app(app)
        
        logger.info(f"Distributed tracing enabled: {otlp_endpoint}")
        
    except Exception as e:
        logger.error(f"Failed to setup distributed tracing: {e}")
        logger.info("Application will continue without tracing")


def get_tracer() -> trace.Tracer:
    """Get the global tracer instance.
    
    Returns:
        Tracer instance, or a no-op tracer if tracing is not enabled
    """
    global _tracer
    if _tracer is None:
        # 返回 no-op tracer if tracing is not setup
        return trace.get_tracer(__name__)
    return _tracer


def create_span(
    name: str,
    attributes: dict = None,
    kind: trace.SpanKind = trace.SpanKind.INTERNAL
):
    """Create a new span for tracing.
    
    Args:
        name: Name of the span
        attributes: Optional attributes to add to the span
        kind: Kind of span (INTERNAL, SERVER, CLIENT, PRODUCER, CONSUMER)
        
    Returns:
        Span context manager
        
    Example:
        with create_span("generate_image", {"model": "flux-1.1-pro"}):
            # ... operation ...
    """
    tracer = get_tracer()
    span = tracer.start_as_current_span(name, kind=kind)
    
    if attributes:
        for key, value in attributes.items():
            span.set_attribute(key, value)
    
    return span


def add_span_event(name: str, attributes: dict = None):
    """
    Add an event to the current span.
    
    Args:
        name: Name of the event
        attributes: Optional attributes for the event
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        current_span.add_event(name, attributes=attributes or {})


def set_span_attribute(key: str, value):
    """ 集合 an attribute on the current span.
    
    Args:
        key: Attribute key
        value: Attribute value
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        current_span.set_attribute(key, value)


def set_span_status(status_code: trace.StatusCode, description: str = None):
    """ 集合 the status of the current span.
    
    Args:
        status_code: Status code (OK, ERROR, UNSET)
        description: Optional status description
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        from opentelemetry.trace import Status
        current_span.set_status(Status(status_code, description))


def record_exception(exception: Exception):
    """ 记录 an exception in the current span.
    
    Args:
        exception: Exception to record
    """
    current_span = trace.get_current_span()
    if current_span and current_span.is_recording():
        current_span.record_exception(exception)
        set_span_status(trace.StatusCode.ERROR, str(exception))


# 装饰器 for tracing functions
def traced(span_name: str = None, attributes: dict = None):
    """ 装饰器 to automatically trace a function.
    
    Args:
        span_name: Optional custom span name (defaults to function name)
        attributes: Optional attributes to add to the span
        
    Example:
        @traced(span_name="generate_image", attributes={"service": "image"})
        async def generate_image(prompt: str, model: str):
            # ... operation ...
    """
    def decorator(func):
        async def async_wrapper(*args, **kwargs):
            name = span_name or func.__name__
            with create_span(name, attributes):
                try:
                    result = await func(*args, **kwargs)
                    set_span_status(trace.StatusCode.OK)
                    return result
                except Exception as e:
                    record_exception(e)
                    raise
        
        def sync_wrapper(*args, **kwargs):
            name = span_name or func.__name__
            with create_span(name, attributes):
                try:
                    result = func(*args, **kwargs)
                    set_span_status(trace.StatusCode.OK)
                    return result
                except Exception as e:
                    record_exception(e)
                    raise
        
        # 返回 appropriate wrapper based on function type
        import asyncio
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper
    
    return decorator


def shutdown_tracing():
    """Shutdown tracing and flush remaining spans.""" 
    global _tracer_provider
    if _tracer_provider:
        _tracer_provider.shutdown()
        logger.info("Distributed tracing shutdown")
