"""
LangChain integration for Visibe.

Uses LangChain's callback system to capture metrics.
"""
import contextvars
import logging
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Union
import uuid

from .base import BaseIntegration
from ..utils import calculate_cost
from .base import TraceSummary

logger = logging.getLogger("visibe")

# Context variable for propagating the active callback through nested graph invocations.
# Set by the root instrumented invoke, read by the Pregel class-level patch.
_active_callback_ctx: contextvars.ContextVar[Optional['LangGraphCallback']] = \
    contextvars.ContextVar('_active_callback_ctx', default=None)

# Module-level ref visible across all threads in the process.
# context.run() isolates contextvars and threading.local() is per-thread,
# but LangGraph executes tool functions in parallel worker threads.
# A plain module attribute is the only mechanism that survives both.
_active_root_callback: Optional['LangGraphCallback'] = None

# ContextVar for propagating the tool ancestry stack across nested graphs.
# Uses IMMUTABLE tuples so each context.run() / copy_context() gets an
# independent snapshot.  Both LangGraph (copy_context inside set_config_context)
# and LangChain (ContextThreadPoolExecutor) copy the current context, so child
# graphs and worker threads inherit the parent's tool stack automatically.
_tool_stack_ctx: contextvars.ContextVar[tuple] = \
    contextvars.ContextVar('_visibe_tool_stack', default=())

# Module-level state for Pregel class-level patching
_pregel_patched = False
_original_pregel_invoke = None
_original_pregel_ainvoke = None
_original_pregel_astream_events = None
_original_pregel_batch = None
_original_pregel_abatch = None
_original_pregel_stream = None
_original_pregel_astream = None
_instrumented_runnables: Dict[int, Any] = {}  # id(runnable) → {'invoke': original, 'ainvoke': original, 'name': str}

# Check if LangChain is available
try:
    from langchain_core.callbacks import BaseCallbackHandler
    from langchain_core.outputs import LLMResult
    from langchain_core.agents import AgentAction, AgentFinish
    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    BaseCallbackHandler = object
    LLMResult = object
    AgentAction = object
    AgentFinish = object

# Try to import OpenAI callback for token tracking (works with streaming)
try:
    from langchain_community.callbacks import get_openai_callback
    OPENAI_CALLBACK_AVAILABLE = True
except ImportError:
    try:
        from langchain.callbacks import get_openai_callback
        OPENAI_CALLBACK_AVAILABLE = True
    except ImportError:
        OPENAI_CALLBACK_AVAILABLE = False
        get_openai_callback = None

# Try to import Bedrock Anthropic callback (for AWS Bedrock users)
try:
    from langchain_community.callbacks import get_bedrock_anthropic_callback
    BEDROCK_CALLBACK_AVAILABLE = True
except ImportError:
    BEDROCK_CALLBACK_AVAILABLE = False
    get_bedrock_anthropic_callback = None

# Try to import tiktoken for token estimation (fallback for non-OpenAI providers)
try:
    import tiktoken
    TIKTOKEN_AVAILABLE = True
except ImportError:
    TIKTOKEN_AVAILABLE = False
    tiktoken = None


def estimate_tokens(text: str, model: str = "gpt-4") -> int:
    """
    Estimate token count using tiktoken.
    Falls back to simple word-based estimate if tiktoken unavailable.
    """
    if not text:
        return 0

    if TIKTOKEN_AVAILABLE:
        try:
            # Try to get encoding for the model
            try:
                encoding = tiktoken.encoding_for_model(model)
            except KeyError:
                # Fall back to cl100k_base (used by GPT-4, GPT-3.5-turbo)
                encoding = tiktoken.get_encoding("cl100k_base")
            return len(encoding.encode(text))
        except Exception:
            pass

    # Simple fallback: ~4 chars per token (rough estimate)
    return len(text) // 4


def _merge_callback_into_config(config, callback):
    """Merge a callback handler into a LangChain config dict, preserving existing callbacks."""
    if config is None:
        config = {}
    else:
        config = dict(config)  # shallow copy to avoid mutating caller's dict

    existing = config.get("callbacks")

    if existing is None:
        # No callbacks yet — just set ours
        config["callbacks"] = [callback]
        return config

    # existing may be a list or a CallbackManager; normalise to a list of handlers
    if isinstance(existing, list):
        handler_list = existing
    elif hasattr(existing, 'handlers'):
        # CallbackManager — pull its handler list
        handler_list = existing.handlers
    else:
        # Unknown type — wrap in a fresh list to be safe
        handler_list = []

    # Don't add duplicate
    if callback not in handler_list:
        config["callbacks"] = list(handler_list) + [callback]
    else:
        # Ensure it's a plain list so downstream code doesn't hit the same issue
        config["callbacks"] = list(handler_list)

    return config


def _patch_pregel_class():
    """One-time class-level patch on Pregel.invoke and Pregel.ainvoke.

    When a trace is active (context var set by the root instrumented invoke),
    nested graph invocations automatically get the callback injected into their
    config — so developers never need to instrument inner graphs.

    This is a no-op when no trace is active, making it safe globally.
    """
    global _pregel_patched, _original_pregel_invoke, _original_pregel_ainvoke
    global _original_pregel_astream_events
    global _original_pregel_batch, _original_pregel_abatch
    global _original_pregel_stream, _original_pregel_astream

    if _pregel_patched:
        return

    try:
        from langgraph.pregel import Pregel
    except ImportError:
        logger.warning(
            "langgraph is not installed — nested graph auto-propagation disabled. "
            "Install with: pip install langgraph"
        )
        return

    _original_pregel_invoke = Pregel.invoke

    def _patched_pregel_invoke(self, input, config=None, **kwargs):
        cb = _active_root_callback or _active_callback_ctx.get()
        # Only inject for nested (non-root) graphs when a trace is active
        if cb is not None and id(self) not in _instrumented_runnables:
            config = _merge_callback_into_config(config, cb)
        return _original_pregel_invoke(self, input, config=config, **kwargs)

    Pregel.invoke = _patched_pregel_invoke

    # Also patch ainvoke for async support
    _original_pregel_ainvoke = Pregel.ainvoke

    async def _patched_pregel_ainvoke(self, input, config=None, **kwargs):
        cb = _active_root_callback or _active_callback_ctx.get()
        if cb is not None and id(self) not in _instrumented_runnables:
            config = _merge_callback_into_config(config, cb)
        return await _original_pregel_ainvoke(self, input, config=config, **kwargs)

    Pregel.ainvoke = _patched_pregel_ainvoke

    # Patch astream_events if available
    if hasattr(Pregel, 'astream_events'):
        _original_pregel_astream_events = Pregel.astream_events

        async def _patched_pregel_astream_events(self, input, *, config=None, **kwargs):
            cb = _active_root_callback or _active_callback_ctx.get()
            if cb is not None and id(self) not in _instrumented_runnables:
                config = _merge_callback_into_config(config, cb)
            async for event in _original_pregel_astream_events(self, input, config=config, **kwargs):
                yield event

        Pregel.astream_events = _patched_pregel_astream_events

    # Patch batch if available
    if hasattr(Pregel, 'batch'):
        _original_pregel_batch = Pregel.batch

        def _patched_pregel_batch(self, inputs, config=None, **kwargs):
            cb = _active_root_callback or _active_callback_ctx.get()
            if cb is not None and id(self) not in _instrumented_runnables:
                if config is None:
                    config = [_merge_callback_into_config(None, cb) for _ in inputs]
                elif isinstance(config, list):
                    config = [_merge_callback_into_config(c, cb) for c in config]
                else:
                    config = [_merge_callback_into_config(config, cb) for _ in inputs]
            return _original_pregel_batch(self, inputs, config=config, **kwargs)

        Pregel.batch = _patched_pregel_batch

    # Patch abatch if available
    if hasattr(Pregel, 'abatch'):
        _original_pregel_abatch = Pregel.abatch

        async def _patched_pregel_abatch(self, inputs, config=None, **kwargs):
            cb = _active_root_callback or _active_callback_ctx.get()
            if cb is not None and id(self) not in _instrumented_runnables:
                if config is None:
                    config = [_merge_callback_into_config(None, cb) for _ in inputs]
                elif isinstance(config, list):
                    config = [_merge_callback_into_config(c, cb) for c in config]
                else:
                    config = [_merge_callback_into_config(config, cb) for _ in inputs]
            return await _original_pregel_abatch(self, inputs, config=config, **kwargs)

        Pregel.abatch = _patched_pregel_abatch

    # Patch stream if available
    if hasattr(Pregel, 'stream'):
        _original_pregel_stream = Pregel.stream

        def _patched_pregel_stream(self, input, config=None, **kwargs):
            cb = _active_root_callback or _active_callback_ctx.get()
            if cb is not None and id(self) not in _instrumented_runnables:
                config = _merge_callback_into_config(config, cb)
            return _original_pregel_stream(self, input, config=config, **kwargs)

        Pregel.stream = _patched_pregel_stream

    # Patch astream if available
    if hasattr(Pregel, 'astream'):
        _original_pregel_astream = Pregel.astream

        async def _patched_pregel_astream(self, input, config=None, **kwargs):
            cb = _active_root_callback or _active_callback_ctx.get()
            if cb is not None and id(self) not in _instrumented_runnables:
                config = _merge_callback_into_config(config, cb)
            async for chunk in _original_pregel_astream(self, input, config=config, **kwargs):
                yield chunk

        Pregel.astream = _patched_pregel_astream

    _pregel_patched = True
    logger.debug("Pregel.invoke/ainvoke patched for automatic callback propagation")


class LangChainIntegration(BaseIntegration):
    """
    LangChain integration for Visibe observability.

    Usage:
        from visibe import Visibe
        from visibe.integrations.langchain import LangChainIntegration

        obs = Visibe(api_url="http://localhost:3000")
        tracker = LangChainIntegration(obs)

        # Use the callback with your agent
        agent = AgentExecutor(agent=..., tools=..., callbacks=[tracker.callback])

        with tracker.track(name="my-workflow"):
            result = agent.invoke({"input": "..."})
    """

    def __init__(self, client):
        """
        Initialize LangChain integration.

        Args:
            client: Visibe client instance
        """
        if not LANGCHAIN_AVAILABLE:
            raise ImportError(
                "LangChain integration requires langchain-core.\n"
                "Install with: pip install langchain-core\n"
            )
        super().__init__(client)
        self._framework = "langchain"
        self._content_limit = None  # use defaults
        self._callback = LangGraphCallback()

    @property
    def callback(self) -> 'LangGraphCallback':
        """Get the callback handler for use with LangChain or LangGraph.

        Available immediately after creating the integration — no need
        to wait for track() to be entered.
        """
        return self._callback

    def track(self, name: str = "langchain-execution",
              content_limit: Optional[int] = None):
        """
        Track a LangChain execution.

        Args:
            name: Human-readable name for this workflow
            content_limit: Max chars for LLM/tool input/output in trace steps.
                          None = use defaults (1000 LLM, 500 tools). 0 = omit content.

        Returns:
            LangChainTracker context manager
        """
        return LangChainTracker(self, name, content_limit=content_limit)

    # ==================== instrument() mode ====================

    def instrument(self, runnable, name: Optional[str] = None,
                   content_limit: Optional[int] = None):
        """Patch a LangChain/LangGraph runnable so every invoke()/ainvoke() sends a trace.

        Also patches the Pregel class so nested sub-graphs automatically
        participate in the same trace without any extra code.

        Args:
            runnable: A compiled LangGraph graph (or any LangChain Runnable)
            name: Optional name for traces (e.g. "trip-planner"). Shown in dashboard.
            content_limit: Max chars for LLM/tool input/output in trace steps.
                          None = defaults (1000 LLM, 500 tools). 0 = omit content.

        Example:
            >>> obs = Visibe()
            >>> obs.instrument(graph, name="trip-planner")
            >>> obs.instrument(graph, name="trip-planner", content_limit=200)
        """
        self._content_limit = content_limit
        runnable_id = id(runnable)
        if runnable_id in _instrumented_runnables:
            if not name:
                logger.debug("Runnable already instrumented — skipping")
                return
            # If explicit name provided, allow re-patching (name will take effect on next invocation)
            logger.debug(f"Runnable already instrumented, updating trace name to: {name}")
            return

        # Ensure class-level Pregel patch is applied (idempotent)
        _patch_pregel_class()

        # Save originals
        original_invoke = runnable.invoke
        original_ainvoke = getattr(runnable, 'ainvoke', None)
        _instrumented_runnables[runnable_id] = {
            'invoke': original_invoke,
            'ainvoke': original_ainvoke,
            'name': name,
        }

        integration = self  # capture for closure
        trace_name = name or f"{integration._framework}-execution"
        _content_limit = content_limit  # capture for closures

        def wrapped_invoke(input, config=None, **kwargs):
            global _active_root_callback
            # If a parent trace is already active, this is a nested sub-graph.
            # Check module-level ref (survives context.run + worker threads).
            parent_cb = _active_root_callback
            if parent_cb is not None:
                config = _merge_callback_into_config(config, parent_cb)
                return original_invoke(input, config=config, **kwargs)

            callback = LangGraphCallback(content_limit=_content_limit)

            # Create trace header
            integration.client.create_trace({
                'trace_id': callback.trace_id,
                'name': trace_name,
                'framework': integration._framework,
                'started_at': datetime.now(timezone.utc).isoformat(),
            })

            # Wire up span streaming
            callback._span_sender = lambda span: integration.client.queue_span(
                callback.trace_id, span
            )

            _active_root_callback = callback
            token = _active_callback_ctx.set(callback)
            config = _merge_callback_into_config(config, callback)

            try:
                result = original_invoke(input, config=config, **kwargs)
                integration._complete_trace(callback, 'completed', trace_name)
                return result
            except Exception as exc:
                callback.status = "failed"
                callback.errors.append({
                    'timestamp': datetime.now(timezone.utc).isoformat(),
                    'error_type': type(exc).__name__,
                    'message': str(exc),
                    'agent_name': callback.current_agent or 'unknown',
                })
                integration._complete_trace(callback, 'failed', trace_name)
                raise
            finally:
                _active_root_callback = None
                _active_callback_ctx.reset(token)

        runnable.invoke = wrapped_invoke

        # Patch ainvoke for async support
        if original_ainvoke is not None:
            async def wrapped_ainvoke(input, config=None, **kwargs):
                global _active_root_callback
                # If a parent trace is already active, propagate it.
                parent_cb = _active_root_callback
                if parent_cb is not None:
                    config = _merge_callback_into_config(config, parent_cb)
                    return await original_ainvoke(input, config=config, **kwargs)

                callback = LangGraphCallback(content_limit=_content_limit)

                integration.client.create_trace({
                    'trace_id': callback.trace_id,
                    'name': trace_name,
                    'framework': integration._framework,
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                callback._span_sender = lambda span: integration.client.queue_span(
                    callback.trace_id, span
                )

                _active_root_callback = callback
                token = _active_callback_ctx.set(callback)
                config = _merge_callback_into_config(config, callback)

                try:
                    result = await original_ainvoke(input, config=config, **kwargs)
                    integration._complete_trace(callback, 'completed', trace_name)
                    return result
                except Exception as exc:
                    callback.status = "failed"
                    callback.errors.append({
                        'timestamp': datetime.now(timezone.utc).isoformat(),
                        'error_type': type(exc).__name__,
                        'message': str(exc),
                        'agent_name': callback.current_agent or 'unknown',
                    })
                    integration._complete_trace(callback, 'failed', trace_name)
                    raise
                finally:
                    _active_root_callback = None
                    _active_callback_ctx.reset(token)

            runnable.ainvoke = wrapped_ainvoke

        # Patch astream_events for streaming event lifecycles
        original_astream_events = getattr(runnable, 'astream_events', None)
        _instrumented_runnables[runnable_id]['astream_events'] = original_astream_events

        if original_astream_events is not None:
            async def wrapped_astream_events(input, *, config=None, **kwargs):
                # If a parent trace is already active, propagate instead of creating a new trace.
                parent_cb = _active_root_callback
                if parent_cb is not None:
                    config = _merge_callback_into_config(config, parent_cb)
                    async for event in original_astream_events(input, config=config, **kwargs):
                        yield event
                    return

                callback = LangGraphCallback(content_limit=_content_limit)

                integration.client.create_trace({
                    'trace_id': callback.trace_id,
                    'name': trace_name,
                    'framework': integration._framework,
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                callback._span_sender = lambda span: integration.client.queue_span(
                    callback.trace_id, span
                )

                token = _active_callback_ctx.set(callback)
                config = _merge_callback_into_config(config, callback)

                try:
                    async for event in original_astream_events(
                        input, config=config, **kwargs
                    ):
                        yield event
                    integration._complete_trace(callback, 'completed', trace_name)
                except Exception as exc:
                    callback.status = "failed"
                    callback.errors.append({
                        'timestamp': datetime.now(timezone.utc).isoformat(),
                        'error_type': type(exc).__name__,
                        'message': str(exc),
                        'agent_name': callback.current_agent or 'unknown',
                    })
                    integration._complete_trace(callback, 'failed', trace_name)
                    raise
                finally:
                    _active_callback_ctx.reset(token)

            runnable.astream_events = wrapped_astream_events

        # Patch batch for grouped execution
        original_batch = getattr(runnable, 'batch', None)
        _instrumented_runnables[runnable_id]['batch'] = original_batch

        if original_batch is not None:
            def wrapped_batch(inputs, config=None, **kwargs):
                # If a parent trace is already active, propagate instead of creating a new trace.
                parent_cb = _active_root_callback
                if parent_cb is not None:
                    if config is None:
                        configs = [_merge_callback_into_config(None, parent_cb) for _ in inputs]
                    elif isinstance(config, list):
                        configs = [_merge_callback_into_config(c, parent_cb) for c in config]
                    else:
                        configs = [_merge_callback_into_config(config, parent_cb) for _ in inputs]
                    return original_batch(inputs, config=configs, **kwargs)

                callback = LangGraphCallback(content_limit=_content_limit)

                integration.client.create_trace({
                    'trace_id': callback.trace_id,
                    'name': trace_name,
                    'framework': integration._framework,
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                callback._span_sender = lambda span: integration.client.queue_span(
                    callback.trace_id, span
                )

                token = _active_callback_ctx.set(callback)

                if config is None:
                    configs = [_merge_callback_into_config(None, callback) for _ in inputs]
                elif isinstance(config, list):
                    configs = [_merge_callback_into_config(c, callback) for c in config]
                else:
                    configs = [_merge_callback_into_config(config, callback) for _ in inputs]

                try:
                    results = original_batch(inputs, config=configs, **kwargs)
                    integration._complete_trace(callback, 'completed', trace_name)
                    return results
                except Exception as exc:
                    callback.status = "failed"
                    callback.errors.append({
                        'timestamp': datetime.now(timezone.utc).isoformat(),
                        'error_type': type(exc).__name__,
                        'message': str(exc),
                        'agent_name': callback.current_agent or 'unknown',
                    })
                    integration._complete_trace(callback, 'failed', trace_name)
                    raise
                finally:
                    _active_callback_ctx.reset(token)

            runnable.batch = wrapped_batch

        # Patch abatch for grouped async execution
        original_abatch = getattr(runnable, 'abatch', None)
        _instrumented_runnables[runnable_id]['abatch'] = original_abatch

        if original_abatch is not None:
            async def wrapped_abatch(inputs, config=None, **kwargs):
                # If a parent trace is already active, propagate instead of creating a new trace.
                parent_cb = _active_root_callback
                if parent_cb is not None:
                    if config is None:
                        configs = [_merge_callback_into_config(None, parent_cb) for _ in inputs]
                    elif isinstance(config, list):
                        configs = [_merge_callback_into_config(c, parent_cb) for c in config]
                    else:
                        configs = [_merge_callback_into_config(config, parent_cb) for _ in inputs]
                    return await original_abatch(inputs, config=configs, **kwargs)

                callback = LangGraphCallback(content_limit=_content_limit)

                integration.client.create_trace({
                    'trace_id': callback.trace_id,
                    'name': trace_name,
                    'framework': integration._framework,
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                callback._span_sender = lambda span: integration.client.queue_span(
                    callback.trace_id, span
                )

                token = _active_callback_ctx.set(callback)

                if config is None:
                    configs = [_merge_callback_into_config(None, callback) for _ in inputs]
                elif isinstance(config, list):
                    configs = [_merge_callback_into_config(c, callback) for c in config]
                else:
                    configs = [_merge_callback_into_config(config, callback) for _ in inputs]

                try:
                    results = await original_abatch(inputs, config=configs, **kwargs)
                    integration._complete_trace(callback, 'completed', trace_name)
                    return results
                except Exception as exc:
                    callback.status = "failed"
                    callback.errors.append({
                        'timestamp': datetime.now(timezone.utc).isoformat(),
                        'error_type': type(exc).__name__,
                        'message': str(exc),
                        'agent_name': callback.current_agent or 'unknown',
                    })
                    integration._complete_trace(callback, 'failed', trace_name)
                    raise
                finally:
                    _active_callback_ctx.reset(token)

            runnable.abatch = wrapped_abatch

        # Patch stream for streaming output
        original_stream = getattr(runnable, 'stream', None)
        _instrumented_runnables[runnable_id]['stream'] = original_stream

        if original_stream is not None:
            def wrapped_stream(input, config=None, **kwargs):
                # If a parent trace is already active, propagate instead of creating a new trace.
                parent_cb = _active_root_callback
                if parent_cb is not None:
                    config = _merge_callback_into_config(config, parent_cb)
                    return original_stream(input, config=config, **kwargs)

                callback = LangGraphCallback(content_limit=_content_limit)

                integration.client.create_trace({
                    'trace_id': callback.trace_id,
                    'name': trace_name,
                    'framework': integration._framework,
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                callback._span_sender = lambda span: integration.client.queue_span(
                    callback.trace_id, span
                )

                token = _active_callback_ctx.set(callback)
                config = _merge_callback_into_config(config, callback)

                try:
                    stream_iter = original_stream(input, config=config, **kwargs)
                    # Wrap the iterator to finalize trace when exhausted
                    return _StreamIteratorWrapper(
                        stream_iter, integration, callback, trace_name, token
                    )
                except Exception as exc:
                    callback.status = "failed"
                    callback.errors.append({
                        'timestamp': datetime.now(timezone.utc).isoformat(),
                        'error_type': type(exc).__name__,
                        'message': str(exc),
                        'agent_name': callback.current_agent or 'unknown',
                    })
                    integration._complete_trace(callback, 'failed', trace_name)
                    _active_callback_ctx.reset(token)
                    raise

            runnable.stream = wrapped_stream

        # Patch astream for async streaming output
        original_astream = getattr(runnable, 'astream', None)
        _instrumented_runnables[runnable_id]['astream'] = original_astream

        if original_astream is not None:
            async def wrapped_astream(input, config=None, **kwargs):
                # If a parent trace is already active, propagate instead of creating a new trace.
                parent_cb = _active_root_callback
                if parent_cb is not None:
                    config = _merge_callback_into_config(config, parent_cb)
                    return await original_astream(input, config=config, **kwargs)

                callback = LangGraphCallback(content_limit=_content_limit)

                integration.client.create_trace({
                    'trace_id': callback.trace_id,
                    'name': trace_name,
                    'framework': integration._framework,
                    'started_at': datetime.now(timezone.utc).isoformat(),
                })

                callback._span_sender = lambda span: integration.client.queue_span(
                    callback.trace_id, span
                )

                token = _active_callback_ctx.set(callback)
                config = _merge_callback_into_config(config, callback)

                try:
                    stream_iter = await original_astream(input, config=config, **kwargs)
                    # Wrap the async iterator to finalize trace when exhausted
                    return _AsyncStreamIteratorWrapper(
                        stream_iter, integration, callback, trace_name, token
                    )
                except Exception as exc:
                    callback.status = "failed"
                    callback.errors.append({
                        'timestamp': datetime.now(timezone.utc).isoformat(),
                        'error_type': type(exc).__name__,
                        'message': str(exc),
                        'agent_name': callback.current_agent or 'unknown',
                    })
                    integration._complete_trace(callback, 'failed', trace_name)
                    _active_callback_ctx.reset(token)
                    raise

            runnable.astream = wrapped_astream

        logger.debug(f"Instrumented {type(runnable).__name__} (id={runnable_id})")

    def uninstrument(self, runnable):
        """Remove instrumentation from a previously instrumented runnable.

        Args:
            runnable: The runnable that was passed to instrument()
        """
        runnable_id = id(runnable)
        entry = _instrumented_runnables.pop(runnable_id, None)
        if entry is None:
            logger.warning("Runnable was not instrumented — nothing to undo")
            return
        runnable.invoke = entry['invoke']
        if entry.get('ainvoke') is not None:
            runnable.ainvoke = entry['ainvoke']
        if entry.get('astream_events') is not None:
            runnable.astream_events = entry['astream_events']
        if entry.get('batch') is not None:
            runnable.batch = entry['batch']
        if entry.get('abatch') is not None:
            runnable.abatch = entry['abatch']
        if entry.get('stream') is not None:
            runnable.stream = entry['stream']
        if entry.get('astream') is not None:
            runnable.astream = entry['astream']
        logger.debug(f"Uninstrumented {type(runnable).__name__} (id={runnable_id})")

    def _complete_trace(self, callback, status: str, name: Optional[str] = None):
        """Send the trace completion PATCH with summary aggregates.

        Spans have already been streamed during execution via callback._span_sender.
        This just closes the trace with final totals.
        """
        metrics = callback.get_metrics()

        total_input_tokens = metrics['total_input_tokens']
        total_output_tokens = metrics['total_output_tokens']
        total_cost = metrics['total_cost']
        total_tokens = total_input_tokens + total_output_tokens

        # Use the first model seen in LLM calls
        model = callback.llm_calls[0]['model'] if callback.llm_calls else None

        summary_data = {
            'status': status,
            'ended_at': metrics['end_time'],
            'duration_ms': metrics['duration_ms'],
            'prompt': callback.prompt or "No prompt captured",
            'model': model,
            'total_cost': total_cost,
            'total_tokens': total_tokens,
            'total_input_tokens': total_input_tokens,
            'total_output_tokens': total_output_tokens,
        }

        success = self.client.complete_trace(callback.trace_id, summary_data)

        duration_s = metrics['duration_ms'] / 1000
        llm_calls = len([s for s in metrics['steps'] if s.get('type') == 'llm_call'])
        tool_calls_count = len([s for s in metrics['steps'] if s.get('type') == 'tool_call'])

        summary = TraceSummary(
            name=name or self._framework,
            status=status,
            llm_calls=llm_calls,
            tool_calls=tool_calls_count,
            total_tokens=total_tokens,
            total_cost=total_cost,
            duration_s=duration_s,
            agents=metrics['agents_executed'],
            sent=success,
        )
        print(f"[Visibe] {summary}")


class LangChainCallback(BaseCallbackHandler):
    """
    LangChain callback handler that captures metrics for Visibe.

    Captures:
    - LLM calls (model, tokens, cost, input/output text)
    - Tool usage (name, input, output, duration, status)
    - Chain/Agent execution
    - Errors
    """

    # Default content limits (characters)
    DEFAULT_LLM_CONTENT_LIMIT = 1000
    DEFAULT_TOOL_CONTENT_LIMIT = 500

    def __init__(self, content_limit: Optional[int] = None):
        """Initialize the callback handler.

        Args:
            content_limit: Max characters for LLM/tool input/output in steps.
                          None = use defaults (1000 for LLM, 500 for tools).
                          0 = omit content entirely. Any positive int = limit for both.
        """
        super().__init__()

        # Content truncation
        if content_limit is not None:
            self._llm_content_limit = content_limit
            self._tool_content_limit = content_limit
        else:
            self._llm_content_limit = self.DEFAULT_LLM_CONTENT_LIMIT
            self._tool_content_limit = self.DEFAULT_TOOL_CONTENT_LIMIT

        # Trace-level data
        self.trace_id = str(uuid.uuid4())
        self.start_time: Optional[datetime] = None
        self.end_time: Optional[datetime] = None
        self.prompt: Optional[str] = None
        self.status = "completed"

        # LLM tracking
        self.llm_calls: List[Dict[str, Any]] = []
        self.current_llm_call: Optional[Dict[str, Any]] = None

        # Tool tracking
        self.tool_calls: List[Dict[str, Any]] = []
        self.current_tool_call: Optional[Dict[str, Any]] = None

        # Agent tracking
        self.agents_executed: List[str] = []
        self.current_agent: Optional[str] = None

        # Errors
        self.errors: List[Dict[str, Any]] = []

        # Steps timeline
        self.steps: List[Dict[str, Any]] = []
        self.step_counter = 0

        # Parent-child tracking: maps run_id → step_id
        self._run_to_step: Dict[str, str] = {}

        # Maps agent_name → step_id of its agent_start span (for parent resolution)
        self._agent_step_ids: Dict[str, str] = {}

        # Span streaming: if set, each step is sent to backend as it's created
        self._span_sender = None

    def _create_step(self, step_type: str, timestamp: datetime,
                     run_id=None, parent_run_id=None,
                     fallback_parent_step_id=None, **kwargs) -> str:
        """Create a step and add it to the timeline.

        Args:
            run_id: LangChain run_id for this callback invocation.
            parent_run_id: LangChain parent_run_id — used to set parent_step_id.
            fallback_parent_step_id: Explicit parent step_id used when
                parent_run_id is None or unresolvable (e.g. sub-graph roots
                invoked through the Pregel patch in init() mode).
        """
        self.step_counter += 1
        step_id = f'step_{self.step_counter}'

        # Resolve parent step from run_id mapping
        parent_step_id = None
        if parent_run_id:
            parent_step_id = self._run_to_step.get(str(parent_run_id))
        if not parent_step_id and fallback_parent_step_id:
            parent_step_id = fallback_parent_step_id

        # Register this run_id → step_id mapping
        if run_id:
            self._run_to_step[str(run_id)] = step_id

        step = {
            'span_id': step_id,
            'type': step_type,
            'timestamp': timestamp.isoformat(),
            **kwargs
        }
        if parent_step_id:
            step['parent_span_id'] = parent_step_id

        self.steps.append(step)

        # Stream span to backend if sender is configured
        if self._span_sender:
            try:
                self._span_sender(step)
            except Exception as e:
                logger.warning(f"Failed to send span {step_id}: {e}")

        return step_id

    def _truncate_tool(self, text: Optional[str]) -> str:
        """Truncate tool input/output text according to content limit."""
        if not text:
            return ""
        if self._tool_content_limit == 0:
            return ""
        if len(text) > self._tool_content_limit:
            return text[:self._tool_content_limit] + '...'
        return text

    # ==================== Chain Callbacks ====================

    def on_chain_start(
        self,
        serialized: Optional[Dict[str, Any]],
        inputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chain starts running."""
        if self.start_time is None:
            self.start_time = datetime.now(timezone.utc)

        # Extract agent/chain name (handle None serialized)
        agent_name = None
        if serialized:
            agent_name = serialized.get("name")
            if not agent_name:
                id_list = serialized.get("id", [])
                if id_list:
                    agent_name = id_list[-1]

        if not agent_name:
            agent_name = "Agent"

        # Capture the input prompt
        if self.prompt is None and inputs:
            if "input" in inputs:
                self.prompt = inputs["input"]
            elif "question" in inputs:
                self.prompt = inputs["question"]
            else:
                # Take first string value as prompt
                for v in inputs.values():
                    if isinstance(v, str):
                        self.prompt = v
                        break

        # Track agent if it's a new one (only track meaningful names)
        if agent_name and agent_name not in self.agents_executed and agent_name != "Agent":
            self.agents_executed.append(agent_name)
            self.current_agent = agent_name

            step_id = self._create_step(
                step_type='agent_start',
                timestamp=datetime.now(timezone.utc),
                run_id=run_id,
                parent_run_id=parent_run_id,
                agent_name=agent_name,
                description=f"Agent started: {agent_name}"
            )
            self._agent_step_ids[agent_name] = step_id
        elif agent_name and agent_name in self.agents_executed:
            self.current_agent = agent_name
            # Register run_id → original agent_start step for parent resolution
            if run_id and agent_name in self._agent_step_ids:
                self._run_to_step[str(run_id)] = self._agent_step_ids[agent_name]
        else:
            # Skipped chain (generic name like "Agent") — propagate run_id
            # so downstream callbacks (tool_start, llm_start) can still
            # resolve their parent_run_id through this intermediate chain.
            if run_id and parent_run_id:
                resolved = self._run_to_step.get(str(parent_run_id))
                if resolved:
                    self._run_to_step[str(run_id)] = resolved

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chain ends."""
        self.end_time = datetime.now(timezone.utc)

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chain errors."""
        self.status = "failed"
        self.end_time = datetime.now(timezone.utc)

        error_info = {
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'error_type': 'CHAIN_ERROR',
            'message': str(error),
            'agent_name': self.current_agent or 'unknown'
        }
        self.errors.append(error_info)

        self._create_step(
            step_type='error',
            timestamp=datetime.now(timezone.utc),
            run_id=run_id,
            parent_run_id=parent_run_id,
            agent_name=self.current_agent or 'unknown',
            error_type='CHAIN_ERROR',
            error_message=str(error),
            description=f"Error: {type(error).__name__}"
        )

    # ==================== LLM Callbacks ====================

    def on_llm_start(
        self,
        serialized: Optional[Dict[str, Any]],
        prompts: List[str],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM starts running."""
        # Extract model name (handle None serialized)
        model = "unknown"
        if serialized and "kwargs" in serialized:
            model = serialized["kwargs"].get("model_name") or serialized["kwargs"].get("model") or "unknown"
        if model == "unknown" and "invocation_params" in kwargs:
            inv_params = kwargs["invocation_params"] or {}
            model = inv_params.get("model_name") or inv_params.get("model") or "unknown"

        # Store current LLM call info
        self.current_llm_call = {
            'start_time': datetime.now(timezone.utc),
            'model': model,
            'input_text': prompts[0] if prompts else "",
            'agent_name': self.current_agent or 'unknown',
            'run_id': run_id,
            'parent_run_id': parent_run_id,
        }

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM ends running."""
        if not self.current_llm_call:
            return

        end_time = datetime.now(timezone.utc)
        start_time = self.current_llm_call['start_time']
        duration_ms = int((end_time - start_time).total_seconds() * 1000)

        # Extract token usage from multiple possible locations
        input_tokens = 0
        output_tokens = 0

        logger.debug(f"on_llm_end: llm_output={response.llm_output}")

        # Try llm_output.token_usage (older format)
        if response.llm_output:
            token_usage = response.llm_output.get("token_usage", {})
            if token_usage:
                input_tokens = token_usage.get("prompt_tokens", 0)
                output_tokens = token_usage.get("completion_tokens", 0)
                logger.debug(f"Found tokens in llm_output.token_usage: {input_tokens}/{output_tokens}")

        # Try generation_info (newer langchain-openai format)
        if input_tokens == 0 and response.generations and response.generations[0]:
            gen = response.generations[0][0]
            logger.debug(f"Generation type: {type(gen).__name__}")

            if hasattr(gen, 'generation_info') and gen.generation_info:
                logger.debug(f"generation_info: {gen.generation_info}")
                usage = gen.generation_info.get('usage_metadata', {})
                if usage:
                    input_tokens = usage.get('input_tokens', 0)
                    output_tokens = usage.get('output_tokens', 0)
                    logger.debug(f"Found tokens in generation_info: {input_tokens}/{output_tokens}")

            # Also try message.usage_metadata for AIMessage
            if input_tokens == 0 and hasattr(gen, 'message'):
                logger.debug(f"Message type: {type(gen.message).__name__}")
                if hasattr(gen.message, 'usage_metadata') and gen.message.usage_metadata:
                    usage = gen.message.usage_metadata
                    logger.debug(f"message.usage_metadata: {usage}")
                    input_tokens = getattr(usage, 'input_tokens', 0)
                    output_tokens = getattr(usage, 'output_tokens', 0)
                    logger.debug(f"Found tokens in message.usage_metadata: {input_tokens}/{output_tokens}")

                # Try response_metadata
                if input_tokens == 0 and hasattr(gen.message, 'response_metadata'):
                    resp_meta = gen.message.response_metadata
                    logger.debug(f"message.response_metadata: {resp_meta}")
                    if resp_meta and 'token_usage' in resp_meta:
                        usage = resp_meta['token_usage']
                        input_tokens = usage.get('prompt_tokens', 0)
                        output_tokens = usage.get('completion_tokens', 0)
                        logger.debug(f"Found tokens in response_metadata: {input_tokens}/{output_tokens}")

        logger.debug(f"Final tokens: input={input_tokens}, output={output_tokens}")

        # Extract output text and model from generation
        output_text = ""
        if response.generations and response.generations[0]:
            gen = response.generations[0][0]
            output_text = gen.text

            # If gen.text is empty, the LLM may have responded with tool calls
            if not output_text and hasattr(gen, 'message') and hasattr(gen.message, 'tool_calls'):
                tool_calls = gen.message.tool_calls
                if tool_calls:
                    parts = []
                    for tc in tool_calls:
                        name = tc.get('name', 'unknown') if isinstance(tc, dict) else getattr(tc, 'name', 'unknown')
                        args = tc.get('args', {}) if isinstance(tc, dict) else getattr(tc, 'args', {})
                        parts.append(f"{name}({args})")
                    output_text = "; ".join(parts)

            # Try to get model from generation_info (more accurate)
            if hasattr(gen, 'generation_info') and gen.generation_info:
                model_from_gen = gen.generation_info.get('model_name')
                if model_from_gen:
                    self.current_llm_call['model'] = model_from_gen

        # Fallback: estimate tokens using tiktoken if we still have 0
        token_source = "api"
        if input_tokens == 0 and output_tokens == 0:
            input_text_for_estimate = self.current_llm_call['input_text']
            model = self.current_llm_call['model']
            input_tokens = estimate_tokens(input_text_for_estimate, model)
            output_tokens = estimate_tokens(output_text, model)
            token_source = "estimated"
            logger.debug(f"Using tiktoken estimation: input={input_tokens}, output={output_tokens}")

        # Calculate cost
        model = self.current_llm_call['model']
        cost = calculate_cost(model, input_tokens, output_tokens)

        # Truncate texts if too long
        input_text = self.current_llm_call['input_text']
        if self._llm_content_limit == 0:
            input_text = ""
            output_text = ""
        else:
            if len(input_text) > self._llm_content_limit:
                input_text = input_text[:self._llm_content_limit] + '...'
            if len(output_text) > self._llm_content_limit:
                output_text = output_text[:self._llm_content_limit] + '...'

        # Create LLM call record
        llm_call = {
            'start_time': start_time.isoformat(),
            'end_time': end_time.isoformat(),
            'model': model,
            'agent_name': self.current_llm_call['agent_name'],
            'input_tokens': input_tokens,
            'output_tokens': output_tokens,
            'total_tokens': input_tokens + output_tokens,
            'cost': cost
        }
        self.llm_calls.append(llm_call)

        # Detect provider from model name
        provider = 'openai'
        model_lower = model.lower()
        if 'claude' in model_lower or 'anthropic' in model_lower:
            provider = 'anthropic'
        elif 'gemini' in model_lower or 'palm' in model_lower:
            provider = 'google'

        # Create step
        self._create_step(
            step_type='llm_call',
            timestamp=start_time,
            run_id=self.current_llm_call.get('run_id'),
            parent_run_id=self.current_llm_call.get('parent_run_id'),
            duration_ms=duration_ms,
            agent_name=self.current_llm_call['agent_name'],
            model=model,
            provider=provider,
            status='success',
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=cost,
            description=f"LLM Call using {model}",
            input_text=input_text,
            output_text=output_text
        )

        self.current_llm_call = None

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM errors."""
        self.status = "failed"

        error_info = {
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'error_type': 'LLM_ERROR',
            'message': str(error),
            'agent_name': self.current_agent or 'unknown'
        }
        self.errors.append(error_info)

        self._create_step(
            step_type='error',
            timestamp=datetime.now(timezone.utc),
            run_id=run_id,
            parent_run_id=parent_run_id,
            agent_name=self.current_agent or 'unknown',
            error_type='LLM_ERROR',
            error_message=str(error),
            description=f"LLM Error: {type(error).__name__}"
        )

        self.current_llm_call = None

    # ==================== Tool Callbacks ====================

    def on_tool_start(
        self,
        serialized: Optional[Dict[str, Any]],
        input_str: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool starts running."""
        tool_name = (serialized.get("name") if serialized else None) or "unknown"

        # Propagate run_id so sub-graph callbacks inside this tool can
        # resolve their parent_run_id back through the tool's ancestry.
        if run_id and parent_run_id:
            resolved = self._run_to_step.get(str(parent_run_id))
            if resolved:
                self._run_to_step[str(run_id)] = resolved

        self.current_tool_call = {
            'start_time': datetime.now(timezone.utc),
            'tool_name': tool_name,
            'input': self._truncate_tool(input_str),
            'agent_name': self.current_agent or 'unknown',
            'run_id': run_id,
            'parent_run_id': parent_run_id,
        }

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool ends running."""
        if not self.current_tool_call:
            return

        end_time = datetime.now(timezone.utc)
        start_time = self.current_tool_call['start_time']
        duration_ms = int((end_time - start_time).total_seconds() * 1000)

        # Extract clean content from output (handles ToolMessage objects)
        if hasattr(output, 'content'):
            output_text = str(output.content)
        else:
            output_text = str(output)
        if self._tool_content_limit == 0:
            output_text = ""
        elif len(output_text) > self._tool_content_limit:
            output_text = output_text[:self._tool_content_limit] + '...'

        # Create tool call record
        tool_call = {
            'tool_name': self.current_tool_call['tool_name'],
            'agent_name': self.current_tool_call['agent_name'],
            'execution_time_ms': duration_ms,
            'status': 'success',
            'start_time': start_time.isoformat(),
            'end_time': end_time.isoformat(),
            'input': self.current_tool_call['input'],
            'output': output_text
        }
        self.tool_calls.append(tool_call)

        # Create step
        self._create_step(
            step_type='tool_call',
            timestamp=start_time,
            run_id=self.current_tool_call.get('run_id'),
            parent_run_id=self.current_tool_call.get('parent_run_id'),
            duration_ms=duration_ms,
            agent_name=self.current_tool_call['agent_name'],
            tool_name=self.current_tool_call['tool_name'],
            status='success',
            description=f"Tool: {self.current_tool_call['tool_name']}",
            input_text=self.current_tool_call['input'],
            output_text=output_text
        )

        self.current_tool_call = None

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool errors."""
        if not self.current_tool_call:
            return

        end_time = datetime.now(timezone.utc)
        start_time = self.current_tool_call['start_time']
        duration_ms = int((end_time - start_time).total_seconds() * 1000)

        # Create tool call record with error
        tool_call = {
            'tool_name': self.current_tool_call['tool_name'],
            'agent_name': self.current_tool_call['agent_name'],
            'execution_time_ms': duration_ms,
            'status': 'failed',
            'start_time': start_time.isoformat(),
            'end_time': end_time.isoformat(),
            'input': self.current_tool_call['input'],
            'output': f"ERROR: {str(error)}"
        }
        self.tool_calls.append(tool_call)

        # Create error record
        error_info = {
            'timestamp': end_time.isoformat(),
            'error_type': 'TOOL_ERROR',
            'message': str(error),
            'agent_name': self.current_tool_call['agent_name'],
            'tool_name': self.current_tool_call['tool_name']
        }
        self.errors.append(error_info)

        # Create steps
        self._create_step(
            step_type='tool_call',
            timestamp=start_time,
            run_id=self.current_tool_call.get('run_id'),
            parent_run_id=self.current_tool_call.get('parent_run_id'),
            duration_ms=duration_ms,
            agent_name=self.current_tool_call['agent_name'],
            tool_name=self.current_tool_call['tool_name'],
            status='failed',
            description=f"Tool: {self.current_tool_call['tool_name']}",
            input_text=self.current_tool_call['input'],
            output_text=f"ERROR: {str(error)}"
        )

        self._create_step(
            step_type='error',
            timestamp=end_time,
            run_id=self.current_tool_call.get('run_id'),
            parent_run_id=self.current_tool_call.get('parent_run_id'),
            agent_name=self.current_tool_call['agent_name'],
            error_type='TOOL_ERROR',
            error_message=str(error),
            tool_name=self.current_tool_call['tool_name'],
            description=f"Tool Error: {self.current_tool_call['tool_name']}"
        )

        self.current_tool_call = None

    # ==================== Agent Callbacks ====================

    def on_agent_action(
        self,
        action: AgentAction,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an agent takes an action."""
        # This is called before tool execution, we can use it for additional context
        pass

    def on_agent_finish(
        self,
        finish: AgentFinish,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an agent finishes."""
        self.end_time = datetime.now(timezone.utc)

    # ==================== Metrics ====================

    def get_metrics(self) -> Dict[str, Any]:
        """Get all captured metrics in a structured format."""
        if not self.start_time:
            self.start_time = datetime.now(timezone.utc)
        if not self.end_time:
            self.end_time = datetime.now(timezone.utc)

        duration_ms = int((self.end_time - self.start_time).total_seconds() * 1000)

        # Build cost breakdown by agent
        cost_breakdown = {}
        for call in self.llm_calls:
            agent = call['agent_name']
            if agent not in cost_breakdown:
                cost_breakdown[agent] = {
                    'agent_name': agent,
                    'model': call['model'],
                    'input_tokens': 0,
                    'output_tokens': 0,
                    'cost': 0.0
                }
            cost_breakdown[agent]['input_tokens'] += call['input_tokens']
            cost_breakdown[agent]['output_tokens'] += call['output_tokens']
            cost_breakdown[agent]['cost'] += call['cost']

        # Calculate totals
        total_cost = sum(agent['cost'] for agent in cost_breakdown.values())
        total_input_tokens = sum(agent['input_tokens'] for agent in cost_breakdown.values())
        total_output_tokens = sum(agent['output_tokens'] for agent in cost_breakdown.values())

        # Sort steps chronologically
        sorted_steps = sorted(self.steps, key=lambda x: x['timestamp'])

        # Enrich tool calls with token estimates
        enriched_tools = []
        for tool in self.tool_calls:
            tool_start = datetime.fromisoformat(tool['start_time'])
            tool_end = datetime.fromisoformat(tool['end_time'])
            tool_agent = tool['agent_name']

            associated_tokens = 0
            associated_cost = 0.0

            for llm in self.llm_calls:
                if llm['agent_name'] != tool_agent:
                    continue

                llm_start = datetime.fromisoformat(llm['start_time'])
                llm_end = datetime.fromisoformat(llm['end_time'])

                # LLM ended within 1 second before tool started
                if llm_end <= tool_start and (tool_start - llm_end).total_seconds() < 1.0:
                    associated_tokens += llm['total_tokens']
                    associated_cost += llm['cost']

                # LLM started within 1 second after tool ended
                if llm_start >= tool_end and (llm_start - tool_end).total_seconds() < 1.0:
                    associated_tokens += llm['total_tokens']
                    associated_cost += llm['cost']

            enriched_tool = {
                'tool_name': tool['tool_name'],
                'agent_name': tool['agent_name'],
                'execution_time_ms': tool['execution_time_ms'],
                'status': tool['status'],
            }
            if tool.get('input'):
                enriched_tool['input'] = tool['input']
            if tool.get('output'):
                enriched_tool['output'] = tool['output']
            if associated_tokens > 0:
                enriched_tool['total_tokens'] = associated_tokens
            if associated_cost > 0:
                enriched_tool['estimated_cost'] = round(associated_cost, 6)
            enriched_tools.append(enriched_tool)

        return {
            'duration_ms': duration_ms,
            'start_time': self.start_time.isoformat(),
            'end_time': self.end_time.isoformat(),
            'total_cost': round(total_cost, 6),
            'total_input_tokens': total_input_tokens,
            'total_output_tokens': total_output_tokens,
            'agent_cost_breakdown': list(cost_breakdown.values()),
            'agents_executed': self.agents_executed,
            'tools_used': enriched_tools,
            'errors': self.errors,
            'steps': sorted_steps
        }


class LangGraphCallback(LangChainCallback):
    """
    LangGraph-specific callback that extends LangChainCallback
    with node-level tracking.

    Inherits all LLM, tool, and error tracking from LangChainCallback.
    Overrides on_chain_start to extract LangGraph node information
    from callback metadata and build hierarchical agent names using
    the _tool_stack_ctx ContextVar (propagated via copy_context).
    """

    def __init__(self, content_limit: Optional[int] = None):
        super().__init__(content_limit=content_limit)
        self.nodes_executed: List[str] = []
        self.current_node: Optional[str] = None

    @staticmethod
    def _current_tool_prefix() -> str:
        """Return the '/'-joined tool ancestry from the ContextVar.

        E.g. if the current context is inside tool 'plan_day' → 'research_activity',
        returns 'plan_day/research_activity'.  Returns '' at root level.
        """
        return '/'.join(_tool_stack_ctx.get())

    def on_chain_start(
        self,
        serialized: Optional[Dict[str, Any]],
        inputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Override to extract LangGraph node information from metadata."""
        if self.start_time is None:
            self.start_time = datetime.now(timezone.utc)

        # Extract LangGraph node name from metadata
        node_name = None
        if metadata and isinstance(metadata, dict):
            node_name = metadata.get("langgraph_node")

        if node_name:
            # Build hierarchical name: prefix with current thread's tool stack
            tool_prefix = self._current_tool_prefix()
            display_name = f"{tool_prefix}/{node_name}" if tool_prefix else node_name

            self.current_node = node_name
            self.current_agent = display_name

            if node_name not in self.nodes_executed:
                self.nodes_executed.append(node_name)

            if display_name not in self.agents_executed:
                self.agents_executed.append(display_name)

                # When a sub-graph is invoked via the Pregel patch (init()
                # mode), its root on_chain_start has parent_run_id=None.
                # Resolve the parent from the containing graph's tools node.
                fallback_parent = None
                if parent_run_id is None and tool_prefix:
                    parts = tool_prefix.split('/')
                    parent_tools_name = ('/'.join(parts[:-1]) + '/tools') if len(parts) > 1 else 'tools'
                    fallback_parent = self._agent_step_ids.get(parent_tools_name)

                step_id = self._create_step(
                    step_type='agent_start',
                    timestamp=datetime.now(timezone.utc),
                    run_id=run_id,
                    parent_run_id=parent_run_id,
                    fallback_parent_step_id=fallback_parent,
                    agent_name=display_name,
                    description=f"Node started: {display_name}",
                )
                self._agent_step_ids[display_name] = step_id
            else:
                # Node seen before — register run_id → original step for parent resolution
                if run_id and display_name in self._agent_step_ids:
                    self._run_to_step[str(run_id)] = self._agent_step_ids[display_name]

            if self.prompt is None and inputs:
                self._extract_prompt(inputs)
        else:
            super().on_chain_start(
                serialized, inputs,
                run_id=run_id, parent_run_id=parent_run_id,
                tags=tags, metadata=metadata,
                **kwargs,
            )

    def on_tool_start(
        self,
        serialized: Optional[Dict[str, Any]],
        input_str: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Override to push tool name onto the ContextVar tool stack."""
        tool_name = (serialized.get("name") if serialized else None) or "unknown"

        # Push onto the ContextVar tool stack (immutable tuple append)
        _tool_stack_ctx.set(_tool_stack_ctx.get() + (tool_name,))

        # Build hierarchical agent name for this tool call
        # The stack already contains this tool, so use the parent's prefix
        # (everything except the last element which is this tool itself)
        stack = _tool_stack_ctx.get()
        parent_prefix = '/'.join(stack[:-1])
        display_agent = f"{parent_prefix}/{tool_name}" if parent_prefix else tool_name
        self.current_agent = display_agent

        # Delegate to parent for the actual tool tracking
        super().on_tool_start(
            serialized, input_str,
            run_id=run_id, parent_run_id=parent_run_id,
            tags=tags, metadata=metadata,
            **kwargs,
        )
        # Override the agent_name the parent stored with our hierarchical name
        if self.current_tool_call:
            self.current_tool_call['agent_name'] = display_agent

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Override to pop tool name from the ContextVar tool stack."""
        # Delegate to parent first (uses current_tool_call)
        super().on_tool_end(output, run_id=run_id, parent_run_id=parent_run_id, **kwargs)
        # Pop the tool from the stack (immutable tuple slice)
        stack = _tool_stack_ctx.get()
        if stack:
            _tool_stack_ctx.set(stack[:-1])

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Override to pop tool name from the ContextVar tool stack on error."""
        super().on_tool_error(error, run_id=run_id, parent_run_id=parent_run_id, **kwargs)
        stack = _tool_stack_ctx.get()
        if stack:
            _tool_stack_ctx.set(stack[:-1])

    def on_llm_start(
        self,
        serialized: Optional[Dict[str, Any]],
        prompts: List[str],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Override to derive agent name from ContextVar tool stack.

        self.current_agent may be stale when LangGraph runs tools in
        parallel threads sharing this callback instance.  Reading from
        the _tool_stack_ctx ContextVar is context-safe and deterministic.
        """
        # Delegate to parent for model extraction & llm_call setup
        super().on_llm_start(
            serialized, prompts,
            run_id=run_id, parent_run_id=parent_run_id,
            tags=tags, metadata=metadata,
            **kwargs,
        )

        # Now fix up the agent_name with hierarchical name
        if self.current_llm_call:
            node_name = None
            if metadata and isinstance(metadata, dict):
                node_name = metadata.get("langgraph_node")
            if node_name:
                tool_prefix = self._current_tool_prefix()
                display_name = f"{tool_prefix}/{node_name}" if tool_prefix else node_name
                self.current_llm_call['agent_name'] = display_name

    def _extract_prompt(self, inputs: Any) -> None:
        """Extract the user's prompt from LangGraph inputs."""
        if not isinstance(inputs, dict):
            return

        if "messages" in inputs:
            msgs = inputs["messages"]
            if isinstance(msgs, list) and msgs:
                last_msg = msgs[-1]
                if isinstance(last_msg, tuple) and len(last_msg) == 2:
                    self.prompt = str(last_msg[1])
                elif hasattr(last_msg, 'content'):
                    self.prompt = str(last_msg.content)
                elif isinstance(last_msg, str):
                    self.prompt = last_msg
        elif "input" in inputs:
            self.prompt = inputs["input"]
        elif "question" in inputs:
            self.prompt = inputs["question"]
        else:
            for v in inputs.values():
                if isinstance(v, str):
                    self.prompt = v
                    break


class LangChainTracker:
    """
    Context manager for tracking LangChain/LangGraph execution.

    Usage:
        with tracker.track(name="my-workflow"):
            result = agent.invoke({"input": "..."}, config={"callbacks": [tracker.callback]})
    """

    def __init__(self, integration: LangChainIntegration, name: str,
                 content_limit: Optional[int] = None):
        """
        Initialize the tracker.

        Args:
            integration: LangChainIntegration instance
            name: Name for this execution trace
            content_limit: Max chars for LLM/tool content. None = use integration default.
        """
        self.integration = integration
        self.name = name
        self._content_limit = content_limit or integration._content_limit
        self.callback = None  # Set in __enter__ (fresh per block)
        self.summary = None
        # Provider-specific callbacks
        self._openai_callback = None
        self._openai_callback_ctx = None
        self._bedrock_callback = None
        self._bedrock_callback_ctx = None

    def __enter__(self):
        """Enter the context manager — create fresh callback, trace header, wire span streaming."""

        # Fresh callback per block to avoid state pollution and support concurrent track() calls
        self.callback = LangGraphCallback(content_limit=self._content_limit)
        self.integration._callback = self.callback  # So integration.callback returns active callback

        # Create trace header on backend
        self.integration.client.create_trace({
            'trace_id': self.callback.trace_id,
            'name': self.name,
            'framework': self.integration._framework,
            'started_at': datetime.now(timezone.utc).isoformat(),
        })

        # Wire up span streaming on the callback
        self.callback._span_sender = lambda span: self.integration.client.queue_span(
            self.callback.trace_id, span
        )

        # Start provider-specific callbacks for accurate token tracking
        if OPENAI_CALLBACK_AVAILABLE:
            self._openai_callback_ctx = get_openai_callback()
            self._openai_callback = self._openai_callback_ctx.__enter__()

        if BEDROCK_CALLBACK_AVAILABLE:
            self._bedrock_callback_ctx = get_bedrock_anthropic_callback()
            self._bedrock_callback = self._bedrock_callback_ctx.__enter__()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit the context manager — complete trace with summary."""
        # Get token data from provider-specific callbacks
        provider_tokens = None
        token_source = None

        if self._openai_callback_ctx:
            try:
                self._openai_callback_ctx.__exit__(exc_type, exc_val, exc_tb)
                if self._openai_callback and self._openai_callback.total_tokens > 0:
                    provider_tokens = {
                        'prompt_tokens': self._openai_callback.prompt_tokens,
                        'completion_tokens': self._openai_callback.completion_tokens,
                        'total_cost': self._openai_callback.total_cost,
                    }
                    token_source = "openai_callback"
            except Exception as e:
                logger.warning(f"Failed to get OpenAI callback data: {e}")

        if self._bedrock_callback_ctx:
            try:
                self._bedrock_callback_ctx.__exit__(exc_type, exc_val, exc_tb)
                if not provider_tokens and self._bedrock_callback:
                    if hasattr(self._bedrock_callback, 'total_tokens') and self._bedrock_callback.total_tokens > 0:
                        provider_tokens = {
                            'prompt_tokens': getattr(self._bedrock_callback, 'prompt_tokens', 0),
                            'completion_tokens': getattr(self._bedrock_callback, 'completion_tokens', 0),
                            'total_cost': getattr(self._bedrock_callback, 'total_cost', 0),
                        }
                        token_source = "bedrock_callback"
            except Exception as e:
                logger.warning(f"Failed to get Bedrock callback data: {e}")

        # Update status if there was an exception
        if exc_type is not None:
            self.callback.status = "failed"
            if not self.callback.errors:
                self.callback.errors.append({
                    'timestamp': datetime.now(timezone.utc).isoformat(),
                    'error_type': exc_type.__name__,
                    'message': str(exc_val),
                    'agent_name': self.callback.current_agent or 'unknown'
                })

        # Collect metrics
        metrics = self.callback.get_metrics()

        # Use provider callback data if available and callback data is empty
        total_input_tokens = metrics['total_input_tokens']
        total_output_tokens = metrics['total_output_tokens']
        total_cost = metrics['total_cost']

        if provider_tokens and (total_input_tokens == 0 or total_cost == 0):
            total_input_tokens = provider_tokens['prompt_tokens']
            total_output_tokens = provider_tokens['completion_tokens']
            total_cost = provider_tokens['total_cost']
            logger.debug(f"Using {token_source} for token tracking")

        has_errors = len(metrics['errors']) > 0
        status = 'failed' if (exc_type is not None or self.callback.status == "failed" or has_errors) else 'completed'

        total_tokens = total_input_tokens + total_output_tokens

        # Use the first model seen in LLM calls
        model = self.callback.llm_calls[0]['model'] if self.callback.llm_calls else None

        # Complete trace with summary
        summary_data = {
            'status': status,
            'ended_at': metrics['end_time'],
            'duration_ms': metrics['duration_ms'],
            'prompt': self.callback.prompt or "No prompt captured",
            'model': model,
            'total_cost': total_cost,
            'total_tokens': total_tokens,
            'total_input_tokens': total_input_tokens,
            'total_output_tokens': total_output_tokens,
        }

        success = self.integration.client.complete_trace(
            self.callback.trace_id, summary_data
        )

        duration_s = metrics['duration_ms'] / 1000
        llm_calls = len([s for s in metrics['steps'] if s.get('type') == 'llm_call'])
        tool_calls = len([s for s in metrics['steps'] if s.get('type') == 'tool_call'])

        self.summary = TraceSummary(
            name=self.name,
            status=status,
            llm_calls=llm_calls,
            tool_calls=tool_calls,
            total_tokens=total_tokens,
            total_cost=total_cost,
            duration_s=duration_s,
            agents=metrics['agents_executed'],
            sent=success,
        )

        # Prepare fresh callback for next track() call
        self.integration._callback = LangGraphCallback(
            content_limit=self.integration._content_limit
        )

        # Don't suppress exceptions
        return False

# ------------------------------------------------------------------
# Stream iterator wrappers
# ------------------------------------------------------------------

class _StreamIteratorWrapper:
    """Wraps a sync stream iterator to finalize trace when exhausted."""

    def __init__(self, stream_iter, integration, callback, trace_name, token):
        self._stream_iter = stream_iter
        self._integration = integration
        self._callback = callback
        self._trace_name = trace_name
        self._token = token
        self._exhausted = False

    def __iter__(self):
        return self

    def __next__(self):
        try:
            return next(self._stream_iter)
        except StopIteration:
            if not self._exhausted:
                self._exhausted = True
                self._integration._complete_trace(
                    self._callback, 'completed', self._trace_name
                )
                _active_callback_ctx.reset(self._token)
            raise

    def __enter__(self):
        if hasattr(self._stream_iter, '__enter__'):
            self._stream_iter.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if not self._exhausted:
            self._exhausted = True
            if exc_type:
                self._callback.status = "failed"
                self._callback.errors.append({
                    'timestamp': datetime.now(timezone.utc).isoformat(),
                    'error_type': exc_type.__name__,
                    'message': str(exc_val),
                })
                self._integration._complete_trace(
                    self._callback, 'failed', self._trace_name
                )
            else:
                self._integration._complete_trace(
                    self._callback, 'completed', self._trace_name
                )
            _active_callback_ctx.reset(self._token)

        if hasattr(self._stream_iter, '__exit__'):
            return self._stream_iter.__exit__(exc_type, exc_val, exc_tb)
        return False


class _AsyncStreamIteratorWrapper:
    """Wraps an async stream iterator to finalize trace when exhausted."""

    def __init__(self, stream_iter, integration, callback, trace_name, token):
        self._stream_iter = stream_iter
        self._integration = integration
        self._callback = callback
        self._trace_name = trace_name
        self._token = token
        self._exhausted = False

    def __aiter__(self):
        return self

    async def __anext__(self):
        try:
            return await self._stream_iter.__anext__()
        except StopAsyncIteration:
            if not self._exhausted:
                self._exhausted = True
                self._integration._complete_trace(
                    self._callback, 'completed', self._trace_name
                )
                _active_callback_ctx.reset(self._token)
            raise

    async def __aenter__(self):
        if hasattr(self._stream_iter, '__aenter__'):
            await self._stream_iter.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if not self._exhausted:
            self._exhausted = True
            if exc_type:
                self._callback.status = "failed"
                self._callback.errors.append({
                    'timestamp': datetime.now(timezone.utc).isoformat(),
                    'error_type': exc_type.__name__,
                    'message': str(exc_val),
                })
                self._integration._complete_trace(
                    self._callback, 'failed', self._trace_name
                )
            else:
                self._integration._complete_trace(
                    self._callback, 'completed', self._trace_name
                )
            _active_callback_ctx.reset(self._token)

        if hasattr(self._stream_iter, '__aexit__'):
            return await self._stream_iter.__aexit__(exc_type, exc_val, exc_tb)
        return False