"""
Tracing and distributed tracing utilities for InferenceKit.
Handles request tracing and performance monitoring.
"""

import asyncio
import uuid
from typing import Dict, Any, Optional, List, Union
from datetime import datetime, timedelta
from contextvars import ContextVar


class Tracer:
    """Main tracer class for distributed tracing."""

    def __init__(
        self,
        enabled: bool = True,
        sampling_rate: float = 1.0,
        max_trace_age: timedelta = timedelta(hours=1),
        **config
    ):
        self.enabled = enabled
        self.sampling_rate = sampling_rate
        self.max_trace_age = max_trace_age
        self.config = config

        # Trace storage
        self.active_traces: Dict[str, "Trace"] = {}
        self.completed_traces: Dict[str, "Trace"] = {}

        # Context variables
        self.current_trace: ContextVar[Optional["Trace"]] = ContextVar("current_trace", default=None)

    def start_trace(
        self,
        operation: str,
        trace_id: Optional[str] = None,
        parent_id: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> "Trace":
        """Start a new trace."""
        if not self.enabled or not self._should_sample():
            return NoOpTrace()

        trace_id = trace_id or str(uuid.uuid4())
        trace = Trace(
            trace_id=trace_id,
            operation=operation,
            parent_id=parent_id,
            attributes=attributes,
            tracer=self,
            **kwargs
        )

        self.active_traces[trace_id] = trace
        self.current_trace.set(trace)

        return trace

    def get_current_trace(self) -> Optional["Trace"]:
        """Get the current trace from context."""
        return self.current_trace.get()

    def end_trace(self, trace_id: str) -> None:
        """End a trace."""
        if trace_id in self.active_traces:
            trace = self.active_traces.pop(trace_id)
            trace.end()
            self.completed_traces[trace_id] = trace
            self._cleanup_old_traces()

    def get_trace(self, trace_id: str) -> Optional["Trace"]:
        """Get a trace by ID."""
        return self.active_traces.get(trace_id) or self.completed_traces.get(trace_id)

    def get_active_traces(self) -> List["Trace"]:
        """Get all active traces."""
        return list(self.active_traces.values())

    def get_completed_traces(self) -> List["Trace"]:
        """Get all completed traces."""
        return list(self.completed_traces.values())

    def _should_sample(self) -> bool:
        """Determine if a trace should be sampled."""
        return self.sampling_rate >= 1.0 or self.sampling_rate >= random.random()

    def _cleanup_old_traces(self) -> None:
        """Clean up old completed traces."""
        now = datetime.utcnow()
        cutoff_time = now - self.max_trace_age

        old_traces = [
            trace_id for trace_id, trace in self.completed_traces.items()
            if trace.end_time and trace.end_time < cutoff_time
        ]

        for trace_id in old_traces:
            del self.completed_traces[trace_id]

    def get_trace_stats(self) -> Dict[str, Any]:
        """Get trace statistics."""
        return {
            "active_traces": len(self.active_traces),
            "completed_traces": len(self.completed_traces),
            "sampling_rate": self.sampling_rate,
        }

    def shutdown(self) -> None:
        """Shutdown the tracer."""
        # End all active traces
        for trace_id in list(self.active_traces.keys()):
            self.end_trace(trace_id)


class Trace:
    """Trace class representing a single trace."""

    def __init__(
        self,
        trace_id: str,
        operation: str,
        parent_id: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
        tracer: Optional[Tracer] = None,
        **kwargs
    ):
        self.trace_id = trace_id
        self.operation = operation
        self.parent_id = parent_id
        self.attributes = attributes or {}
        self.tracer = tracer
        self.start_time = datetime.utcnow()
        self.end_time: Optional[datetime] = None
        self.duration: Optional[float] = None
        self.events: List[Dict[str, Any]] = []
        self.errors: List[Dict[str, Any]] = []
        self.links: List[Dict[str, Any]] = []
        self.resource_usage: List[Dict[str, Any]] = []

    def log_event(
        self,
        event_name: str,
        attributes: Optional[Dict[str, Any]] = None,
        timestamp: Optional[datetime] = None
    ) -> None:
        """Log an event in the trace."""
        event = {
            "event_name": event_name,
            "timestamp": timestamp or datetime.utcnow(),
            "attributes": attributes or {},
        }
        self.events.append(event)

    def log_error(
        self,
        error_name: str,
        error_message: str,
        error_details: Optional[Dict[str, Any]] = None,
        timestamp: Optional[datetime] = None
    ) -> None:
        """Log an error in the trace."""
        error = {
            "error_name": error_name,
            "error_message": error_message,
            "error_details": error_details or {},
            "timestamp": timestamp or datetime.utcnow(),
        }
        self.errors.append(error)

    def add_link(
        self,
        trace_id: str,
        operation: str,
        attributes: Optional[Dict[str, Any]] = None
    ) -> None:
        """Add a link to another trace."""
        link = {
            "trace_id": trace_id,
            "operation": operation,
            "attributes": attributes or {},
        }
        self.links.append(link)

    def record_resource_usage(
        self,
        resource_type: str,
        usage: float,
        unit: str = "bytes",
        timestamp: Optional[datetime] = None
    ) -> None:
        """Record resource usage."""
        usage_record = {
            "resource_type": resource_type,
            "usage": usage,
            "unit": unit,
            "timestamp": timestamp or datetime.utcnow(),
        }
        self.resource_usage.append(usage_record)

    def end(self) -> None:
        """End the trace."""
        if not self.end_time:
            self.end_time = datetime.utcnow()
            self.duration = (self.end_time - self.start_time).total_seconds()

    def get_summary(self) -> Dict[str, Any]:
        """Get a summary of the trace."""
        return {
            "trace_id": self.trace_id,
            "operation": self.operation,
            "parent_id": self.parent_id,
            "start_time": self.start_time.isoformat(),
            "end_time": self.end_time.isoformat() if self.end_time else None,
            "duration": self.duration,
            "attributes": self.attributes,
            "events_count": len(self.events),
            "errors_count": len(self.errors),
            "links_count": len(self.links),
            "resource_usage_count": len(self.resource_usage),
        }

    def get_events(self) -> List[Dict[str, Any]]:
        """Get all events in the trace."""
        return self.events.copy()

    def get_errors(self) -> List[Dict[str, Any]]:
        """Get all errors in the trace."""
        return self.errors.copy()

    def get_links(self) -> List[Dict[str, Any]]:
        """Get all links in the trace."""
        return self.links.copy()

    def get_resource_usage(self) -> List[Dict[str, Any]]:
        """Get all resource usage records."""
        return self.resource_usage.copy()


class NoOpTrace:
    """No-op trace for when tracing is disabled."""

    def log_event(self, *args, **kwargs) -> None:
        pass

    def log_error(self, *args, **kwargs) -> None:
        pass

    def add_link(self, *args, **kwargs) -> None:
        pass

    def record_resource_usage(self, *args, **kwargs) -> None:
        pass

    def end(self) -> None:
        pass

    def get_summary(self) -> Dict[str, Any]:
        return {}

    def get_events(self) -> List[Dict[str, Any]]:
        return []

    def get_errors(self) -> List[Dict[str, Any]]:
        return []

    def get_links(self) -> List[Dict[str, Any]]:
        return []

    def get_resource_usage(self) -> List[Dict[str, Any]]:
        return []


class TraceContext:
    """Context manager for trace context."""

    def __init__(self, tracer: Tracer, operation: str, **kwargs):
        self.tracer = tracer
        self.operation = operation
        self.kwargs = kwargs
        self.trace: Optional[Trace] = None

    def __enter__(self):
        self.trace = self.tracer.start_trace(self.operation, **self.kwargs)
        return self.trace

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.trace:
            if exc_type:
                self.trace.log_error(
                    error_name=str(exc_type.__name__),
                    error_message=str(exc_val),
                    error_details={}
                )
            self.trace.end()
            if self.tracer:
                self.tracer.end_trace(self.trace.trace_id)


class TraceDecorator:
    """Decorator for tracing functions."""

    def __init__(self, tracer: Tracer, operation: str, **kwargs):
        self.tracer = tracer
        self.operation = operation
        self.kwargs = kwargs

    def __call__(self, func):
        async def wrapper(*args, **kwargs):
            async with TraceContext(self.tracer, self.operation, **self.kwargs):
                return await func(*args, **kwargs)
        return wrapper


class TracePropagation:
    """Utilities for trace propagation across services."""

    @staticmethod
    def inject_trace_context(trace: Trace, carrier: Dict[str, str]) -> None:
        """Inject trace context into a carrier."""
        carrier["trace_id"] = trace.trace_id
        carrier["parent_id"] = trace.parent_id
        carrier["operation"] = trace.operation

    @staticmethod
    def extract_trace_context(carrier: Dict[str, str]) -> Dict[str, Any]:
        """Extract trace context from a carrier."""
        return {
            "trace_id": carrier.get("trace_id"),
            "parent_id": carrier.get("parent_id"),
            "operation": carrier.get("operation"),
        }