import time
import uuid
import logging
import functools
import asyncio
import json
from contextvars import ContextVar
from typing import Any, Dict, List, Optional, Union, Callable
from dataclasses import dataclass, field, asdict
from datetime import datetime

# Context variable to store the current trace context
_trace_context: ContextVar[Optional['TraceContext']] = ContextVar("_trace_context", default=None)
# Context variable to store a callback for capturing spans
_trace_callback: ContextVar[Optional[Callable[[Dict[str, Any]], None]]] = ContextVar("_trace_callback", default=None)

logger = logging.getLogger("langxchange.tracing")

def set_trace_callback(callback: Optional[Callable[[Dict[str, Any]], None]]):
    """Set a callback to be called whenever a span ends."""
    _trace_callback.set(callback)

@dataclass
class TraceSpan:
    name: str
    trace_id: str
    span_id: str
    parent_id: Optional[str] = None
    start_time: float = field(default_factory=time.time)
    end_time: Optional[float] = None
    status: str = "pending"
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    @property
    def duration_ms(self) -> Optional[float]:
        if self.end_time:
            return (self.end_time - self.start_time) * 1000
        return None

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        d["duration_ms"] = self.duration_ms
        d["start_time_iso"] = datetime.fromtimestamp(self.start_time).isoformat()
        if self.end_time:
            d["end_time_iso"] = datetime.fromtimestamp(self.end_time).isoformat()
        return d

class TraceContext:
    def __init__(self, trace_id: Optional[str] = None):
        self.trace_id = trace_id or str(uuid.uuid4())
        self.spans: List[TraceSpan] = []
        self.active_span_stack: List[str] = []

    def start_span(self, name: str, metadata: Optional[Dict[str, Any]] = None) -> TraceSpan:
        parent_id = self.active_span_stack[-1] if self.active_span_stack else None
        span_id = str(uuid.uuid4())[:8]
        span = TraceSpan(
            name=name,
            trace_id=self.trace_id,
            span_id=span_id,
            parent_id=parent_id,
            metadata=metadata or {}
        )
        self.spans.append(span)
        self.active_span_stack.append(span_id)
        return span

    def end_span(self, span: TraceSpan, status: str = "success", metadata: Optional[Dict[str, Any]] = None):
        span.end_time = time.time()
        span.status = status
        if metadata:
            span.metadata.update(metadata)
        
        if self.active_span_stack and self.active_span_stack[-1] == span.span_id:
            self.active_span_stack.pop()
        
        # Log the completed span as structural JSON
        span_data = span.to_dict()
        logger.info(f"TRACE_SPAN: {json.dumps(span_data)}")
        
        # Call the callback if registered
        callback = _trace_callback.get()
        if callback:
            try:
                callback(span_data)
            except Exception as e:
                logger.error(f"Error in trace callback: {e}")

def get_trace_context() -> Optional[TraceContext]:
    """Get the current trace context from context variables."""
    return _trace_context.get()

def start_trace(trace_id: Optional[str] = None) -> TraceContext:
    """Initialize a new trace context and set it as active."""
    ctx = TraceContext(trace_id=trace_id)
    _trace_context.set(ctx)
    return ctx

class trace_span:
    """Context manager and decorator for tracing spans."""
    def __init__(self, name: str, metadata: Optional[Dict[str, Any]] = None):
        self.name = name
        self.metadata = metadata or {}
        self.span: Optional[TraceSpan] = None
        self.ctx: Optional[TraceContext] = None

    def __enter__(self):
        self.ctx = get_trace_context()
        if self.ctx:
            self.span = self.ctx.start_span(self.name, self.metadata)
        return self.span

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.ctx and self.span:
            status = "success" if exc_type is None else "error"
            if exc_type:
                self.span.metadata["error"] = str(exc_val)
            self.ctx.end_span(self.span, status=status)

    async def __aenter__(self):
        return self.__enter__()

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return self.__exit__(exc_type, exc_val, exc_tb)

    def __call__(self, func: Callable):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Try to use existing context
            ctx = get_trace_context()
            if not ctx:
                return func(*args, **kwargs)
                
            with trace_span(self.name or func.__name__, self.metadata):
                return func(*args, **kwargs)

        @functools.wraps(func)
        async def awrapper(*args, **kwargs):
            ctx = get_trace_context()
            if not ctx:
                return await func(*args, **kwargs)
                
            async with trace_span(self.name or func.__name__, self.metadata):
                return await func(*args, **kwargs)

        return awrapper if asyncio.iscoroutinefunction(func) else wrapper
