"""
Distributed tracing support using OpenTelemetry.

Provides:
- Span creation for operations
- Latency and throughput tracking
- Integration with observability backends
"""

import functools
import time
from contextlib import contextmanager
from typing import Any, Callable, Dict, Optional, TypeVar

from .logging import get_logger

logger = get_logger(__name__)

# Type variable for generic functions
F = TypeVar("F", bound=Callable[..., Any])

# Global tracer instance
_tracer = None
_enabled = False


class Span:
    """Represents a tracing span."""
    
    def __init__(self, name: str, attributes: Optional[Dict[str, Any]] = None):
        self.name = name
        self.attributes = attributes or {}
        self.start_time: Optional[float] = None
        self.end_time: Optional[float] = None
        self._otel_span = None
    
    def __enter__(self) -> "Span":
        self.start_time = time.perf_counter()
        if _enabled and _tracer:
            self._otel_span = _tracer.start_span(self.name)
            for key, value in self.attributes.items():
                self._otel_span.set_attribute(key, value)
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.perf_counter()
        duration_ms = (self.end_time - self.start_time) * 1000
        
        if self._otel_span:
            self._otel_span.set_attribute("duration_ms", duration_ms)
            if exc_type:
                self._otel_span.set_attribute("error", True)
                self._otel_span.set_attribute("error.type", exc_type.__name__)
                self._otel_span.set_attribute("error.message", str(exc_val))
            self._otel_span.end()
        
        logger.debug(
            f"Span '{self.name}' completed",
            additional={
                "duration_ms": duration_ms,
                "attributes": self.attributes,
                "error": exc_type is not None
            }
        )
    
    def set_attribute(self, key: str, value: Any) -> None:
        """Set a span attribute."""
        self.attributes[key] = value
        if self._otel_span:
            self._otel_span.set_attribute(key, value)
    
    def add_event(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> None:
        """Add an event to the span."""
        if self._otel_span:
            self._otel_span.add_event(name, attributes or {})
    
    @property
    def duration_ms(self) -> Optional[float]:
        """Get span duration in milliseconds."""
        if self.start_time and self.end_time:
            return (self.end_time - self.start_time) * 1000
        return None


class Tracer:
    """Main tracer class for creating spans."""
    
    def __init__(
        self,
        service_name: str = "readanybook",
        endpoint: Optional[str] = None,
        enabled: bool = True
    ):
        """
        Initialize the tracer.
        
        Args:
            service_name: Name of the service for tracing
            endpoint: OTLP endpoint URL
            enabled: Whether tracing is enabled
        """
        global _tracer, _enabled
        
        self.service_name = service_name
        self.endpoint = endpoint
        _enabled = enabled
        
        if enabled and endpoint:
            self._setup_opentelemetry()
        else:
            _tracer = None
            logger.info("Tracing disabled or no endpoint configured")
    
    def _setup_opentelemetry(self) -> None:
        """Set up OpenTelemetry tracing."""
        global _tracer
        
        try:
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider
            from opentelemetry.sdk.trace.export import BatchSpanProcessor
            from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
            from opentelemetry.sdk.resources import Resource
            
            resource = Resource.create({
                "service.name": self.service_name,
                "service.version": "0.1.0"
            })
            
            provider = TracerProvider(resource=resource)
            
            if self.endpoint:
                exporter = OTLPSpanExporter(endpoint=self.endpoint)
                processor = BatchSpanProcessor(exporter)
                provider.add_span_processor(processor)
            
            trace.set_tracer_provider(provider)
            _tracer = trace.get_tracer(self.service_name)
            
            logger.info(f"OpenTelemetry tracing initialized for {self.service_name}")
            
        except ImportError:
            logger.warning(
                "OpenTelemetry packages not installed. "
                "Install with: pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp"
            )
            _tracer = None
    
    def start_span(self, name: str, attributes: Optional[Dict[str, Any]] = None) -> Span:
        """Start a new span."""
        return Span(name, attributes)
    
    @contextmanager
    def span(self, name: str, attributes: Optional[Dict[str, Any]] = None):
        """Context manager for creating a span."""
        with Span(name, attributes) as s:
            yield s


def trace_function(
    name: Optional[str] = None,
    attributes: Optional[Dict[str, Any]] = None
) -> Callable[[F], F]:
    """
    Decorator to trace a function.
    
    Args:
        name: Span name (defaults to function name)
        attributes: Additional span attributes
        
    Usage:
        @trace_function()
        def my_function():
            pass
            
        @trace_function(name="custom_name", attributes={"key": "value"})
        def another_function():
            pass
    """
    def decorator(func: F) -> F:
        span_name = name or func.__name__
        span_attrs = attributes or {}
        
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            with Span(span_name, {**span_attrs, "function": func.__name__}):
                return func(*args, **kwargs)
        
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            with Span(span_name, {**span_attrs, "function": func.__name__}):
                return await func(*args, **kwargs)
        
        import asyncio
        if asyncio.iscoroutinefunction(func):
            return async_wrapper  # type: ignore
        return wrapper  # type: ignore
    
    return decorator


class MetricsCollector:
    """Simple metrics collection for latency and throughput."""
    
    def __init__(self):
        self._metrics: Dict[str, list] = {}
    
    def record(self, name: str, value: float) -> None:
        """Record a metric value."""
        if name not in self._metrics:
            self._metrics[name] = []
        self._metrics[name].append(value)
    
    def get_stats(self, name: str) -> Dict[str, float]:
        """Get statistics for a metric."""
        values = self._metrics.get(name, [])
        if not values:
            return {}
        
        import statistics
        return {
            "count": len(values),
            "mean": statistics.mean(values),
            "median": statistics.median(values),
            "min": min(values),
            "max": max(values),
            "stdev": statistics.stdev(values) if len(values) > 1 else 0
        }
    
    def clear(self, name: Optional[str] = None) -> None:
        """Clear metrics."""
        if name:
            self._metrics.pop(name, None)
        else:
            self._metrics.clear()


# Global metrics collector
metrics = MetricsCollector()
