"""
Distributed Tracing Support

OpenTelemetry-ready tracing for ALTRUS system.
"""

import time
import uuid
from typing import Dict, Optional, List, Any
from dataclasses import dataclass, field
from enum import Enum


class SpanKind(Enum):
    """Span kind"""
    INTERNAL = "INTERNAL"
    SERVER = "SERVER"
    CLIENT = "CLIENT"


@dataclass
class Span:
    """Trace span"""
    span_id: str
    trace_id: str
    parent_span_id: Optional[str]
    name: str
    kind: SpanKind
    start_time: float
    end_time: Optional[float] = None
    attributes: Dict[str, Any] = field(default_factory=dict)
    events: List[Dict] = field(default_factory=list)
    status: str = "OK"

    def set_attribute(self, key: str, value: Any):
        """Set span attribute"""
        self.attributes[key] = value

    def add_event(self, name: str, **attributes):
        """Add event to span"""
        self.events.append({
            "name": name,
            "timestamp": time.time(),
            "attributes": attributes
        })

    def end(self, status: str = "OK"):
        """End the span"""
        self.end_time = time.time()
        self.status = status

    def duration_ms(self) -> float:
        """Get span duration in milliseconds"""
        if self.end_time:
            return (self.end_time - self.start_time) * 1000
        return 0.0

    def to_dict(self) -> Dict:
        """Convert span to dictionary"""
        return {
            "span_id": self.span_id,
            "trace_id": self.trace_id,
            "parent_span_id": self.parent_span_id,
            "name": self.name,
            "kind": self.kind.value,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": self.duration_ms(),
            "attributes": self.attributes,
            "events": self.events,
            "status": self.status
        }


class Tracer:
    """
    Distributed tracer for ALTRUS.

    OpenTelemetry-compatible interface for future integration.
    """

    def __init__(self, service_name: str = "altrus"):
        self.service_name = service_name
        self._spans: List[Span] = []
        self._active_spans: Dict[str, Span] = {}

    def start_span(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        parent_span_id: Optional[str] = None,
        trace_id: Optional[str] = None,
        **attributes
    ) -> Span:
        """Start a new span"""
        span_id = str(uuid.uuid4())

        if trace_id is None:
            trace_id = str(uuid.uuid4())

        span = Span(
            span_id=span_id,
            trace_id=trace_id,
            parent_span_id=parent_span_id,
            name=name,
            kind=kind,
            start_time=time.time(),
            attributes={
                "service.name": self.service_name,
                **attributes
            }
        )

        self._spans.append(span)
        self._active_spans[span_id] = span

        return span

    def end_span(self, span: Span, status: str = "OK"):
        """End a span"""
        span.end(status)
        if span.span_id in self._active_spans:
            del self._active_spans[span.span_id]

    def get_spans(self, trace_id: Optional[str] = None) -> List[Span]:
        """Get all spans, optionally filtered by trace_id"""
        if trace_id:
            return [s for s in self._spans if s.trace_id == trace_id]
        return self._spans.copy()

    def get_trace(self, trace_id: str) -> List[Dict]:
        """Get complete trace as list of span dictionaries"""
        spans = self.get_spans(trace_id)
        return [s.to_dict() for s in spans]

    def trace_intent_execution(
        self,
        intent_id: str,
        intent_name: str,
        capability: str
    ) -> Span:
        """Start tracing intent execution"""
        return self.start_span(
            f"intent.execute.{intent_name}",
            kind=SpanKind.SERVER,
            intent_id=intent_id,
            intent_name=intent_name,
            capability=capability
        )

    def trace_fault_recovery(
        self,
        module_id: str,
        fault_type: str
    ) -> Span:
        """Start tracing fault recovery"""
        return self.start_span(
            f"fault.recovery.{fault_type}",
            kind=SpanKind.INTERNAL,
            module_id=module_id,
            fault_type=fault_type
        )

    def trace_ledger_write(self, block_index: int) -> Span:
        """Start tracing ledger write"""
        return self.start_span(
            "ledger.write_block",
            kind=SpanKind.INTERNAL,
            block_index=block_index
        )


# Global tracer instance
_global_tracer: Optional[Tracer] = None


def get_tracer(service_name: str = "altrus") -> Tracer:
    """Get or create global tracer"""
    global _global_tracer
    if _global_tracer is None:
        _global_tracer = Tracer(service_name)
    return _global_tracer


# Context manager for automatic span lifecycle
class trace_span:
    """Context manager for tracing a code block"""

    def __init__(
        self,
        name: str,
        kind: SpanKind = SpanKind.INTERNAL,
        **attributes
    ):
        self.name = name
        self.kind = kind
        self.attributes = attributes
        self.span: Optional[Span] = None
        self.tracer = get_tracer()

    def __enter__(self) -> Span:
        self.span = self.tracer.start_span(
            self.name,
            self.kind,
            **self.attributes
        )
        return self.span

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.span:
            status = "ERROR" if exc_type else "OK"
            if exc_type:
                self.span.add_event(
                    "exception",
                    exception_type=str(exc_type),
                    exception_message=str(exc_val)
                )
            self.tracer.end_span(self.span, status)


# Example usage
if __name__ == "__main__":
    tracer = get_tracer()

    # Manual span management
    span = tracer.start_span("example_operation", operation_type="demo")
    span.set_attribute("user_id", "user_123")
    span.add_event("processing_started")
    time.sleep(0.1)
    span.add_event("processing_complete")
    tracer.end_span(span)

    # Context manager
    with trace_span("another_operation", request_id="req_456") as span:
        span.add_event("step_1")
        time.sleep(0.05)
        span.add_event("step_2")

    # Print traces
    for span in tracer.get_spans():
        print(f"Span: {span.name}, Duration: {span.duration_ms():.2f}ms")
