"""
LangGraph instrumentation core.

Single source of truth for all LangGraph span/trace creation logic.
Both the auto-instrumentation handler and the manual handler delegate here.

Design decisions:
- All spans use single SPAN_CREATE with complete data (backend compat).
- datetime.now(timezone.utc) everywhere (no naive datetimes).
- Cost calculation via cost_tracking.calculate_cost().
- Error classification via error_detection.ErrorDetector.
- Session/thread IDs extracted from configurable + env var.
- I/O measured and truncated via shared.io_utils.
"""

import logging
import os
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from ...shared.io_utils import sanitize_and_measure
from ...shared.token_extraction import extract_from_langchain_response

logger = logging.getLogger(__name__)

# Framework internals to skip when encountered as node names
_INTERNAL_NAMES = frozenset({
    "RunnableSequence", "CompiledStateGraph", "Pregel",
    "ChannelWrite", "ChannelRead", "__start__", "__end__",
})


def _normalize_model(name: str) -> str:
    """Strip provider prefixes from model names (e.g. 'models/gemini-2.5-flash' -> 'gemini-2.5-flash')."""
    if name.startswith("models/"):
        return name[7:]
    return name


# ---------------------------------------------------------------------------
# Per-span state dataclasses
# ---------------------------------------------------------------------------

@dataclass
class _TraceState:
    """Per-run state for a root graph execution."""
    trace: Any  # TraceContext from auto_instrument.trace
    workflow_span_id: str
    workflow_name: str
    start_time: datetime
    input_data: Any
    session_id: Optional[str] = None
    thread_id: Optional[str] = None
    turn_count: int = 0
    total_tool_calls: int = 0


@dataclass
class _NodeState:
    span_id: str
    trace_id: str
    parent_id: str
    node_name: str
    start_time: datetime
    input_data: Any
    input_size: int = 0


@dataclass
class _LLMState:
    span_id: str
    trace_id: str
    parent_id: str
    model_name: str
    start_time: datetime
    input_data: Any
    input_size: int = 0
    # Streaming support
    first_token_time: Optional[datetime] = None
    chunk_count: int = 0


@dataclass
class _ToolState:
    span_id: str
    trace_id: str
    parent_id: str
    tool_name: str
    start_time: datetime
    input_data: Any
    input_size: int = 0


@dataclass
class _RetrieverState:
    span_id: str
    trace_id: str
    parent_id: str
    retriever_name: str
    query: str
    start_time: datetime


# ---------------------------------------------------------------------------
# Core
# ---------------------------------------------------------------------------

class LangGraphCore:
    """Shared instrumentation logic for all LangGraph handlers.

    Handles concurrent runs via run_id-keyed state so a single instance
    can serve many parallel graph executions.
    """

    def __init__(self, graph_schema: Any = None):
        self._graph_schema = graph_schema

        # Per-run bookkeeping (keyed by root run_id)
        self._runs: Dict[str, _TraceState] = {}
        # child run_id -> root run_id
        self._run_to_root: Dict[str, str] = {}
        # run_id -> span state
        self._node_spans: Dict[str, _NodeState] = {}
        self._llm_spans: Dict[str, _LLMState] = {}
        self._tool_spans: Dict[str, _ToolState] = {}
        self._retriever_spans: Dict[str, _RetrieverState] = {}

        # Error detection (lazy import to avoid circular)
        self._error_detector = None
        # Drift detection
        self._drift_detector = None

    def _get_error_detector(self):
        if self._error_detector is None:
            try:
                from .error_detection import get_error_detector
                self._error_detector = get_error_detector()
            except ImportError:
                pass
        return self._error_detector

    def _get_drift_detector(self):
        if self._drift_detector is None:
            try:
                from .drift_detection import DriftDetector
                self._drift_detector = DriftDetector()
            except ImportError:
                pass
        return self._drift_detector

    # ------------------------------------------------------------------
    # Trace lifecycle
    # ------------------------------------------------------------------

    def start_trace(
        self,
        run_id: str,
        trace: Any,
        workflow_name: str,
        inputs: Any,
        *,
        session_id: Optional[str] = None,
        thread_id: Optional[str] = None,
    ) -> None:
        """Register a new root graph execution."""
        # Extract session/thread from inputs configurable
        if not thread_id and isinstance(inputs, dict):
            configurable = inputs.get("configurable", {})
            if isinstance(configurable, dict):
                thread_id = configurable.get("thread_id")
        if not session_id:
            session_id = os.environ.get("AIGIE_SESSION_ID")

        # Wire up session context
        try:
            from .session import get_or_create_session_context
            sess = get_or_create_session_context(trace_name=workflow_name)
            sess.increment_graph_run()
            if not session_id:
                session_id = sess.trace_id
        except ImportError:
            pass

        input_data, _ = sanitize_and_measure(inputs) if isinstance(inputs, dict) else ({"input": str(inputs)[:500]}, 0)

        lg_span_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc)

        self._runs[run_id] = _TraceState(
            trace=trace,
            workflow_span_id=lg_span_id,
            workflow_name=workflow_name,
            start_time=now,
            input_data=input_data,
            session_id=session_id,
            thread_id=thread_id,
        )
        self._run_to_root[run_id] = run_id

        # Wire drift detector
        dd = self._get_drift_detector()
        if dd:
            dd.capture_initial_input(inputs)
            if isinstance(inputs, dict):
                system_prompt = (
                    inputs.get("system_prompt")
                    or inputs.get("system_message")
                    or inputs.get("system")
                    or inputs.get("instructions")
                )
                if system_prompt:
                    dd.capture_system_prompt(str(system_prompt))

        logger.debug(f"LangGraph core: trace started for '{workflow_name}'")

    def end_trace(self, run_id: str, outputs: Any) -> List[Dict[str, Any]]:
        """Finalize a root graph execution. Returns batch events for sync delivery."""
        run_data = self._runs.pop(run_id, None)
        if not run_data:
            return []

        end_time = datetime.now(timezone.utc)
        duration = (end_time - run_data.start_time).total_seconds()

        output_data, output_size = sanitize_and_measure(outputs) if isinstance(outputs, dict) else ({"output": str(outputs)[:500]}, 0)

        # Drift finalization
        drift_report = self._finalize_drift(run_data, output=outputs)

        # Workflow span
        workflow_span = {
            "id": run_data.workflow_span_id,
            "trace_id": str(run_data.trace.id),
            "name": "LangGraph",
            "type": "workflow",
            "input": run_data.input_data,
            "output": output_data,
            "status": "success",
            "metadata": self._build_workflow_metadata(run_data, drift_report),
            "start_time": run_data.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

        trace_update = {
            "id": str(run_data.trace.id),
            "name": run_data.workflow_name,
            "status": "success",
            "metadata": self._build_workflow_metadata(run_data, drift_report),
        }

        # Clean up child mappings
        self._cleanup_children(run_id)

        logger.debug(f"LangGraph core: trace completed for '{run_data.workflow_name}'")

        return [
            {"type": "span-create", "id": workflow_span["id"], "timestamp": workflow_span["start_time"], "body": workflow_span},
            {"type": "trace-update", "id": str(run_data.trace.id), "timestamp": end_time.isoformat(), "body": trace_update},
        ]

    def error_trace(self, run_id: str, error: BaseException) -> List[Dict[str, Any]]:
        """Finalize a root graph execution that errored. Returns batch events."""
        run_data = self._runs.pop(run_id, None)
        if not run_data:
            return []

        end_time = datetime.now(timezone.utc)
        duration = (end_time - run_data.start_time).total_seconds()

        # Classify error
        error_meta = self._classify_error(error, f"graph:{run_data.workflow_name}")

        workflow_span = {
            "id": run_data.workflow_span_id,
            "trace_id": str(run_data.trace.id),
            "name": "LangGraph",
            "type": "workflow",
            "input": run_data.input_data,
            "status": "failed",
            "error": str(error),
            "error_message": str(error),
            "metadata": {
                "framework": "langgraph",
                "workflow_name": run_data.workflow_name,
                **({"error_detection": error_meta} if error_meta else {}),
            },
            "start_time": run_data.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

        # Promote error classification to top-level span fields
        if error_meta:
            workflow_span["error_type"] = error_meta.get("type")
            workflow_span["error_severity"] = error_meta.get("severity")
            workflow_span["error_is_transient"] = error_meta.get("is_transient")

        self._cleanup_children(run_id)

        return [
            {"type": "span-create", "id": workflow_span["id"], "timestamp": workflow_span["start_time"], "body": workflow_span},
        ]

    # ------------------------------------------------------------------
    # Node spans
    # ------------------------------------------------------------------

    def start_node(
        self,
        run_id: str,
        parent_run_id: str,
        node_name: str,
        inputs: Any,
    ) -> bool:
        """Start a node span. Returns False if this node should be skipped."""
        if node_name in _INTERNAL_NAMES:
            return False

        # Register in run-to-root mapping
        self._run_to_root[run_id] = self._run_to_root.get(parent_run_id, parent_run_id)
        root_id = self._run_to_root.get(run_id)
        if not root_id or root_id not in self._runs:
            return False

        run_data = self._runs[root_id]
        run_data.turn_count += 1

        input_data, input_size = sanitize_and_measure(inputs) if isinstance(inputs, dict) else ({"input": str(inputs)[:500]}, 0)

        self._node_spans[run_id] = _NodeState(
            span_id=str(uuid.uuid4()),
            trace_id=str(run_data.trace.id),
            parent_id=run_data.workflow_span_id,
            node_name=node_name,
            start_time=datetime.now(timezone.utc),
            input_data=input_data,
            input_size=input_size,
        )

        # Drift detection
        dd = self._get_drift_detector()
        if dd:
            dd.capture_planning_node_output(node_name, inputs)

        # Session context
        try:
            from .session import get_session_context
            sess = get_session_context()
            if sess:
                sess.increment_node_execution(node_name)
        except ImportError:
            pass

        return True

    def end_node(self, run_id: str, outputs: Any) -> Optional[Dict[str, Any]]:
        """End a node span. Returns the span payload or None."""
        span_info = self._node_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()

        output_data, output_size = sanitize_and_measure(outputs) if isinstance(outputs, dict) else ({"output": str(outputs)[:500]}, 0)

        # Drift detection
        dd = self._get_drift_detector()
        if dd:
            dd.record_node_execution(
                span_info.node_name,
                duration_ms=duration * 1000,
                output_state=outputs if isinstance(outputs, dict) else None,
                is_error=False,
            )

        return {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": span_info.node_name,
            "type": "chain",
            "input": span_info.input_data,
            "output": output_data,
            "status": "success",
            "metadata": {
                "nodeType": "langgraph_node",
                "input_size_bytes": span_info.input_size,
                "output_size_bytes": output_size,
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

    def error_node(self, run_id: str, error: BaseException) -> Optional[Dict[str, Any]]:
        """End a node span with error. Returns the span payload or None."""
        span_info = self._node_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()

        error_meta = self._classify_error(error, f"node:{span_info.node_name}")

        # Drift detection
        dd = self._get_drift_detector()
        if dd:
            dd.record_node_execution(
                span_info.node_name,
                duration_ms=duration * 1000,
                is_error=True,
            )

        span = {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": span_info.node_name,
            "type": "chain",
            "input": span_info.input_data,
            "status": "failed",
            "error": str(error),
            "error_message": str(error),
            "metadata": {
                "nodeType": "langgraph_node",
                **({"error_detection": error_meta} if error_meta else {}),
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

        if error_meta:
            span["error_type"] = error_meta.get("type")
            span["error_severity"] = error_meta.get("severity")
            span["error_is_transient"] = error_meta.get("is_transient")

        return span

    # ------------------------------------------------------------------
    # LLM/Generation spans
    # ------------------------------------------------------------------

    def start_llm(
        self,
        run_id: str,
        parent_run_id: Optional[str],
        model_name: str,
        prompts: Any,
    ) -> bool:
        """Start an LLM span. Returns False if skipped."""
        pid = parent_run_id
        self._run_to_root[run_id] = self._run_to_root.get(pid, pid) if pid else None
        root_id = self._run_to_root.get(run_id)
        if not root_id or root_id not in self._runs:
            return False

        run_data = self._runs[root_id]

        # Parent: prefer the node span, fall back to workflow span
        parent_id = run_data.workflow_span_id
        if pid and pid in self._node_spans:
            parent_id = self._node_spans[pid].span_id

        input_data, input_size = sanitize_and_measure({"prompts": prompts[:2] if prompts else []})

        self._llm_spans[run_id] = _LLMState(
            span_id=str(uuid.uuid4()),
            trace_id=str(run_data.trace.id),
            parent_id=str(parent_id),
            model_name=model_name,
            start_time=datetime.now(timezone.utc),
            input_data=input_data,
            input_size=input_size,
        )

        # Session context
        try:
            from .session import get_session_context
            sess = get_session_context()
            if sess:
                sess.increment_llm_calls()
                sess.track_model(model_name)
        except ImportError:
            pass

        return True

    def end_llm(self, run_id: str, response: Any) -> Optional[Dict[str, Any]]:
        """End an LLM span. Returns the span payload or None."""
        span_info = self._llm_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()

        # Canonical token extraction
        usage = extract_from_langchain_response(response)
        model_name = _normalize_model(usage.model or span_info.model_name)

        # Build structured output: response text + token usage
        output_data: Dict[str, Any] = {}
        if usage.response_text:
            output_data["role"] = "assistant"
            output_data["content"] = usage.response_text[:2000]

        token_usage: Dict[str, Any] = {}
        if usage.total_tokens:
            token_usage = {
                "input_tokens": usage.input_tokens,
                "output_tokens": usage.output_tokens,
                "total_tokens": usage.total_tokens,
            }

        # Cost calculation
        cost_data = self._calculate_cost(token_usage, model_name)

        # Streaming TTFT
        streaming_meta = {}
        if span_info.first_token_time:
            ttft_ms = (span_info.first_token_time - span_info.start_time).total_seconds() * 1000
            streaming_meta = {
                "ttft_ms": round(ttft_ms, 2),
                "chunk_count": span_info.chunk_count,
            }

        # Session context - add tokens
        try:
            from .session import get_session_context
            sess = get_session_context()
            if sess:
                sess.add_tokens(usage.input_tokens, usage.output_tokens)
                if cost_data:
                    sess.add_cost(cost_data.get("total_cost", 0))
        except ImportError:
            pass

        latency = duration

        span_data = {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": model_name,
            "type": "llm",
            "model": model_name,
            "input": span_info.input_data,
            "output": sanitize_and_measure(output_data)[0] if output_data else {},
            "status": "success",
            "status_message": "completed",
            "metadata": {
                "provider": "langchain",
                "model": model_name,
                "ls_model_name": model_name,
                **({"token_usage": token_usage} if token_usage else {}),
                **({"cost": cost_data["total_cost"]} if cost_data else {}),
                **({"input_cost": cost_data["input_cost"]} if cost_data else {}),
                **({"output_cost": cost_data["output_cost"]} if cost_data else {}),
                **streaming_meta,
                "input_size_bytes": span_info.input_size,
                "latency_seconds": round(latency, 3),
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

        if token_usage.get("input_tokens"):
            span_data["prompt_tokens"] = token_usage["input_tokens"]
        if token_usage.get("output_tokens"):
            span_data["completion_tokens"] = token_usage["output_tokens"]
        if token_usage.get("total_tokens"):
            span_data["total_tokens"] = token_usage["total_tokens"]
        if cost_data:
            span_data["input_cost"] = cost_data.get("input_cost", 0)
            span_data["output_cost"] = cost_data.get("output_cost", 0)
            span_data["total_cost"] = cost_data.get("total_cost", 0)

        return span_data

    def error_llm(self, run_id: str, error: BaseException) -> Optional[Dict[str, Any]]:
        """End an LLM span with error. Returns the span payload or None."""
        span_info = self._llm_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()
        model_name = _normalize_model(span_info.model_name)

        error_meta = self._classify_error(error, f"llm:{model_name}")

        span = {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": model_name,
            "type": "llm",
            "model": model_name,
            "input": span_info.input_data,
            "status": "failed",
            "error": str(error),
            "error_message": str(error),
            "metadata": {
                "provider": "langchain",
                "model": model_name,
                "ls_model_name": model_name,
                **({"error_detection": error_meta} if error_meta else {}),
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

        if error_meta:
            span["error_type"] = error_meta.get("type")
            span["error_severity"] = error_meta.get("severity")
            span["error_is_transient"] = error_meta.get("is_transient")

        return span

    def llm_new_token(self, run_id: str, token: str) -> None:
        """Track a streaming token for TTFT calculation."""
        span_info = self._llm_spans.get(run_id)
        if not span_info:
            return
        span_info.chunk_count += 1
        if span_info.first_token_time is None:
            span_info.first_token_time = datetime.now(timezone.utc)

    # ------------------------------------------------------------------
    # Tool spans
    # ------------------------------------------------------------------

    def start_tool(
        self,
        run_id: str,
        parent_run_id: Optional[str],
        tool_name: str,
        input_str: Any,
    ) -> bool:
        """Start a tool span. Returns False if skipped."""
        pid = parent_run_id
        self._run_to_root[run_id] = self._run_to_root.get(pid, pid) if pid else None
        root_id = self._run_to_root.get(run_id)
        if not root_id or root_id not in self._runs:
            return False

        run_data = self._runs[root_id]
        run_data.total_tool_calls += 1

        parent_id = run_data.workflow_span_id
        if pid and pid in self._node_spans:
            parent_id = self._node_spans[pid].span_id

        # Try to parse tool input as structured data
        if input_str:
            if isinstance(input_str, dict):
                input_data, input_size = sanitize_and_measure(input_str)
            elif isinstance(input_str, str):
                try:
                    import json as _json
                    input_data, input_size = sanitize_and_measure(_json.loads(input_str))
                except (ValueError, TypeError):
                    input_data, input_size = sanitize_and_measure({"input": input_str[:500]})
            else:
                input_data, input_size = sanitize_and_measure({"input": str(input_str)[:500]})
        else:
            input_data, input_size = {}, 0

        self._tool_spans[run_id] = _ToolState(
            span_id=str(uuid.uuid4()),
            trace_id=str(run_data.trace.id),
            parent_id=str(parent_id),
            tool_name=tool_name,
            start_time=datetime.now(timezone.utc),
            input_data=input_data,
            input_size=input_size,
        )

        # Session context
        try:
            from .session import get_session_context
            sess = get_session_context()
            if sess:
                sess.increment_tool_calls()
        except ImportError:
            pass

        return True

    def end_tool(self, run_id: str, output: Any) -> Optional[Dict[str, Any]]:
        """End a tool span. Returns the span payload or None."""
        span_info = self._tool_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()

        # Extract clean tool output
        if output:
            output_str = getattr(output, "content", None) or str(output)
            output_data, output_size = sanitize_and_measure({"output": str(output_str)[:1000]})
        else:
            output_data, output_size = {}, 0

        # Drift detection
        dd = self._get_drift_detector()
        if dd:
            dd.record_tool_use(
                span_info.tool_name,
                span_info.input_data,
                duration_ms=duration * 1000,
            )

        return {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": str(span_info.tool_name),
            "type": "tool",
            "input": span_info.input_data,
            "output": output_data,
            "status": "success",
            "metadata": {
                "nodeType": "tool",
                "input_size_bytes": span_info.input_size,
                "output_size_bytes": output_size,
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

    def error_tool(self, run_id: str, error: BaseException) -> Optional[Dict[str, Any]]:
        """End a tool span with error. Returns the span payload or None."""
        span_info = self._tool_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()

        error_meta = self._classify_error(error, f"tool:{span_info.tool_name}")

        span = {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": str(span_info.tool_name),
            "type": "tool",
            "input": span_info.input_data,
            "status": "failed",
            "error": str(error),
            "error_message": str(error),
            "metadata": {
                "nodeType": "tool",
                **({"error_detection": error_meta} if error_meta else {}),
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

        if error_meta:
            span["error_type"] = error_meta.get("type")
            span["error_severity"] = error_meta.get("severity")
            span["error_is_transient"] = error_meta.get("is_transient")

        return span

    # ------------------------------------------------------------------
    # Retriever spans (NEW)
    # ------------------------------------------------------------------

    def start_retriever(
        self,
        run_id: str,
        parent_run_id: Optional[str],
        query: str,
        serialized: Dict[str, Any],
    ) -> bool:
        """Start a retriever span. Returns False if skipped."""
        pid = parent_run_id
        self._run_to_root[run_id] = self._run_to_root.get(pid, pid) if pid else None
        root_id = self._run_to_root.get(run_id)
        if not root_id or root_id not in self._runs:
            return False

        run_data = self._runs[root_id]

        parent_id = run_data.workflow_span_id
        if pid and pid in self._node_spans:
            parent_id = self._node_spans[pid].span_id

        retriever_name = serialized.get("name", "Retriever") if serialized else "Retriever"

        self._retriever_spans[run_id] = _RetrieverState(
            span_id=str(uuid.uuid4()),
            trace_id=str(run_data.trace.id),
            parent_id=str(parent_id),
            retriever_name=retriever_name,
            query=str(query)[:500],
            start_time=datetime.now(timezone.utc),
        )
        return True

    def end_retriever(self, run_id: str, documents: Any) -> Optional[Dict[str, Any]]:
        """End a retriever span. Returns the span payload or None."""
        span_info = self._retriever_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()

        # Extract document metadata
        doc_count = 0
        sources: List[str] = []
        total_content_size = 0
        relevance_scores: List[float] = []

        if documents:
            doc_list = documents if isinstance(documents, list) else [documents]
            doc_count = len(doc_list)
            for doc in doc_list[:20]:  # cap at 20
                if hasattr(doc, "metadata") and doc.metadata:
                    src = doc.metadata.get("source") or doc.metadata.get("url")
                    if src:
                        sources.append(str(src)[:200])
                    score = doc.metadata.get("relevance_score") or doc.metadata.get("score")
                    if score is not None:
                        try:
                            relevance_scores.append(float(score))
                        except (TypeError, ValueError):
                            pass
                if hasattr(doc, "page_content"):
                    total_content_size += len(str(doc.page_content))

        output_data = {
            "document_count": doc_count,
            "sources": sources[:10],
            "total_content_size": total_content_size,
        }
        if relevance_scores:
            output_data["relevance_scores"] = relevance_scores[:10]

        return {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": f"Retriever: {span_info.retriever_name}",
            "type": "retriever",
            "input": {"query": span_info.query},
            "output": output_data,
            "status": "success",
            "metadata": {
                "nodeType": "retriever",
                "document_count": doc_count,
                "total_content_size": total_content_size,
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

    def error_retriever(self, run_id: str, error: BaseException) -> Optional[Dict[str, Any]]:
        """End a retriever span with error. Returns the span payload or None."""
        span_info = self._retriever_spans.pop(run_id, None)
        if not span_info:
            return None

        end_time = datetime.now(timezone.utc)
        duration = (end_time - span_info.start_time).total_seconds()

        error_meta = self._classify_error(error, f"retriever:{span_info.retriever_name}")

        span = {
            "id": span_info.span_id,
            "trace_id": span_info.trace_id,
            "parent_id": span_info.parent_id,
            "name": f"Retriever: {span_info.retriever_name}",
            "type": "retriever",
            "input": {"query": span_info.query},
            "status": "failed",
            "error": str(error),
            "error_message": str(error),
            "metadata": {
                "nodeType": "retriever",
                **({"error_detection": error_meta} if error_meta else {}),
            },
            "start_time": span_info.start_time.isoformat(),
            "end_time": end_time.isoformat(),
            "duration_ns": int(duration * 1_000_000_000),
        }

        if error_meta:
            span["error_type"] = error_meta.get("type")
            span["error_severity"] = error_meta.get("severity")
            span["error_is_transient"] = error_meta.get("is_transient")

        return span

    # ------------------------------------------------------------------
    # Helpers (internal)
    # ------------------------------------------------------------------

    def resolve_root(self, run_id: str, parent_run_id: Optional[str]) -> Optional[str]:
        """Register run_id in the run-to-root mapping and return root_id."""
        pid = str(parent_run_id) if parent_run_id else None
        if pid:
            self._run_to_root[run_id] = self._run_to_root.get(pid, pid)
        return self._run_to_root.get(run_id)

    def get_node_span_id(self, run_id: str) -> Optional[str]:
        """Get the span_id for a node span by run_id."""
        ns = self._node_spans.get(run_id)
        return ns.span_id if ns else None

    def _calculate_cost(self, token_usage: Dict[str, Any], model_name: str) -> Optional[Dict[str, Any]]:
        """Calculate cost from token usage."""
        if not token_usage.get("total_tokens"):
            return None
        try:
            from ...cost_tracking import calculate_cost, UsageMetadata
            usage_meta = UsageMetadata(
                input_tokens=token_usage.get("input_tokens", 0),
                output_tokens=token_usage.get("output_tokens", 0),
                total_tokens=token_usage.get("total_tokens", 0),
                model=str(model_name),
            )
            cost_result = calculate_cost(usage_meta)
            if cost_result:
                return {
                    "input_cost": float(cost_result.input_cost),
                    "output_cost": float(cost_result.output_cost),
                    "total_cost": float(cost_result.total_cost),
                }
        except Exception:
            pass
        return None

    def _classify_error(self, error: BaseException, source: str) -> Optional[Dict[str, Any]]:
        """Classify an error using the ErrorDetector."""
        detector = self._get_error_detector()
        if not detector:
            return None
        try:
            detected = detector.detect_from_exception(error, source)
            if detected:
                return {
                    "type": detected.error_type.value,
                    "severity": detected.severity.value,
                    "is_transient": detected.is_transient,
                    "message": detected.message,
                }
        except Exception:
            pass
        return None

    def _finalize_drift(self, run_data: _TraceState, output: Any = None) -> Optional[Dict[str, Any]]:
        """Finalize drift detection and return drift report."""
        dd = self._get_drift_detector()
        if not dd:
            return None
        try:
            drifts = dd.finalize(
                total_duration_ms=0,
                total_tokens=0,
                total_cost=0,
                final_output=str(output)[:500] if output else None,
                final_state=output if isinstance(output, dict) else None,
            )
            return {
                "plan": dd.plan.to_dict(),
                "execution": dd.execution.to_dict(),
                "drifts": [d.to_dict() for d in drifts],
                "drift_count": len(drifts),
            }
        except Exception as e:
            logger.debug(f"Drift finalization error: {e}")
            return None

    def _build_workflow_metadata(self, run_data: _TraceState, drift_report: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Build metadata dict for the workflow span / trace update."""
        meta: Dict[str, Any] = {
            "framework": "langgraph",
            "workflow_name": run_data.workflow_name,
            "turn_count": run_data.turn_count,
            "total_tool_calls": run_data.total_tool_calls,
        }
        if run_data.session_id:
            meta["session_id"] = run_data.session_id
        if run_data.thread_id:
            meta["thread_id"] = run_data.thread_id
        if drift_report:
            meta["drift_report"] = drift_report
            plan_data = drift_report.get("plan")
            if plan_data:
                meta["agent_plan"] = plan_data
        return meta

    def _cleanup_children(self, root_id: str) -> None:
        """Remove all child mappings for a root run_id."""
        stale = [k for k, v in self._run_to_root.items() if v == root_id]
        for k in stale:
            self._run_to_root.pop(k, None)
