"""
Tracing — structured execution traces for debugging and observability.

Design notes:
  - Built-in lightweight tracer (no dependencies required)
  - Optional OpenTelemetry integration for production use
  - Every agent call, tool call, and handoff creates a Span
  - Spans form a tree mirroring the execution hierarchy
  - Export to JSON, console, or OpenTelemetry collector
"""

from __future__ import annotations

import logging
import time
from collections.abc import AsyncIterator, Iterator
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from typing import Any

from nexus.types import new_id

logger = logging.getLogger("nexus.observability.tracing")


@dataclass
class Span:
    """
    A single span in the execution trace.

    Spans form a tree: each span has a parent (except the root).
    """

    id: str = field(default_factory=new_id)
    parent_id: str | None = None
    trace_id: str = ""
    name: str = ""
    kind: str = "internal"  # "agent", "tool", "llm", "handoff", "internal"

    # Timing
    start_time: float = 0.0
    end_time: float = 0.0

    # Status
    status: str = "ok"  # "ok", "error"
    error_message: str = ""

    # Attributes
    attributes: dict[str, Any] = field(default_factory=dict)

    # Events (timestamped annotations)
    events: list[dict[str, Any]] = field(default_factory=list)

    @property
    def duration_ms(self) -> float:
        if self.end_time > 0 and self.start_time > 0:
            return (self.end_time - self.start_time) * 1000
        return 0.0

    def set_attribute(self, key: str, value: Any) -> None:
        self.attributes[key] = value

    def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None:
        self.events.append({
            "name": name,
            "timestamp": time.time(),
            "attributes": attributes or {},
        })

    def set_error(self, error: str) -> None:
        self.status = "error"
        self.error_message = error

    def finish(self) -> None:
        self.end_time = time.time()

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "parent_id": self.parent_id,
            "trace_id": self.trace_id,
            "name": self.name,
            "kind": self.kind,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "duration_ms": round(self.duration_ms, 2),
            "status": self.status,
            "error_message": self.error_message,
            "attributes": self.attributes,
            "events": self.events,
        }


class Tracer:
    """
    Lightweight execution tracer for Nexus.

    Creates spans for agent calls, tool invocations, LLM calls, and handoffs.
    Spans form a tree that can be exported for debugging.

    Usage:
        tracer = Tracer()

        with tracer.span("agent:researcher", kind="agent") as span:
            span.set_attribute("agent_id", "abc123")

            with tracer.span("llm:gpt-4o", kind="llm", parent=span) as llm_span:
                llm_span.set_attribute("tokens", 1500)

            with tracer.span("tool:search", kind="tool", parent=span) as tool_span:
                tool_span.set_attribute("query", "AI agents 2026")

        # Export
        for span in tracer.spans:
            print(span.to_dict())
    """

    def __init__(self, trace_id: str | None = None) -> None:
        self.trace_id = trace_id or new_id()
        self._spans: list[Span] = []
        self._active_span: Span | None = None
        self._otel_tracer: Any = None  # OpenTelemetry tracer (optional)

    @property
    def spans(self) -> list[Span]:
        return list(self._spans)

    @contextmanager
    def span(
        self,
        name: str,
        kind: str = "internal",
        parent: Span | None = None,
        attributes: dict[str, Any] | None = None,
    ) -> Iterator[Span]:
        """Create a new span as a context manager."""
        s = Span(
            trace_id=self.trace_id,
            name=name,
            kind=kind,
            parent_id=(parent.id if parent else (self._active_span.id if self._active_span else None)),
            start_time=time.time(),
        )
        if attributes:
            s.attributes.update(attributes)

        prev_active = self._active_span
        self._active_span = s

        try:
            yield s
        except Exception as exc:
            s.set_error(str(exc))
            raise
        finally:
            s.finish()
            self._spans.append(s)
            self._active_span = prev_active

    @asynccontextmanager
    async def async_span(
        self,
        name: str,
        kind: str = "internal",
        parent: Span | None = None,
        attributes: dict[str, Any] | None = None,
    ) -> AsyncIterator[Span]:
        """Create a new span as an async context manager."""
        s = Span(
            trace_id=self.trace_id,
            name=name,
            kind=kind,
            parent_id=(parent.id if parent else (self._active_span.id if self._active_span else None)),
            start_time=time.time(),
        )
        if attributes:
            s.attributes.update(attributes)

        prev_active = self._active_span
        self._active_span = s

        try:
            yield s
        except Exception as exc:
            s.set_error(str(exc))
            raise
        finally:
            s.finish()
            self._spans.append(s)
            self._active_span = prev_active

    def try_init_otel(self) -> bool:
        """
        Try to initialize OpenTelemetry integration.

        Returns True if successful, False if OTel is not installed.
        """
        try:
            from opentelemetry import trace
            from opentelemetry.sdk.trace import TracerProvider
            from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor

            provider = TracerProvider()
            provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
            trace.set_tracer_provider(provider)
            self._otel_tracer = trace.get_tracer("nexus")
            logger.info("OpenTelemetry integration initialized")
            return True
        except ImportError:
            logger.debug("OpenTelemetry not installed, using built-in tracer only")
            return False

    def to_tree_string(self) -> str:
        """Generate a human-readable tree representation of the trace."""
        if not self._spans:
            return "(empty trace)"

        # Build parent → children map
        children: dict[str | None, list[Span]] = {}
        for s in self._spans:
            children.setdefault(s.parent_id, []).append(s)

        lines: list[str] = []

        def render(parent_id: str | None, indent: int) -> None:
            for s in children.get(parent_id, []):
                prefix = "  " * indent + ("└─ " if indent > 0 else "")
                status_icon = "✓" if s.status == "ok" else "✗"
                lines.append(
                    f"{prefix}{status_icon} {s.name} ({s.kind}) "
                    f"[{s.duration_ms:.0f}ms]"
                )
                if s.error_message:
                    lines.append(f"{'  ' * (indent + 1)}  ERROR: {s.error_message}")
                render(s.id, indent + 1)

        # Find root spans (no parent)
        root_parent_ids = {s.parent_id for s in self._spans} - {s.id for s in self._spans}
        for root_pid in root_parent_ids:
            render(root_pid, 0)

        return "\n".join(lines)

    def export_json(self) -> list[dict[str, Any]]:
        """Export all spans as JSON-serializable dicts."""
        return [s.to_dict() for s in self._spans]

    def reset(self) -> None:
        """Clear all recorded spans."""
        self._spans.clear()
        self._active_span = None
