"""
Tracer - Trace agent execution with nested spans.

Provides:
- Context manager-based span creation
- Automatic metric aggregation
- Thread-safe trace management
- Integration with storage backends
"""

import asyncio
import logging
import uuid
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime
from typing import Any, Optional

from .models import Span, Trace, SpanKind, SpanStatus

logger = logging.getLogger(__name__)


class Tracer:
    """
    Trace agent execution with nested spans.

    Provides a simple API for creating traces and spans that automatically
    track timing, metrics, and parent-child relationships.

    Inspired by OpenTelemetry and AgentOps patterns.

    Example:
        tracer = Tracer()

        # Start a trace for a task
        trace = tracer.start_trace("task-123", "claude-code")

        # Create spans for operations
        async with tracer.span("generate_code", SpanKind.LLM_CALL) as span:
            result = await llm.generate(prompt)
            span.output_data = {"result": result}
            span.output_tokens = 500

        # End the trace
        tracer.end_trace()
    """

    def __init__(self, storage: Optional["TraceStorage"] = None):
        """
        Initialize the tracer.

        Args:
            storage: Optional storage backend for persisting traces
        """
        self._storage = storage
        self._current_trace: Optional[Trace] = None
        self._span_stack: list[Span] = []
        self._lock = asyncio.Lock()

    @property
    def current_trace(self) -> Optional[Trace]:
        """Get the current active trace."""
        return self._current_trace

    @property
    def current_span(self) -> Optional[Span]:
        """Get the current active span."""
        return self._span_stack[-1] if self._span_stack else None

    def start_trace(
        self,
        task_id: str,
        agent_id: str,
        tags: Optional[dict[str, str]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> Trace:
        """
        Start a new trace for a task.

        Args:
            task_id: Task identifier
            agent_id: Agent identifier
            tags: Optional tags for filtering
            metadata: Optional additional metadata

        Returns:
            The created Trace object
        """
        if self._current_trace:
            logger.warning("Starting new trace while another is active")
            self.end_trace(SpanStatus.CANCELLED)

        trace = Trace(
            trace_id=str(uuid.uuid4()),
            task_id=task_id,
            agent_id=agent_id,
            start_time=datetime.now(),
            tags=tags or {},
            metadata=metadata or {},
        )
        self._current_trace = trace
        self._span_stack = []

        logger.debug(f"Started trace {trace.trace_id} for task {task_id}")
        return trace

    def end_trace(self, status: SpanStatus = SpanStatus.SUCCESS) -> Optional[Trace]:
        """
        End the current trace.

        Args:
            status: Final status of the trace

        Returns:
            The completed trace or None
        """
        if not self._current_trace:
            logger.warning("No active trace to end")
            return None

        # Close any open spans
        while self._span_stack:
            span = self._span_stack.pop()
            if not span.is_complete:
                span.complete(SpanStatus.CANCELLED)

        trace = self._current_trace
        trace.complete(status)

        # Save to storage
        if self._storage:
            try:
                self._storage.save_trace(trace)
            except Exception as e:
                logger.error(f"Failed to save trace: {e}")

        self._current_trace = None
        logger.debug(f"Ended trace {trace.trace_id} with status {status.value}")
        return trace

    @asynccontextmanager
    async def span(
        self,
        name: str,
        kind: SpanKind,
        input_data: Optional[dict[str, Any]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ):
        """
        Create a span context manager (async version).

        Args:
            name: Span name
            kind: Type of span
            input_data: Optional input data
            metadata: Optional metadata

        Yields:
            The created Span object
        """
        span = self._create_span(name, kind, input_data, metadata)

        try:
            yield span
            span.complete(SpanStatus.SUCCESS)
        except Exception as e:
            span.complete(SpanStatus.ERROR, error=str(e))
            raise
        finally:
            self._pop_span()

    @contextmanager
    def span_sync(
        self,
        name: str,
        kind: SpanKind,
        input_data: Optional[dict[str, Any]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ):
        """
        Create a span context manager (sync version).

        Args:
            name: Span name
            kind: Type of span
            input_data: Optional input data
            metadata: Optional metadata

        Yields:
            The created Span object
        """
        span = self._create_span(name, kind, input_data, metadata)

        try:
            yield span
            span.complete(SpanStatus.SUCCESS)
        except Exception as e:
            span.complete(SpanStatus.ERROR, error=str(e))
            raise
        finally:
            self._pop_span()

    def _create_span(
        self,
        name: str,
        kind: SpanKind,
        input_data: Optional[dict[str, Any]] = None,
        metadata: Optional[dict[str, Any]] = None,
    ) -> Span:
        """Create and push a new span."""
        if not self._current_trace:
            raise RuntimeError("No active trace - call start_trace first")

        span = Span(
            span_id=str(uuid.uuid4()),
            trace_id=self._current_trace.trace_id,
            name=name,
            kind=kind,
            start_time=datetime.now(),
            parent_span_id=self._span_stack[-1].span_id if self._span_stack else None,
            input_data=input_data or {},
            metadata=metadata or {},
        )

        # Add to parent or root
        if self._span_stack:
            self._span_stack[-1].add_child(span)
        else:
            self._current_trace.root_span = span

        self._span_stack.append(span)
        return span

    def _pop_span(self) -> Optional[Span]:
        """Pop the current span from the stack."""
        if self._span_stack:
            return self._span_stack.pop()
        return None

    def record_tokens(self, input_tokens: int = 0, output_tokens: int = 0) -> None:
        """Record token usage on the current span."""
        if self.current_span:
            self.current_span.input_tokens += input_tokens
            self.current_span.output_tokens += output_tokens

    def record_cost(self, cost_usd: float) -> None:
        """Record cost on the current span."""
        if self.current_span:
            self.current_span.cost_usd += cost_usd

    def add_metadata(self, key: str, value: Any) -> None:
        """Add metadata to the current span."""
        if self.current_span:
            self.current_span.metadata[key] = value

    def set_error(self, error: str) -> None:
        """Set error on the current span."""
        if self.current_span:
            self.current_span.error = error
            self.current_span.status = SpanStatus.ERROR


class TracerContext:
    """
    Thread-local tracer context for managing active tracer.

    Allows accessing the current tracer from anywhere in the call stack.
    """

    _current: Optional[Tracer] = None

    @classmethod
    def get_current(cls) -> Optional[Tracer]:
        """Get the current tracer."""
        return cls._current

    @classmethod
    def set_current(cls, tracer: Optional[Tracer]) -> None:
        """Set the current tracer."""
        cls._current = tracer


# Convenience functions
def get_tracer() -> Optional[Tracer]:
    """Get the current tracer from context."""
    return TracerContext.get_current()


def start_trace(
    task_id: str,
    agent_id: str,
    tags: Optional[dict[str, str]] = None,
) -> Trace:
    """Start a trace using the current tracer."""
    tracer = get_tracer()
    if not tracer:
        raise RuntimeError("No tracer configured")
    return tracer.start_trace(task_id, agent_id, tags)


def end_trace(status: SpanStatus = SpanStatus.SUCCESS) -> Optional[Trace]:
    """End the current trace."""
    tracer = get_tracer()
    if tracer:
        return tracer.end_trace(status)
    return None


@asynccontextmanager
async def span(
    name: str,
    kind: SpanKind,
    input_data: Optional[dict[str, Any]] = None,
):
    """Create a span using the current tracer."""
    tracer = get_tracer()
    if not tracer:
        # No-op if no tracer configured
        yield None
        return

    async with tracer.span(name, kind, input_data) as s:
        yield s


# Module-level tracer instance
_tracer: Optional[Tracer] = None


def get_global_tracer() -> Tracer:
    """Get or create the global tracer instance."""
    global _tracer
    if _tracer is None:
        _tracer = Tracer()
        TracerContext.set_current(_tracer)
    return _tracer


def set_global_tracer(tracer: Optional[Tracer]) -> None:
    """Set the global tracer instance."""
    global _tracer
    _tracer = tracer
    TracerContext.set_current(tracer)
