"""OpenTelemetry span exporter for AlbusOS pathway execution.

Creates real OTel spans from the ``PathwayEvent`` stream, maintaining
parent/child hierarchy.  Requires the ``otel`` optional dependency group::

    pip install albusos[otel]

Usage::

    from core.observability.otel import OTelSpanExporter

    exporter = OTelSpanExporter()
    vm.add_listener(exporter)
    # ... run pathway ...

The exporter creates spans for:

- Pathway execution (root span)
- Node execution (child of pathway)
- Tool calls (child of node)
- LLM decisions (span events on node)
"""

from __future__ import annotations

import contextlib
import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable

    from core.pathways.models.execution import PathwayEvent

logger = logging.getLogger(__name__)

try:
    from opentelemetry import trace
    from opentelemetry.trace import StatusCode

    _HAS_OTEL = True
except ImportError:
    _HAS_OTEL = False


class OTelSpanExporter:
    """VMEventEmitter listener that creates OpenTelemetry spans.

    Wraps ``opentelemetry.trace.get_tracer("albusos")`` and creates
    spans from ``PathwayEvent`` objects, maintaining the execution
    hierarchy.

    The exporter is resilient: all event handlers are wrapped in
    ``contextlib.suppress`` so telemetry failures never break execution.

    Parameters
    ----------
    service_name
        The OTel ``service.name`` resource attribute.
    tracer_name
        The OTel tracer name.
    """

    def __init__(
        self,
        *,
        service_name: str = "albusos",
        tracer_name: str = "albusos",
    ) -> None:
        if not _HAS_OTEL:
            raise ImportError(
                "opentelemetry-api and opentelemetry-sdk are required. "
                "Install with: pip install albusos[otel]"
            )
        self._tracer = trace.get_tracer(tracer_name)
        self._service_name = service_name

        # Active span tracking: span_id -> OTel span
        self._spans: dict[str, trace.Span] = {}
        # Root span for the execution
        self._root_span: trace.Span | None = None

    # ------------------------------------------------------------------
    # Listener interface
    # ------------------------------------------------------------------

    def __call__(self, event: PathwayEvent) -> None:
        """Dispatch incoming pathway event to the appropriate handler."""
        event_type = getattr(event, "event_type", None)
        handler = _OTEL_DISPATCH.get(event_type)
        if handler:
            with contextlib.suppress(Exception):
                handler(self, event)

    # ------------------------------------------------------------------
    # Per-event handlers
    # ------------------------------------------------------------------

    def _on_execution_started(self, event: PathwayEvent) -> None:
        data = event.data or {}
        execution_id = str(data.get("execution_id", ""))
        pathway_id = event.pathway_id or ""
        span_id = str(data.get("span_id", f"exec:{execution_id}"))

        self._root_span = self._tracer.start_span(
            name=f"pathway:{pathway_id}",
            attributes={
                "service.name": self._service_name,
                "albusos.pathway_id": pathway_id,
                "albusos.execution_id": execution_id,
            },
        )
        self._spans[span_id] = self._root_span

    def _on_execution_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        span_id = str(data.get("span_id", ""))

        span = self._spans.pop(span_id, None) or self._root_span
        if span:
            span.set_status(StatusCode.OK)
            span.end()
            if span is self._root_span:
                self._root_span = None

    def _on_node_started(self, event: PathwayEvent) -> None:
        data = event.data or {}
        node_id = event.node_id or ""
        node_type = str(data.get("node_type", "unknown"))
        span_id = str(data.get("span_id", f"node:{node_id}"))
        parent_span_id = str(data.get("parent_span_id", ""))

        # Try to create as child of parent span
        ctx = None
        parent = self._spans.get(parent_span_id) or self._root_span
        if parent:
            ctx = trace.set_span_in_context(parent)

        span = self._tracer.start_span(
            name=f"node:{node_id}",
            context=ctx,
            attributes={
                "albusos.node_id": node_id,
                "albusos.node_type": node_type,
                "albusos.pathway_id": event.pathway_id or "",
            },
        )
        self._spans[span_id] = span

    def _on_node_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        span_id = str(data.get("span_id", ""))
        duration_ms = data.get("duration_ms")

        span = self._spans.pop(span_id, None)
        if span:
            if isinstance(duration_ms, (int, float)):
                span.set_attribute("albusos.duration_ms", duration_ms)
            span.set_status(StatusCode.OK)
            span.end()

    def _on_node_failed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        span_id = str(data.get("span_id", ""))
        error_msg = event.error or str(data.get("error", ""))

        span = self._spans.pop(span_id, None)
        if span:
            span.set_status(StatusCode.ERROR, error_msg)
            span.record_exception(Exception(error_msg))
            span.end()

    def _on_tool_called(self, event: PathwayEvent) -> None:
        data = event.data or {}
        tool_name = str(data.get("tool_name", "unknown"))
        call_id = str(data.get("call_id", ""))
        node_id = event.node_id or ""
        node_span_id = str(data.get("parent_span_id", f"node:{node_id}"))

        ctx = None
        parent = self._spans.get(node_span_id)
        if parent:
            ctx = trace.set_span_in_context(parent)

        span = self._tracer.start_span(
            name=f"tool:{tool_name}",
            context=ctx,
            attributes={
                "albusos.tool_name": tool_name,
                "albusos.call_id": call_id,
            },
        )
        if call_id:
            self._spans[f"tool:{call_id}"] = span

    def _on_tool_completed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        call_id = str(data.get("call_id", ""))
        span = self._spans.pop(f"tool:{call_id}", None)
        if span:
            span.set_status(StatusCode.OK)
            span.end()

    def _on_tool_failed(self, event: PathwayEvent) -> None:
        data = event.data or {}
        call_id = str(data.get("call_id", ""))
        error_msg = event.error or str(data.get("error", ""))
        span = self._spans.pop(f"tool:{call_id}", None)
        if span:
            span.set_status(StatusCode.ERROR, error_msg)
            span.record_exception(Exception(error_msg))
            span.end()

    def _on_llm_decision(self, event: PathwayEvent) -> None:
        data = event.data or {}
        node_id = event.node_id or ""

        # Find the parent node span and add a span event
        parent_key = next(
            (k for k in self._spans if k.startswith("node:") and k.endswith(f":{node_id}")),
            None,
        )
        parent = self._spans.get(parent_key) if parent_key else None
        if parent:
            parent.add_event(
                "llm_decision",
                attributes={
                    "round": int(data.get("round", 0)),
                    "tools_chosen": str(data.get("tools_chosen", [])),
                    "model": str(data.get("model", "")),
                },
            )

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------

    def finish(self) -> None:
        """End all remaining open spans (call on error/shutdown)."""
        for span in self._spans.values():
            with contextlib.suppress(Exception):
                span.end()
        self._spans.clear()
        if self._root_span:
            with contextlib.suppress(Exception):
                self._root_span.end()
            self._root_span = None


# Dispatch table
_OTEL_DISPATCH: dict[str, Callable[[OTelSpanExporter, PathwayEvent], None]] = {
    "execution_started": OTelSpanExporter._on_execution_started,
    "execution_completed": OTelSpanExporter._on_execution_completed,
    "node_started": OTelSpanExporter._on_node_started,
    "node_completed": OTelSpanExporter._on_node_completed,
    "node_failed": OTelSpanExporter._on_node_failed,
    "tool_called": OTelSpanExporter._on_tool_called,
    "tool_completed": OTelSpanExporter._on_tool_completed,
    "tool_failed": OTelSpanExporter._on_tool_failed,
    "llm_decision": OTelSpanExporter._on_llm_decision,
}


__all__ = [
    "OTelSpanExporter",
]
