"""
Strands Agents hook handler for Aigie SDK.

Implements HookProvider to automatically trace Strands agent invocations,
tool calls, LLM calls, and multi-agent orchestrations.

Includes comprehensive error detection and drift monitoring.
"""

import logging
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Dict, List

from ...buffer import EventType
from .config import StrandsConfig
from .cost_tracking import calculate_strands_cost
from .drift_detection import DriftDetector
from .error_detection import DetectedError, ErrorDetector

logger = logging.getLogger(__name__)


def _utc_now() -> datetime:
    """Get current time in UTC with timezone info."""
    return datetime.now(timezone.utc)


def _utc_isoformat() -> str:
    """Get current UTC time as ISO format string."""
    return datetime.now(timezone.utc).isoformat()


if TYPE_CHECKING:
    try:
        from strands.agent.agent_result import AgentResult
        from strands.hooks import (
            AfterInvocationEvent,
            AfterModelCallEvent,
            AfterMultiAgentInvocationEvent,
            AfterNodeCallEvent,
            AfterToolCallEvent,
            BeforeInvocationEvent,
            BeforeModelCallEvent,
            BeforeMultiAgentInvocationEvent,
            BeforeNodeCallEvent,
            BeforeToolCallEvent,
            HookProvider,
            HookRegistry,
            MessageAddedEvent,
        )
        from strands.types.tools import ToolResult, ToolUse
    except ImportError:
        pass


class StrandsHandler:
    """
    Strands Agents handler for Aigie tracing.

    Implements HookProvider to automatically trace:
    - Agent invocations (BeforeInvocationEvent → AfterInvocationEvent)
    - Tool calls (BeforeToolCallEvent → AfterToolCallEvent)
    - LLM calls (BeforeModelCallEvent → AfterModelCallEvent)
    - Multi-agent orchestrations (BeforeMultiAgentInvocationEvent → AfterMultiAgentInvocationEvent)
    - Node executions (BeforeNodeCallEvent → AfterNodeCallEvent)

    Example:
        >>> from strands import Agent
        >>> from aigie.integrations.strands import StrandsHandler
        >>>
        >>> handler = StrandsHandler()
        >>> agent = Agent(tools=[...], hooks=[handler])
        >>> result = agent("What is the capital of France?")
    """

    def __init__(
        self,
        config: StrandsConfig | None = None,
        trace_name: str | None = None,
        metadata: Dict[str, Any] | None = None,
        tags: List[str] | None = None,
        user_id: str | None = None,
        session_id: str | None = None,
    ):
        """
        Initialize Strands handler.

        Args:
            config: Configuration for tracing behavior
            trace_name: Name for the trace (default: agent name)
            metadata: Additional metadata to attach
            tags: Tags to apply to trace and spans
            user_id: User ID for the trace
            session_id: Session ID for the trace
        """
        self.config = config or StrandsConfig.from_env()
        self.trace_name = trace_name
        self.metadata = metadata or {}
        self.tags = tags or []
        self.user_id = user_id
        self.session_id = session_id

        # State tracking
        self.trace_id: str | None = None
        self.agent_span_id: str | None = None
        self.tool_map: Dict[str, Dict[str, Any]] = {}  # tool_use_id -> {spanId, startTime}
        self.model_call_map: Dict[str, Dict[str, Any]] = {}  # model_span_id -> {startTime, ...}
        self.model_span_id: str | None = None  # Keep for backward compat, points to latest
        self.model_start_time: datetime | None = None  # Keep for backward compat
        self.multi_agent_map: Dict[
            str, Dict[str, Any]
        ] = {}  # orchestrator_id -> {spanId, startTime}
        self.node_map: Dict[str, Dict[str, Any]] = {}  # node_id -> {spanId, startTime}

        # Current context for parent relationships
        self._current_parent_span_id: str | None = None
        self._parent_span_stack: List[str] = []  # Stack for nested multi-agent/nodes
        self._aigie = None

        # Statistics
        self._total_tool_calls = 0
        self._total_input_tokens = 0
        self._total_output_tokens = 0
        self._total_cost = 0.0

        # Per-call token tracking - stores accumulated usage at start of each model call
        # This allows us to calculate delta (per-call) tokens
        self._model_call_start_tokens: Dict[str, int] | None = None

        # Deferred LLM span - queued in _on_after_model_call, created in next
        # _on_before_model_call or _on_after_invocation when tokens are available
        self._pending_llm_span: Dict[str, Any] | None = None
        self._llm_call_count: int = 0

        # Error tracking
        self._has_errors = False
        self._error_messages: List[str] = []

        # Error detection and monitoring
        self._error_detector = ErrorDetector()
        self._drift_detector = DriftDetector()
        self._detected_errors: List[DetectedError] = []
        self._invocation_start_time: datetime | None = None

        # Depth tracking for flow view - maps span_id to depth
        self._span_depth_map: Dict[str, int] = {}

        # Real-time remediation engine (shared across SDK integrations)
        self._remediation_engine = None
        self._intervention_dispatcher = None
        if self.config.enable_realtime_remediation:
            aigie = self._get_aigie()
            if aigie and aigie._initialized:
                api_url = getattr(aigie, "_api_url", None) or getattr(aigie, "api_url", None)
                api_key = getattr(aigie, "_api_key", None) or getattr(aigie, "api_key", None)
                if api_url:
                    from ...realtime.remediation_engine import RemediationEngine

                    self._remediation_engine = RemediationEngine(
                        api_url=api_url,
                        api_key=api_key or "",
                        query_timeout=self.config.remediation_query_timeout,
                    )
                # Grab gateway intervention dispatcher if available
                self._intervention_dispatcher = getattr(aigie, "_intervention_dispatcher", None)

    def _get_depth_for_parent(self, parent_id: str | None) -> int:
        """Calculate depth based on parent span's depth."""
        if not parent_id:
            return 0  # Root level
        parent_depth = self._span_depth_map.get(parent_id, 0)
        return parent_depth + 1

    def _register_span_depth(self, span_id: str, parent_id: str | None) -> int:
        """Register a span's depth and return it."""
        depth = self._get_depth_for_parent(parent_id)
        self._span_depth_map[span_id] = depth
        return depth

    def _get_aigie(self):
        """Lazy load Aigie client."""
        if self._aigie is None:
            from ...client import get_aigie

            self._aigie = get_aigie()
        return self._aigie

    def register_hooks(self, registry: "HookRegistry") -> None:
        """
        Register hook callbacks with the Strands hook registry.

        This method is called by Strands when the handler is added to an agent.

        Args:
            registry: The hook registry to register callbacks with
        """
        if not self.config.enabled:
            return

        try:
            from strands.hooks import (
                AfterInvocationEvent,
                AfterModelCallEvent,
                AfterMultiAgentInvocationEvent,
                AfterNodeCallEvent,
                AfterToolCallEvent,
                BeforeInvocationEvent,
                BeforeModelCallEvent,
                BeforeMultiAgentInvocationEvent,
                BeforeNodeCallEvent,
                BeforeToolCallEvent,
                MessageAddedEvent,
            )
        except ImportError:
            logger.warning("[AIGIE] Strands hooks not available - cannot register callbacks")
            return

        # Core agent lifecycle
        if self.config.trace_agents:
            registry.add_callback(BeforeInvocationEvent, self._on_before_invocation)
            registry.add_callback(AfterInvocationEvent, self._on_after_invocation)
            registry.add_callback(MessageAddedEvent, self._on_message_added)

        # Tool execution
        if self.config.trace_tools:
            registry.add_callback(BeforeToolCallEvent, self._on_before_tool_call)
            registry.add_callback(AfterToolCallEvent, self._on_after_tool_call)

        # Model invocation
        if self.config.trace_llm_calls:
            registry.add_callback(BeforeModelCallEvent, self._on_before_model_call)
            registry.add_callback(AfterModelCallEvent, self._on_after_model_call)

        # Multi-agent
        if self.config.trace_multi_agent:
            registry.add_callback(BeforeMultiAgentInvocationEvent, self._on_before_multi_agent)
            registry.add_callback(AfterMultiAgentInvocationEvent, self._on_after_multi_agent)
            registry.add_callback(BeforeNodeCallEvent, self._on_before_node_call)
            registry.add_callback(AfterNodeCallEvent, self._on_after_node_call)

        # BidiAgent streaming support (optional)
        if self.config.trace_streaming:
            self._register_streaming_hooks(registry)

    # Core agent lifecycle hooks

    async def _on_before_invocation(self, event: "BeforeInvocationEvent") -> None:
        """Handle BeforeInvocationEvent - create trace and agent span."""
        if not self.config.enabled or not self.config.trace_agents:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            # Check if we're already in a trace (from context or previous invocation)
            # This allows multiple agent invocations to share the same trace
            existing_trace_id = None

            # Check Aigie context for existing trace
            try:
                from ...auto_instrument.trace import get_current_trace

                current_trace = get_current_trace()
                if current_trace and hasattr(current_trace, "id"):
                    existing_trace_id = current_trace.id
            except Exception:
                pass  # Context not available, continue

            # If we already have a trace_id (from previous invocation or context), reuse it
            if existing_trace_id:
                self.trace_id = existing_trace_id
                trace_already_exists = True
            elif self.trace_id:
                # Reuse existing trace_id from handler (for nested invocations)
                trace_already_exists = True
            else:
                # Generate new trace ID only if we don't have one
                self.trace_id = str(uuid.uuid4())
                trace_already_exists = False

            # Reset invocation-specific state (but keep trace_id)
            self._has_errors = False
            self._error_messages = []
            self._total_tool_calls = 0
            self._total_input_tokens = 0
            self._total_output_tokens = 0
            self._total_cost = 0.0
            self.agent_span_id = None
            self.model_span_id = None
            self.model_start_time = None
            self._model_call_start_tokens = None  # Reset per-call token tracking
            self._pending_llm_span = None  # Deferred LLM span awaiting token data
            self._llm_call_count = 0  # Track model call count for token distribution
            self.tool_map.clear()
            self.model_call_map.clear()  # Clear model call tracking map
            self.multi_agent_map.clear()
            self.node_map.clear()
            self._current_parent_span_id = None
            self._parent_span_stack.clear()

            # Reset remediation engine state
            if self._remediation_engine:
                self._remediation_engine.reset()

            # Determine trace name
            agent_name = getattr(event.agent, "name", "Strands Agent")
            trace_name = self.trace_name or agent_name

            # Only create trace if it doesn't already exist
            if not trace_already_exists:
                # Create trace
                trace_data = {
                    "id": self.trace_id,
                    "name": trace_name,
                    "metadata": {
                        "framework": "strands",
                        "agent_id": getattr(event.agent, "agent_id", None),
                        "agent_name": agent_name,
                        **self.metadata,
                    },
                    "tags": self.tags,
                    "start_time": _utc_now().isoformat(),
                }

                if self.user_id:
                    trace_data["user_id"] = self.user_id
                if self.session_id:
                    trace_data["session_id"] = self.session_id

                # Create trace first
                if aigie._buffer:
                    await aigie._buffer.add(EventType.TRACE_CREATE, trace_data)

                # Set trace in context so other handlers can reuse it
                try:
                    from ...auto_instrument.trace import set_current_trace

                    # Create a simple trace context object with the trace_id
                    # This allows subsequent agent invocations to find and reuse the trace
                    class SimpleTraceContext:
                        def __init__(self, trace_id: str, trace_name: str):
                            self.id = trace_id
                            self.name = trace_name

                    set_current_trace(SimpleTraceContext(self.trace_id, trace_name))
                except Exception as e:
                    logger.debug(f"[AIGIE] Could not set trace in context: {e}")

                # Flush immediately to ensure trace is created before spans
                # This prevents orphan spans
                await aigie._buffer.flush()
            else:
                logger.debug(f"[AIGIE] Reusing existing trace: {self.trace_id}")

            # Create agent span at start so child spans (LLM, Tool) can reference it as parent
            # NOTE: We create the span here and complete it in _on_after_invocation
            # Unlike LLM/Tool spans which are created complete at the end, Agent spans need
            # to exist before children are created so they can reference the parent_id
            self.agent_span_id = str(uuid.uuid4())
            self._invocation_start_time = _utc_now()

            # Reset error/drift detection for this invocation
            self._error_detector = ErrorDetector()
            self._drift_detector = DriftDetector()
            self._detected_errors = []

            # Register agent span depth (root level = 0)
            agent_depth = self._register_span_depth(self.agent_span_id, None)

            # Capture system prompt and model for drift detection
            system_prompt = getattr(event.agent, "system_prompt", None)
            if system_prompt:
                self._drift_detector.capture_system_prompt(str(system_prompt))

            # Capture model ID in the plan
            if hasattr(event.agent, "model"):
                model_id = self._extract_model_id(event.agent.model)
                if model_id:
                    self._drift_detector.plan.model = model_id

            # Store agent span data for completion
            self._agent_span_data = {
                "agent_name": agent_name,
                "agent_id": getattr(event.agent, "agent_id", None),
                "depth": agent_depth,
            }

            # Create agent span (parent for LLM/Tool child spans)
            agent_span_data = {
                "id": self.agent_span_id,
                "trace_id": self.trace_id,
                "parent_id": None,  # Root span under trace
                "name": f"Agent: {agent_name}",
                "type": "agent",
                "start_time": self._invocation_start_time.isoformat(),
                "metadata": {
                    "framework": "strands",
                    "agent_id": getattr(event.agent, "agent_id", None),
                    "agent_name": agent_name,
                    "depth": agent_depth,
                },
                "tags": self.tags,
                "depth": agent_depth,
            }

            if self.user_id:
                agent_span_data["user_id"] = self.user_id
            if self.session_id:
                agent_span_data["session_id"] = self.session_id

            # Capture input if enabled
            if self.config.capture_inputs and event.messages:
                # Truncate messages if needed
                messages_repr = str(event.messages)
                if len(messages_repr) > self.config.max_content_length:
                    messages_repr = messages_repr[: self.config.max_content_length] + "..."
                agent_span_data["input"] = messages_repr
                self._agent_span_data["input"] = messages_repr

            if aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, agent_span_data)
            self._current_parent_span_id = self.agent_span_id

            # Subscribe to gateway push interventions for this trace
            if self._intervention_dispatcher and self.trace_id:
                self._intervention_dispatcher.subscribe_trace(self.trace_id)

            logger.debug(f"[AIGIE] Trace started: {trace_name} (id={self.trace_id})")

        except Exception as e:
            logger.error(f"[AIGIE] Error in _on_before_invocation: {e}", exc_info=True)

    async def _on_after_invocation(self, event: "AfterInvocationEvent") -> None:
        """Handle AfterInvocationEvent - complete agent span and trace."""
        if not self.config.enabled or not self.config.trace_agents:
            return

        if not self.agent_span_id:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            # Finalize any pending LLM span from the last model call
            # At this point, accumulated_usage is fully updated (result.metrics has final totals)
            if self._pending_llm_span:
                await self._finalize_pending_llm_span(event.agent)

            # Determine if invocation succeeded or failed
            result: AgentResult | None = event.result
            has_error = result is None or self._has_errors or len(self._error_messages) > 0

            # Extract metrics from result or agent
            # Strategy 1: Try result.metrics.accumulated_usage
            usage_found = False
            if result and hasattr(result, "metrics"):
                metrics = result.metrics
                if hasattr(metrics, "accumulated_usage") and metrics.accumulated_usage:
                    usage = metrics.accumulated_usage
                    self._total_input_tokens = usage.get("inputTokens", 0) or 0
                    self._total_output_tokens = usage.get("outputTokens", 0) or 0
                    usage_found = True
                    logger.debug(
                        f"[AIGIE] Got total tokens from result.metrics: input={self._total_input_tokens}, output={self._total_output_tokens}"
                    )

            # Strategy 2: Try agent.event_loop_metrics.accumulated_usage
            if not usage_found:
                accumulated = self._get_current_accumulated_usage(event.agent)
                if accumulated.get("inputTokens", 0) > 0 or accumulated.get("outputTokens", 0) > 0:
                    self._total_input_tokens = accumulated.get("inputTokens", 0)
                    self._total_output_tokens = accumulated.get("outputTokens", 0)
                    usage_found = True
                    logger.debug(
                        f"[AIGIE] Got total tokens from agent metrics: input={self._total_input_tokens}, output={self._total_output_tokens}"
                    )

            # Strategy 3: Try result.usage (some models return usage directly on result)
            if not usage_found and result:
                if hasattr(result, "usage") and result.usage:
                    usage = result.usage
                    if isinstance(usage, dict):
                        self._total_input_tokens = (
                            usage.get("inputTokens", 0) or usage.get("input_tokens", 0) or 0
                        )
                        self._total_output_tokens = (
                            usage.get("outputTokens", 0) or usage.get("output_tokens", 0) or 0
                        )
                    else:
                        self._total_input_tokens = (
                            getattr(usage, "inputTokens", 0)
                            or getattr(usage, "input_tokens", 0)
                            or 0
                        )
                        self._total_output_tokens = (
                            getattr(usage, "outputTokens", 0)
                            or getattr(usage, "output_tokens", 0)
                            or 0
                        )
                    if self._total_input_tokens > 0 or self._total_output_tokens > 0:
                        usage_found = True
                        logger.debug(
                            f"[AIGIE] Got total tokens from result.usage: input={self._total_input_tokens}, output={self._total_output_tokens}"
                        )

            # Log warning if no usage found
            if not usage_found:
                logger.warning("[AIGIE] Could not extract token usage from result or agent metrics")

            # Calculate cost
            if self._total_input_tokens > 0 or self._total_output_tokens > 0:
                model_id = None
                if hasattr(event.agent, "model"):
                    model_id = self._extract_model_id(event.agent.model)

                try:
                    self._total_cost = calculate_strands_cost(
                        model_id=model_id,
                        input_tokens=self._total_input_tokens,
                        output_tokens=self._total_output_tokens,
                    )
                    logger.debug(
                        f"[AIGIE] Calculated total cost: ${self._total_cost:.6f} for model={model_id}"
                    )
                except Exception as cost_err:
                    logger.warning(
                        f"[AIGIE] Cost calculation failed for model {model_id}: {cost_err}"
                    )
                    self._total_cost = 0.0

            # Determine status
            status = "error" if has_error else "success"
            error_message = None
            if self._error_messages:
                error_message = "; ".join(self._error_messages[:3])  # Limit to first 3 errors
            elif result is None:
                error_message = "Agent invocation returned no result"

            # Calculate duration
            agent_end_time = _utc_now()
            agent_duration_ms = 0
            if self._invocation_start_time:
                agent_duration_ms = (
                    agent_end_time - self._invocation_start_time
                ).total_seconds() * 1000

            # Re-create agent span as COMPLETE with all data (single SPAN_CREATE)
            # NOTE: The initial SPAN_CREATE in _on_before_invocation was needed so child
            # spans could reference parent_id. Now we send a complete replacement with
            # end_time, status, metrics, etc. The backend treats a second SPAN_CREATE
            # with the same ID as an upsert.
            agent_span_data = self._agent_span_data or {}
            total_tokens = self._total_input_tokens + self._total_output_tokens
            complete_agent_span = {
                "id": self.agent_span_id,
                "trace_id": self.trace_id,
                "parent_id": None,  # Root span
                "name": f"Agent: {agent_span_data.get('agent_name', 'Strands Agent')}",
                "type": "agent",
                "start_time": self._invocation_start_time.isoformat()
                if self._invocation_start_time
                else agent_end_time.isoformat(),
                "end_time": agent_end_time.isoformat(),
                "duration_ns": int(agent_duration_ms * 1_000_000),
                "status": status,
                "depth": agent_span_data.get("depth", 0),
                "tags": self.tags,
                "prompt_tokens": self._total_input_tokens,
                "completion_tokens": self._total_output_tokens,
                "total_tokens": total_tokens,
                "total_cost": self._total_cost,
                "token_usage": {
                    "prompt_tokens": self._total_input_tokens,
                    "completion_tokens": self._total_output_tokens,
                    "total_tokens": total_tokens,
                    "unit": "TOKENS",
                },
                "usage": {
                    "prompt_tokens": self._total_input_tokens,
                    "completion_tokens": self._total_output_tokens,
                    "total_tokens": total_tokens,
                    "input_tokens": self._total_input_tokens,
                    "output_tokens": self._total_output_tokens,
                },
                "metadata": {
                    "framework": "strands",
                    "agent_id": agent_span_data.get("agent_id"),
                    "agent_name": agent_span_data.get("agent_name", "Strands Agent"),
                    "depth": agent_span_data.get("depth", 0),
                    "total_tool_calls": self._total_tool_calls,
                    "total_input_tokens": self._total_input_tokens,
                    "total_output_tokens": self._total_output_tokens,
                    "total_tokens": total_tokens,
                    "total_cost": self._total_cost,
                    "duration_ms": agent_duration_ms,
                    "status": status,
                },
            }

            if self.user_id:
                complete_agent_span["user_id"] = self.user_id
            if self.session_id:
                complete_agent_span["session_id"] = self.session_id

            # Add input from initial capture
            if agent_span_data.get("input"):
                complete_agent_span["input"] = agent_span_data["input"]

            # Add error info
            if error_message:
                complete_agent_span["error"] = error_message
                complete_agent_span["error_message"] = error_message
                complete_agent_span["error_type"] = "AgentError"
                complete_agent_span["metadata"]["error"] = error_message
                complete_agent_span["metadata"]["error_type"] = "AgentError"

            # Capture output if enabled
            if self.config.capture_outputs and result:
                output_repr = str(result)
                if len(output_repr) > self.config.max_content_length:
                    output_repr = output_repr[: self.config.max_content_length] + "..."
                complete_agent_span["output"] = output_repr
                complete_agent_span["metadata"]["result"] = output_repr

            if aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, complete_agent_span)

            # Update trace with aggregated metrics
            # Note: We update the trace but don't set end_time if this is part of a larger workflow
            # The trace will be finalized when the entire workflow completes
            trace_output = {}
            if self.config.capture_outputs and result:
                output_repr = str(result)
                if len(output_repr) > self.config.max_content_length:
                    output_repr = output_repr[: self.config.max_content_length] + "..."
                trace_output["response"] = output_repr

            # Finalize drift detection and get all detected drifts
            end_time = _utc_now()
            duration_ms = 0
            if self._invocation_start_time:
                duration_ms = (end_time - self._invocation_start_time).total_seconds() * 1000
            total_tokens = self._total_input_tokens + self._total_output_tokens

            # Extract final output for drift detection
            final_output = None
            if self.config.capture_outputs and result:
                final_output = str(result)[:500]

            detected_drifts = self._drift_detector.finalize(
                total_duration_ms=duration_ms,
                total_tokens=total_tokens,
                total_cost=self._total_cost,
                final_output=final_output,
            )

            # Build monitoring data
            monitoring_data = {
                "drift_detection": {
                    "plan": self._drift_detector.plan.to_dict()
                    if self._drift_detector.plan
                    else None,
                    "execution": self._drift_detector.execution.to_dict()
                    if self._drift_detector.execution
                    else None,
                    "detected_drifts": [d.to_dict() for d in detected_drifts],
                    "drift_count": len(detected_drifts),
                },
                "error_detection": {
                    "stats": self._error_detector.stats.to_dict(),
                    "detected_errors": [e.to_dict() for e in self._detected_errors],
                    "error_count": len(self._detected_errors),
                },
            }

            # Log monitoring summary
            if detected_drifts:
                logger.info(
                    f"[AIGIE] Drift detection summary: {len(detected_drifts)} drifts detected"
                )
            if self._detected_errors:
                logger.info(
                    f"[AIGIE] Error detection summary: {len(self._detected_errors)} errors detected"
                )

            # Update trace with metrics from this invocation
            # Note: If multiple invocations share the same trace, each will update it
            # The backend should aggregate or the last update will be the final state
            trace_output["monitoring"] = {
                "drift_count": len(detected_drifts),
                "error_count": len(self._detected_errors),
                "plan_captured": self._drift_detector._plan_captured,
            }

            # Calculate trace duration
            trace_duration_ns = 0
            if self._invocation_start_time:
                trace_duration_ns = int(
                    (end_time - self._invocation_start_time).total_seconds() * 1_000_000_000
                )

            # Build execution plan summary
            agent_name = getattr(event.agent, "name", "Strands Agent")
            execution_plan = {
                "agent": agent_name,
                "model_calls": self._llm_call_count,
                "tool_calls": self._total_tool_calls,
                "turn_count": self._llm_call_count,
                "total_tokens": total_tokens,
                "total_cost": self._total_cost,
                "duration_ms": duration_ms,
                "status": status,
                "error_count": len(self._detected_errors),
                "drift_count": len(detected_drifts),
            }

            # Add plan data from drift detector if available
            plan_data = self._drift_detector.plan.to_dict() if self._drift_detector.plan else {}
            if plan_data.get("planned_steps"):
                execution_plan["planned_steps"] = plan_data["planned_steps"]
            if plan_data.get("expected_tools"):
                execution_plan["expected_tools"] = plan_data["expected_tools"]

            # Include name to restore it if backend auto-create overwrote it
            agent_name = agent_span_data.get("agent_name", "Strands Agent")
            trace_name = self.trace_name or agent_name

            trace_update = {
                "id": self.trace_id,
                "name": trace_name,
                "status": status,  # Status from this invocation
                "end_time": end_time.isoformat(),  # Update end_time on each invocation
                "duration_ns": trace_duration_ns,  # Duration in nanoseconds
                "output": trace_output,
                # Token/cost fields for backend aggregation display
                "total_tokens": self._total_input_tokens + self._total_output_tokens,
                "prompt_tokens": self._total_input_tokens,
                "completion_tokens": self._total_output_tokens,
                "total_cost": self._total_cost,
                "metadata": {
                    "total_tool_calls": self._total_tool_calls,
                    "total_input_tokens": self._total_input_tokens,
                    "total_output_tokens": self._total_output_tokens,
                    "total_tokens": self._total_input_tokens + self._total_output_tokens,
                    "total_cost": self._total_cost,
                    "last_agent": agent_name,
                    "turn_count": self._llm_call_count,
                    "execution_plan": execution_plan,
                    "monitoring": monitoring_data,
                    "realtime_remediation": {
                        "enabled": self.config.enable_realtime_remediation,
                        "mode": self.config.remediation_mode,
                        "applied_count": self._remediation_engine.applied_count
                        if self._remediation_engine
                        else 0,
                        "results": [
                            r.to_dict()
                            for r in (
                                self._remediation_engine.results[:10]
                                if self._remediation_engine
                                else []
                            )
                        ],
                    }
                    if self.config.enable_realtime_remediation
                    else None,
                },
            }

            # Add agent_plan to trace metadata if available
            if plan_data:
                trace_update["metadata"]["agent_plan"] = plan_data

            if error_message:
                trace_update["error"] = error_message
                trace_update["error_message"] = error_message

            if aigie._buffer:
                await aigie._buffer.add(EventType.TRACE_UPDATE, trace_update)

            # Clean up any pending spans (tools, models, nodes, multi-agent)
            await self._complete_pending_spans()

            # Flush everything, then re-send TRACE_CREATE to restore name.
            # Backend's _process_span_event_list auto-creates traces with "Auto-created..."
            # using upsert with name=EXCLUDED.name. A final TRACE_CREATE after all spans
            # ensures the correct name wins.
            await aigie._buffer.flush()
            await aigie._buffer.add(
                EventType.TRACE_CREATE, {"id": self.trace_id, "name": trace_name}
            )
            await aigie._buffer.flush()

            logger.debug(f"[AIGIE] Trace completed: {self.trace_id} (status={status})")

        except Exception as e:
            logger.error(f"[AIGIE] Error in _on_after_invocation: {e}", exc_info=True)
            # Still try to complete trace with error status
            try:
                if self.agent_span_id and self.trace_id:
                    now = _utc_now()
                    agent_span_data = getattr(self, "_agent_span_data", {}) or {}
                    error_span = {
                        "id": self.agent_span_id,
                        "trace_id": self.trace_id,
                        "parent_id": None,
                        "name": f"Agent: {agent_span_data.get('agent_name', 'Strands Agent')}",
                        "type": "agent",
                        "start_time": self._invocation_start_time.isoformat()
                        if self._invocation_start_time
                        else now.isoformat(),
                        "end_time": now.isoformat(),
                        "status": "error",
                        "error": str(e),
                        "error_message": str(e),
                    }
                    await aigie._buffer.add(EventType.SPAN_CREATE, error_span)

                    trace_error = {
                        "id": self.trace_id,
                        "status": "error",
                        "end_time": now.isoformat(),
                        "error": str(e),
                        "error_message": str(e),
                    }
                    await aigie._buffer.add(EventType.TRACE_UPDATE, trace_error)
            except Exception:
                pass  # Best effort
        finally:
            # Unsubscribe from gateway push interventions
            if self._intervention_dispatcher and self.trace_id:
                self._intervention_dispatcher.unsubscribe_trace(self.trace_id)
            # Clear span depth map to prevent memory leak in long-running agents
            self._span_depth_map.clear()

    async def _on_message_added(self, event: "MessageAddedEvent") -> None:
        """Handle MessageAddedEvent - track message additions."""
        if not self.config.enabled or not self.config.capture_messages:
            return

        # Messages are tracked as part of the agent span
        # This hook can be used for fine-grained message tracking if needed
        pass

    # Tool execution hooks

    async def _on_before_tool_call(self, event: "BeforeToolCallEvent") -> None:
        """Handle BeforeToolCallEvent - store data for complete span creation in _on_after_tool_call.

        NOTE: We don't create the span here. Instead, we create the COMPLETE span
        in _on_after_tool_call with a single SPAN_CREATE. This avoids relying on
        SPAN_UPDATE which the backend doesn't process correctly (end_time stays null).
        """
        if not self.config.enabled or not self.config.trace_tools:
            return

        if not self._current_parent_span_id or not self.trace_id:
            return

        try:
            tool_use: ToolUse = event.tool_use
            tool_use_id = tool_use.get("toolUseId", str(uuid.uuid4()))
            tool_name = tool_use.get("name", "unknown_tool")

            span_id = str(uuid.uuid4())
            start_time = _utc_now()

            # Calculate depth for flow view ordering
            tool_depth = self._register_span_depth(span_id, self._current_parent_span_id)

            # Store data for complete span creation in _on_after_tool_call
            self.tool_map[tool_use_id] = {
                "spanId": span_id,
                "startTime": start_time,
                "toolName": tool_name,
                "depth": tool_depth,
                "tool_input": tool_use.get("input", {}),
                "parentId": self._current_parent_span_id,
            }

            self._total_tool_calls += 1

            # Pre-tool interception: run local rules + gateway signal (observe only)
            aigie = self._get_aigie()
            if aigie and aigie._initialized:
                try:
                    intercept = await aigie.intercept_before_tool(
                        tool_name=tool_name,
                        tool_args=tool_use.get("input", {}),
                        trace_id=self.trace_id,
                        span_id=span_id,
                    )
                    if intercept.get("decision") != "allow":
                        logger.info(
                            f"[AIGIE] Pre-tool signal for {tool_name}: "
                            f"{intercept.get('decision')} — {intercept.get('reason')}"
                        )
                except Exception as ie:
                    logger.debug(f"[AIGIE] Pre-tool intercept error (non-fatal): {ie}")

            # Poll for pending gateway push intervention (non-blocking)
            if self._intervention_dispatcher and self.trace_id:
                signal = self._intervention_dispatcher.pop_pending(self.trace_id)
                if signal:
                    logger.info(
                        f"[AIGIE] Intervention received for {tool_name}: "
                        f"type={signal.intervention_type}, reason={signal.reason}"
                    )
                    import asyncio as _asyncio

                    result = await self._intervention_dispatcher.process(
                        signal, tool_name, tool_use.get("input", {})
                    )
                    if result.delay_ms:
                        await _asyncio.sleep(result.delay_ms / 1000.0)
                    self.tool_map[tool_use_id]["_pending_intervention"] = signal

            logger.debug(f"[AIGIE] Tool call started: {tool_name} (span_id={span_id})")

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_before_tool_call: {e}")

    async def _on_after_tool_call(self, event: "AfterToolCallEvent") -> None:
        """Handle AfterToolCallEvent - update tool span with completion data."""
        if not self.config.enabled or not self.config.trace_tools:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        # Extract tool_use_id before try block so it's accessible in finally
        tool_use: ToolUse = event.tool_use
        tool_use_id = tool_use.get("toolUseId")

        # If no toolUseId, try to find matching tool by name from recent entries
        if not tool_use_id:
            tool_name = tool_use.get("name", "unknown_tool")
            # Look for most recent tool with matching name
            for tid, tdata in reversed(list(self.tool_map.items())):
                if tdata.get("toolName") == tool_name:
                    tool_use_id = tid
                    break
            if not tool_use_id:
                logger.warning(
                    f"[AIGIE] No toolUseId found for tool {tool_name}, skipping duration tracking"
                )
                return

        try:
            tool_data = self.tool_map.get(tool_use_id)
            if not tool_data:
                return

            span_id = tool_data["spanId"]
            tool_name = tool_data["toolName"]

            # Determine status
            status = "success"
            is_error = False
            error_msg = None
            if event.exception:
                status = "error"
                is_error = True
                error_msg = str(event.exception)
            elif isinstance(event.result, dict) and event.result.get("status") == "error":
                # Strands sets result.status='error' when tool raises,
                # but may not set event.exception (tool errors are caught
                # and converted to error results for the LLM to handle)
                status = "error"
                is_error = True
                content = event.result.get("content", [])
                if content and isinstance(content, list):
                    for item in content:
                        if isinstance(item, dict) and "text" in item:
                            error_msg = item["text"]
                            break
                if not error_msg:
                    error_msg = str(event.result)
            elif event.cancel_message:
                status = "cancelled"

            if is_error:
                self._has_errors = True
                if error_msg and error_msg not in self._error_messages:
                    self._error_messages.append(error_msg)

            # Calculate duration
            start_time = tool_data["startTime"]
            end_time = _utc_now()
            duration = (end_time - start_time).total_seconds()
            duration_ms = duration * 1000

            # Error detection - check tool result for errors
            result_str = str(event.result) if event.result else ""
            detected_error = self._error_detector.detect_from_tool_result(
                tool_name=tool_name,
                tool_use_id=tool_use_id,
                result=result_str,
                is_error_flag=is_error,
                duration_ms=duration_ms,
            )
            if detected_error:
                self._detected_errors.append(detected_error)
                logger.debug(
                    f"[AIGIE] Error detected in tool {tool_name}: {detected_error.error_type.value}"
                )

            # Record for drift detection
            self._drift_detector.record_tool_use(
                tool_name=tool_name,
                tool_input=tool_data.get("tool_input", {}),
                duration_ms=duration_ms,
                is_error=is_error,
            )

            # Real-time remediation: if error detected and remediation enabled,
            # query platform for a matching fix and inject corrective guidance
            # into the tool result BEFORE it goes back to the LLM.
            if is_error and error_msg and self._remediation_engine:
                try:
                    rem_result = await self._remediation_engine.evaluate(
                        error_msg,
                        tool_name,
                        span_id,
                        self.trace_id or "",
                        mode=self.config.remediation_mode,
                    )
                    if rem_result and self.config.remediation_mode == "autonomous":
                        applied_programmatic = False
                        # Try programmatic fix via AutoFixApplicator first
                        if rem_result.action_type:
                            try:
                                fix_action = self._remediation_engine.to_fix_action(rem_result)
                                if fix_action:
                                    from ...interceptor.protocols import InterceptionContext
                                    from ...realtime.auto_fix import AutoFixApplicator

                                    applicator = AutoFixApplicator()
                                    ctx = InterceptionContext(
                                        provider="strands",
                                        model="",
                                        messages=[],
                                        response_content=error_msg,
                                    )
                                    fix_result = await applicator.apply_fixes(ctx, [fix_action])
                                    if fix_result.success:
                                        applied_programmatic = True
                                        logger.warning(
                                            f"[AIGIE] Programmatic fix applied: {rem_result.action_type} on {tool_name} "
                                            f"(strategy={fix_result.strategy.value}, "
                                            f"latency={fix_result.latency_ms:.0f}ms)"
                                        )
                                        if fix_result.modified_response:
                                            self._inject_guidance(
                                                event,
                                                str(
                                                    fix_result.modified_response.get("content", "")
                                                ),
                                            )
                            except Exception as fix_err:
                                logger.debug(
                                    f"[AIGIE] Programmatic fix failed, falling back to guidance: {fix_err}"
                                )

                        # Fall back to text guidance if no programmatic fix applied
                        if not applied_programmatic:
                            self._inject_guidance(event, rem_result.guidance_text)
                        self._remediation_engine.mark_applied(rem_result)
                        logger.warning(
                            f"[AIGIE] Real-time correction applied: {rem_result.error_type} on {tool_name} "
                            f"(strategy={rem_result.strategy}, rate={rem_result.success_rate:.0%}, "
                            f"query={rem_result.query_ms:.0f}ms, programmatic={applied_programmatic})"
                        )
                    elif rem_result:
                        logger.info(
                            f"[AIGIE] Remediation available for {rem_result.error_type} on {tool_name}: "
                            f"strategy={rem_result.strategy}, rate={rem_result.success_rate:.0%} "
                            f"(mode=recommendation, not applied)"
                        )
                    if rem_result:
                        await self._report_remediation(span_id, rem_result)
                except Exception as rem_err:
                    logger.debug(f"[AIGIE] Remediation failed (non-fatal): {rem_err}")

            # Post-tool interception: consult backend for additional fix signals
            if aigie and aigie._initialized and is_error:
                try:
                    post_intercept = await aigie.intercept_after_tool(
                        tool_name=tool_name,
                        result=event.result,
                        error=error_msg,
                        error_type=type(event.exception).__name__
                        if event.exception
                        else ("ToolError" if is_error else None),
                        trace_id=self.trace_id,
                        span_id=span_id,
                        duration_ms=duration_ms,
                    )
                    if post_intercept.get("fixes") and self.config.remediation_mode == "autonomous":
                        for fix in post_intercept["fixes"]:
                            if fix.get("confidence", 0) >= 0.7:
                                guidance = fix.get("reason") or fix.get("parameters", {}).get(
                                    "guidance"
                                )
                                if guidance:
                                    self._inject_guidance(event, guidance)
                                    logger.info(
                                        f"[AIGIE] Backend fix applied: {fix.get('action_type')} "
                                        f"(confidence={fix.get('confidence'):.0%})"
                                    )
                                    break
                except Exception as pie:
                    logger.debug(f"[AIGIE] Post-tool intercept error (non-fatal): {pie}")

            # Handle push interventions stored from _on_before_tool_call
            if tool_data and self.config.remediation_mode == "autonomous":
                pending = tool_data.get("_pending_intervention")
                if pending:
                    try:
                        if pending.intervention_type == "inject_correction":
                            guidance = pending.payload.get("guidance", pending.reason)
                            self._inject_guidance(event, guidance)
                        elif pending.intervention_type == "break_loop":
                            self._inject_guidance(
                                event,
                                f"[Aigie: Stop] {pending.reason}. Do not retry this action.",
                            )
                        elif pending.intervention_type == "redirect":
                            alt = pending.payload.get("alternative_tool", "")
                            self._inject_guidance(
                                event,
                                f"[Aigie: Use {alt} instead] {pending.reason}",
                            )
                        logger.info(
                            f"[AIGIE] Push intervention applied: type={pending.intervention_type}"
                        )
                    except Exception as int_err:
                        logger.debug(
                            f"[AIGIE] Intervention injection failed (non-fatal): {int_err}"
                        )

            # Create COMPLETE tool span with all data in a single SPAN_CREATE
            # This avoids the SPAN_UPDATE backend issue where end_time stays null
            span_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "parent_id": tool_data.get("parentId", self._current_parent_span_id),
                "name": f"Tool: {tool_name}",
                "type": "tool",
                "start_time": start_time.isoformat(),
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": status,
                "is_error": is_error,
                "depth": tool_data.get("depth", 1),
                "tags": self.tags,
                "metadata": {
                    "framework": "strands",
                    "tool_name": tool_name,
                    "tool_use_id": tool_use_id,
                    "depth": tool_data.get("depth", 1),
                    "duration_ms": duration_ms,
                    "status": status,
                },
            }

            if self.user_id:
                span_data["user_id"] = self.user_id
            if self.session_id:
                span_data["session_id"] = self.session_id

            # Capture tool input if enabled
            if self.config.capture_inputs:
                tool_input = tool_data.get("tool_input", {})
                input_repr = str(tool_input)
                if len(input_repr) > self.config.max_content_length:
                    input_repr = input_repr[: self.config.max_content_length] + "..."
                span_data["input"] = input_repr
                span_data["metadata"]["tool_input"] = input_repr

            # Capture tool output if enabled
            if self.config.capture_outputs:
                if event.result:
                    result_repr = str(event.result)
                    if len(result_repr) > self.config.max_tool_result_length:
                        result_repr = result_repr[: self.config.max_tool_result_length] + "..."
                    span_data["output"] = result_repr
                    span_data["metadata"]["tool_result"] = result_repr

            # Add error info
            if is_error and error_msg:
                span_data["error"] = error_msg
                span_data["error_message"] = error_msg
                span_data["metadata"]["error"] = error_msg
                if event.exception:
                    span_data["error_type"] = type(event.exception).__name__
                    span_data["metadata"]["error_type"] = type(event.exception).__name__
                else:
                    span_data["error_type"] = "ToolError"

            # Create complete span (single SPAN_CREATE with all data)
            logger.debug(f"[AIGIE] Creating complete Tool span: id={span_id}")
            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

            logger.debug(
                f"[AIGIE] Tool span created: {tool_name} (id={span_id}, status={status}, duration_ms={duration_ms:.2f})"
            )

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_after_tool_call: {e}")
        finally:
            # Always cleanup map entry to prevent memory leaks
            if tool_use_id in self.tool_map:
                del self.tool_map[tool_use_id]

    # Real-time remediation helpers (Strands-specific)

    def _inject_guidance(self, event: "AfterToolCallEvent", guidance_text: str) -> None:
        """Inject corrective guidance into the Strands tool result.

        Strands' AfterToolCallEvent.result is writable — we append guidance
        to the content list so the LLM sees both the error AND the fix.
        """
        try:
            result = event.result
            if isinstance(result, dict):
                content = result.get("content", [])
                if isinstance(content, list):
                    content.append({"text": guidance_text})
                else:
                    result["content"] = [{"text": str(content)}, {"text": guidance_text}]
                event.result = result
        except Exception as e:
            logger.warning(f"[AIGIE] Failed to apply remediation to event.result: {e}")

    async def _report_remediation(self, span_id: str, result) -> None:
        """Report a remediation result to the platform via SPAN_UPDATE metadata
        and POST to /remediation/results for closed-loop learning."""
        try:
            aigie = self._get_aigie()
            if not aigie or not aigie._initialized:
                return

            # Path 1: Span metadata (for trace detail view)
            if aigie._buffer:
                await aigie._buffer.add(
                    EventType.SPAN_UPDATE,
                    {
                        "id": span_id,
                        "trace_id": self.trace_id,
                        "metadata": {
                            "realtime_remediation": result.to_dict(),
                        },
                    },
                )

            # Path 2: POST to /remediation/results (shared engine handles this)
            if self._remediation_engine:
                await self._remediation_engine.report_result(result, self.trace_id)
        except Exception:
            pass  # Best-effort reporting

    # Model invocation hooks

    async def _on_before_model_call(self, event: "BeforeModelCallEvent") -> None:
        """Handle BeforeModelCallEvent - finalize previous pending LLM span and prepare for new call.

        KEY INSIGHT: Strands updates accumulated_usage AFTER firing AfterModelCallEvent.
        So in AfterModelCallEvent, the accumulated_usage doesn't yet include the current call.
        But by the time the NEXT BeforeModelCallEvent fires, accumulated_usage IS updated.

        Strategy:
        1. If there's a pending LLM span from the previous call, finalize it now
           (accumulated_usage now reflects the previous call's tokens)
        2. Capture current accumulated_usage as start_tokens for the new call
        """
        if not self.config.enabled or not self.config.trace_llm_calls:
            return

        if not self._current_parent_span_id or not self.trace_id:
            return

        try:
            # Step 0: Inject pending fix from async judge (quality correction from previous step)
            # The async judge may have found issues (hallucination, drift, etc.) in the
            # previous LLM response and queued a PendingFix. Inject it as system context
            # so the LLM can self-correct.
            try:
                aigie = self._get_aigie()
                if aigie and aigie._initialized and self.trace_id:
                    async_judge = getattr(aigie, "_async_judge_runner", None)
                    if async_judge:
                        fix = async_judge.pop_pending_fix(self.trace_id)
                        if fix and hasattr(event, "messages") and event.messages is not None:
                            # BeforeModelCallEvent.messages is read-only in Strands,
                            # so we log the fix for now. The actual injection happens
                            # via the agent's message history if writable.
                            logger.warning(
                                f"[AIGIE] Quality fix from judge (score={fix.judge_score:.2f}, "
                                f"cat={fix.judge_category}): {fix.fix_content[:200]}"
                            )
                            # Try to inject into agent's message list if accessible
                            if hasattr(event, "agent") and hasattr(event.agent, "messages"):
                                try:
                                    agent_msgs = event.agent.messages
                                    if isinstance(agent_msgs, list):
                                        agent_msgs.append(
                                            {"role": "user", "content": [{"text": fix.fix_content}]}
                                        )
                                        logger.info(
                                            "[AIGIE] Injected quality fix into agent messages"
                                        )
                                except Exception as inj_err:
                                    logger.debug(
                                        f"[AIGIE] Could not inject fix into messages: {inj_err}"
                                    )
            except Exception as fix_err:
                logger.debug(f"[AIGIE] Pending fix check skipped: {fix_err}")

            # Step 1: Finalize any pending LLM span from the previous model call
            # At this point, accumulated_usage has been updated with the previous call's tokens
            if self._pending_llm_span:
                await self._finalize_pending_llm_span(event.agent)

            # Step 2: Prepare for the new model call
            span_id = str(uuid.uuid4())
            start_time = _utc_now()
            self._llm_call_count += 1

            # Capture accumulated usage BEFORE this model call
            start_tokens = self._get_current_accumulated_usage(event.agent)

            # Get model info
            model_id = None
            if hasattr(event.agent, "model"):
                model_id = self._extract_model_id(event.agent.model)

            # Calculate depth for flow view ordering
            llm_depth = self._register_span_depth(span_id, self._current_parent_span_id)

            # Store in map for concurrent call support
            self.model_call_map[span_id] = {
                "startTime": start_time,
                "startTokens": start_tokens,
                "modelId": model_id,
                "depth": llm_depth,
                "parentId": self._current_parent_span_id,
            }

            # Keep backward compat references (points to latest)
            self.model_span_id = span_id
            self.model_start_time = start_time
            self._model_call_start_tokens = start_tokens

            logger.debug(
                f"[AIGIE] LLM call #{self._llm_call_count} started: model={model_id} (id={span_id})"
            )

        except Exception as e:
            logger.error(f"[AIGIE] Error in _on_before_model_call: {e}", exc_info=True)

    async def _on_after_model_call(self, event: "AfterModelCallEvent") -> None:
        """Handle AfterModelCallEvent - queue LLM span for deferred creation.

        KEY INSIGHT: Strands updates accumulated_usage AFTER this callback fires.
        So we can't calculate per-call tokens here (delta would be 0).

        Instead, we queue the span data and create it later:
        - In the next _on_before_model_call (accumulated_usage is now updated)
        - Or in _on_after_invocation (for the last model call)
        """
        if not self.config.enabled or not self.config.trace_llm_calls:
            return

        if not self.model_span_id:
            return

        span_id = self.model_span_id

        try:
            model_data = self.model_call_map.get(span_id, {})
            start_time = model_data.get("startTime") or self.model_start_time
            start_tokens = model_data.get("startTokens") or self._model_call_start_tokens
            parent_id = model_data.get("parentId") or self._current_parent_span_id
            llm_depth = model_data.get("depth", 1)

            # Determine status
            status = "success"
            is_error = False
            error_str = None
            if event.exception:
                status = "error"
                is_error = True
                self._has_errors = True
                error_str = str(event.exception)
                if error_str and error_str not in self._error_messages:
                    self._error_messages.append(error_str)

            # Extract model info
            model_id = model_data.get("modelId")
            if not model_id and hasattr(event.agent, "model"):
                model_id = self._extract_model_id(event.agent.model)

            end_time = _utc_now()

            # Try to get tokens directly from stop_response (works for some providers)
            response_usage = self._extract_usage_from_stop_response(event.stop_response)
            event_usage = self._extract_usage_from_event(event) if not response_usage else None

            # Capture structured output (JSON format matching other integrations)
            output_repr = None
            stop_reason = None
            if event.stop_response:
                stop_reason = str(event.stop_response.stop_reason)
                if self.config.capture_outputs:
                    response_text = str(event.stop_response.message)
                    if len(response_text) > self.config.max_content_length:
                        response_text = response_text[: self.config.max_content_length] + "..."
                    # Build structured output JSON
                    structured_output = {
                        "model": model_id,
                        "status": status,
                        "response": response_text,
                        "finish_reason": stop_reason,
                    }
                    # Add token usage if available from response
                    if response_usage:
                        structured_output["token_usage"] = {
                            "input_tokens": response_usage.get("inputTokens", 0),
                            "output_tokens": response_usage.get("outputTokens", 0),
                            "total_tokens": response_usage.get("inputTokens", 0)
                            + response_usage.get("outputTokens", 0),
                        }
                    output_repr = structured_output

            # Capture structured input (JSON format matching other integrations)
            input_repr = None
            if self.config.capture_inputs and hasattr(event, "messages") and event.messages:
                # Build structured input JSON
                prompts = []
                try:
                    for msg in event.messages:
                        if isinstance(msg, dict):
                            role = msg.get("role", "user")
                            content = msg.get("content", "")
                            if isinstance(content, list):
                                # Handle multi-part content (text + images)
                                text_parts = [
                                    p.get("text", "")
                                    for p in content
                                    if isinstance(p, dict) and "text" in p
                                ]
                                content = " ".join(text_parts) if text_parts else str(content)
                            prompts.append(
                                {
                                    "role": str(role),
                                    "content": str(content)[: self.config.max_content_length],
                                }
                            )
                        else:
                            prompts.append(
                                {
                                    "role": "user",
                                    "content": str(msg)[: self.config.max_content_length],
                                }
                            )
                except Exception:
                    prompts = [
                        {
                            "role": "user",
                            "content": str(event.messages)[: self.config.max_content_length],
                        }
                    ]

                structured_input = {
                    "model": model_id,
                    "prompts": prompts,
                }
                input_repr = structured_input

            # Queue the span data - will be created later when tokens are available
            self._pending_llm_span = {
                "span_id": span_id,
                "start_time": start_time,
                "end_time": end_time,
                "start_tokens": start_tokens,
                "model_id": model_id,
                "parent_id": parent_id,
                "depth": llm_depth,
                "status": status,
                "is_error": is_error,
                "error_str": error_str,
                "error_type": "LLMError" if is_error else None,
                "response_usage": response_usage,
                "event_usage": event_usage,
                "output": output_repr,
                "input": input_repr,
                "stop_reason": stop_reason,
            }

            logger.debug(
                f"[AIGIE] LLM span queued for deferred creation: {span_id} (model={model_id})"
            )

            # ── Quality evaluation: Fast Detect + Async Judge (non-blocking) ──
            # Type A errors (rate limit, timeout) are handled in _on_after_tool_call.
            # Here we handle Type B quality issues (hallucination, drift, lazy) on LLM output.
            try:
                aigie = self._get_aigie()
                if aigie and aigie._initialized:
                    output_text = None
                    input_msgs = None

                    # Extract output text for quality check
                    if output_repr and isinstance(output_repr, dict):
                        output_text = output_repr.get("response")
                    elif event.stop_response:
                        output_text = str(event.stop_response.message)

                    # Extract input messages for judge context
                    if hasattr(event, "messages") and event.messages:
                        input_msgs = event.messages

                    if output_text and self.trace_id:
                        # Step 1: Run fast detector for signal hints (<5ms, sync)
                        signal_hints = []
                        fast_detector = getattr(aigie, "_fast_detector", None)
                        if fast_detector:
                            try:
                                detection_result = fast_detector.detect(
                                    output=output_text,
                                    error=error_str,
                                    span_type="llm",
                                    metadata={"model": model_id},
                                )
                                signal_hints = getattr(detection_result, "signal_hints", [])
                                if signal_hints:
                                    logger.debug(
                                        f"[AIGIE] Quality signals detected: {signal_hints}"
                                    )
                            except Exception as fd_err:
                                logger.debug(f"[AIGIE] Fast detection skipped: {fd_err}")

                        # Step 2: Fire async judge with signal hints (background, non-blocking)
                        async_judge = getattr(aigie, "_async_judge_runner", None)
                        if async_judge and (signal_hints or is_error):
                            try:
                                async_judge.fire_and_forget(
                                    trace_id=self.trace_id,
                                    span_id=span_id,
                                    span_type="llm",
                                    output_content=output_text,
                                    input_messages=input_msgs,
                                    metadata={"model": model_id, "provider": "strands"},
                                    error=error_str,
                                    signal_hints=signal_hints,
                                )
                                logger.debug(
                                    f"[AIGIE] Async judge fired for LLM span {span_id} (signals={signal_hints})"
                                )
                            except Exception as aj_err:
                                logger.debug(f"[AIGIE] Async judge fire skipped: {aj_err}")
            except Exception as qe_err:
                logger.debug(f"[AIGIE] Quality evaluation skipped (non-fatal): {qe_err}")

        except Exception as e:
            logger.error(f"[AIGIE] Error in _on_after_model_call: {e}", exc_info=True)
        finally:
            # Clean up model span state
            if span_id in self.model_call_map:
                del self.model_call_map[span_id]
            self.model_span_id = None
            self.model_start_time = None
            self._model_call_start_tokens = None

    async def _finalize_pending_llm_span(self, agent: Any) -> None:
        """Finalize a pending LLM span by calculating tokens and creating it.

        Called from _on_before_model_call (for intermediate calls) or
        _on_after_invocation (for the last call). At these points,
        accumulated_usage has been updated with the call's tokens.
        """
        pending = self._pending_llm_span
        if not pending:
            return

        self._pending_llm_span = None  # Clear immediately to prevent double-finalization

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            span_id = pending["span_id"]
            start_time = pending["start_time"]
            end_time = pending["end_time"]
            model_id = pending["model_id"]
            parent_id = pending["parent_id"]
            llm_depth = pending["depth"]
            status = pending["status"]
            pending["is_error"]

            # Determine tokens - try direct extraction first, then delta
            model_input_tokens = 0
            model_output_tokens = 0

            response_usage = pending.get("response_usage")
            event_usage = pending.get("event_usage")

            if response_usage and (
                response_usage.get("inputTokens", 0) > 0
                or response_usage.get("outputTokens", 0) > 0
            ):
                model_input_tokens = response_usage.get("inputTokens", 0)
                model_output_tokens = response_usage.get("outputTokens", 0)
                logger.debug(
                    f"[AIGIE] LLM tokens from stop_response: in={model_input_tokens}, out={model_output_tokens}"
                )
            elif event_usage and (
                event_usage.get("inputTokens", 0) > 0 or event_usage.get("outputTokens", 0) > 0
            ):
                model_input_tokens = event_usage.get("inputTokens", 0)
                model_output_tokens = event_usage.get("outputTokens", 0)
                logger.debug(
                    f"[AIGIE] LLM tokens from event: in={model_input_tokens}, out={model_output_tokens}"
                )
            else:
                # Delta calculation - NOW accumulated_usage includes this call's tokens
                start_tokens = pending.get("start_tokens", {"inputTokens": 0, "outputTokens": 0})
                current_usage = self._get_current_accumulated_usage(agent)
                model_input_tokens = max(
                    0, current_usage["inputTokens"] - start_tokens.get("inputTokens", 0)
                )
                model_output_tokens = max(
                    0, current_usage["outputTokens"] - start_tokens.get("outputTokens", 0)
                )
                logger.debug(
                    f"[AIGIE] LLM tokens from delta: in={model_input_tokens}, out={model_output_tokens} "
                    f"(start={start_tokens}, current={current_usage})"
                )

            total_tokens = model_input_tokens + model_output_tokens

            # Calculate cost
            try:
                model_cost = calculate_strands_cost(
                    model_id=model_id,
                    input_tokens=model_input_tokens,
                    output_tokens=model_output_tokens,
                )
            except Exception:
                model_cost = 0.0
            input_cost = model_cost * 0.4 if model_cost else 0
            output_cost = model_cost * 0.6 if model_cost else 0

            # Calculate duration
            duration = 0.0
            if start_time:
                duration = (end_time - start_time).total_seconds()
            duration_ms = duration * 1000

            # Build span data
            metadata = {
                "service": "llm",
                "framework": "strands",
                "model": model_id,
                "model_id": model_id,
                "depth": llm_depth,
                "duration_ms": duration_ms,
                "status": status,
                "input_tokens": model_input_tokens,
                "output_tokens": model_output_tokens,
                "prompt_tokens": model_input_tokens,
                "completion_tokens": model_output_tokens,
                "total_tokens": total_tokens,
                "input_cost": input_cost,
                "output_cost": output_cost,
                "total_cost": model_cost,
                "cost": model_cost,
                "token_usage": {
                    "prompt_tokens": model_input_tokens,
                    "completion_tokens": model_output_tokens,
                    "total_tokens": total_tokens,
                    "input_tokens": model_input_tokens,
                    "output_tokens": model_output_tokens,
                    "input_cost": input_cost,
                    "output_cost": output_cost,
                    "total_cost": model_cost,
                },
            }

            span_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "parent_id": parent_id,
                "name": f"LLM: {model_id}" if model_id else "LLM",
                "type": "llm",
                "start_time": start_time.isoformat() if start_time else end_time.isoformat(),
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration_ms * 1_000_000),
                "status": status,
                "metadata": metadata,
                "token_usage": {
                    "prompt_tokens": model_input_tokens,
                    "completion_tokens": model_output_tokens,
                    "total_tokens": total_tokens,
                    "unit": "TOKENS",
                    "input_cost": input_cost,
                    "output_cost": output_cost,
                    "total_cost": model_cost,
                },
                "usage": {
                    "prompt_tokens": model_input_tokens,
                    "completion_tokens": model_output_tokens,
                    "total_tokens": total_tokens,
                    "input_tokens": model_input_tokens,
                    "output_tokens": model_output_tokens,
                    "input_cost": input_cost,
                    "output_cost": output_cost,
                    "total_cost": model_cost,
                },
                "prompt_tokens": model_input_tokens,
                "completion_tokens": model_output_tokens,
                "total_tokens": total_tokens,
                "input_cost": input_cost,
                "output_cost": output_cost,
                "total_cost": model_cost,
            }

            if model_id:
                span_data["model"] = model_id
            if self.user_id:
                span_data["user_id"] = self.user_id
            if self.session_id:
                span_data["session_id"] = self.session_id
            if pending.get("input"):
                span_data["input"] = pending["input"]
            if pending.get("output"):
                span_data["output"] = pending["output"]
            if pending.get("stop_reason"):
                span_data["metadata"]["stop_reason"] = pending["stop_reason"]
            if pending.get("error_str"):
                span_data["error"] = pending["error_str"]
                span_data["error_message"] = pending["error_str"]
                span_data["metadata"]["error"] = pending["error_str"]
                span_data["error_type"] = pending.get("error_type", "LLMError")
                span_data["metadata"]["error_type"] = pending.get("error_type", "LLMError")

            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)
            logger.debug(
                f"[AIGIE] LLM span created (deferred): {span_id} "
                f"(in={model_input_tokens}, out={model_output_tokens}, status={status})"
            )

        except Exception as e:
            logger.error(f"[AIGIE] Error finalizing pending LLM span: {e}", exc_info=True)

    # Multi-agent hooks

    async def _on_before_multi_agent(self, event: "BeforeMultiAgentInvocationEvent") -> None:
        """Handle BeforeMultiAgentInvocationEvent - create orchestrator span."""
        if not self.config.enabled or not self.config.trace_multi_agent:
            return

        if not self.trace_id:
            # If no trace_id, create one (for nested multi-agent scenarios)
            self.trace_id = str(uuid.uuid4())

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            orchestrator = event.source
            orchestrator_id = id(orchestrator)
            orchestrator_type = type(orchestrator).__name__

            span_id = str(uuid.uuid4())
            start_time = _utc_now()

            # Calculate depth for flow view ordering
            multi_agent_depth = self._register_span_depth(span_id, self._current_parent_span_id)

            self.multi_agent_map[orchestrator_id] = {
                "spanId": span_id,
                "startTime": start_time,
                "type": orchestrator_type,
                "depth": multi_agent_depth,
            }

            span_data = {
                "id": span_id,
                "trace_id": self.trace_id or str(uuid.uuid4()),
                "parent_id": self._current_parent_span_id,  # Use parent_id not parent_span_id
                "name": f"Multi-Agent: {orchestrator_type}",
                "type": "multi_agent",
                "start_time": start_time.isoformat(),
                "metadata": {
                    "orchestrator_type": orchestrator_type,
                    "orchestrator_id": str(orchestrator_id),
                    "depth": multi_agent_depth,
                },
                "depth": multi_agent_depth,  # For flow view ordering
                "tags": self.tags,
                "status": "running",
            }

            if self.user_id:
                span_data["user_id"] = self.user_id
            if self.session_id:
                span_data["session_id"] = self.session_id

            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

            # Store previous parent and push to stack
            if self._current_parent_span_id:
                self._parent_span_stack.append(self._current_parent_span_id)
            self._current_parent_span_id = span_id

            logger.debug(
                f"[AIGIE] Multi-agent orchestrator span created: {orchestrator_type} (id={span_id})"
            )

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_before_multi_agent: {e}")

    async def _on_after_multi_agent(self, event: "AfterMultiAgentInvocationEvent") -> None:
        """Handle AfterMultiAgentInvocationEvent - complete orchestrator span."""
        if not self.config.enabled or not self.config.trace_multi_agent:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        # Extract orchestrator_id before try block so it's accessible in finally
        orchestrator = event.source
        orchestrator_id = id(orchestrator)

        try:
            orchestrator_data = self.multi_agent_map.get(orchestrator_id)
            if not orchestrator_data:
                return

            span_id = orchestrator_data["spanId"]
            orchestrator_type = orchestrator_data["type"]
            start_time = orchestrator_data["startTime"]
            end_time = _utc_now()
            duration = (end_time - start_time).total_seconds()

            # Check for errors - AfterMultiAgentInvocationEvent might have error info
            status = "success"
            is_error = False
            error_message = None
            # Note: AfterMultiAgentInvocationEvent doesn't have exception field in Strands
            # but we track errors from child nodes/agents
            if self._has_errors:
                status = "error"
                is_error = True
                if self._error_messages:
                    error_message = "; ".join(self._error_messages[:2])

            # Calculate duration in ms
            duration_ms = duration * 1000

            update_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": status,
                "is_error": is_error,
                "metadata": {
                    "orchestrator_type": orchestrator_type,
                    "duration_ms": duration_ms,
                    "status": status,
                },
            }

            if error_message:
                update_data["error"] = error_message
                update_data["error_message"] = error_message
                update_data["metadata"]["error"] = error_message

            await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)

            # Restore previous parent from stack
            if self._parent_span_stack:
                self._current_parent_span_id = self._parent_span_stack.pop()
            else:
                # Fallback to agent span
                self._current_parent_span_id = self.agent_span_id

            logger.debug(
                f"[AIGIE] Multi-agent orchestrator span completed: {orchestrator_type} (id={span_id}, status={status})"
            )

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_after_multi_agent: {e}")
        finally:
            # Always cleanup map entry to prevent memory leaks
            if orchestrator_id in self.multi_agent_map:
                del self.multi_agent_map[orchestrator_id]

    async def _on_before_node_call(self, event: "BeforeNodeCallEvent") -> None:
        """Handle BeforeNodeCallEvent - create node span."""
        if not self.config.enabled or not self.config.trace_multi_agent:
            return

        if not self._current_parent_span_id or not self.trace_id:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            node_id = event.node_id
            orchestrator = event.source
            orchestrator_type = type(orchestrator).__name__

            span_id = str(uuid.uuid4())
            start_time = _utc_now()

            # Calculate depth for flow view ordering
            node_depth = self._register_span_depth(span_id, self._current_parent_span_id)

            self.node_map[node_id] = {
                "spanId": span_id,
                "startTime": start_time,
                "nodeId": node_id,
                "depth": node_depth,
            }

            span_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "parent_id": self._current_parent_span_id,  # Use parent_id not parent_span_id
                "name": f"Node: {node_id}",
                "type": "node",
                "start_time": start_time.isoformat(),
                "metadata": {
                    "node_id": node_id,
                    "orchestrator_type": orchestrator_type,
                    "depth": node_depth,
                },
                "tags": self.tags,
                "status": "running",
                "depth": node_depth,  # For flow view ordering
            }

            if self.user_id:
                span_data["user_id"] = self.user_id
            if self.session_id:
                span_data["session_id"] = self.session_id

            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

            logger.debug(f"[AIGIE] Node span created: {node_id} (id={span_id})")

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_before_node_call: {e}")

    async def _on_after_node_call(self, event: "AfterNodeCallEvent") -> None:
        """Handle AfterNodeCallEvent - complete node span."""
        if not self.config.enabled or not self.config.trace_multi_agent:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        # Extract node_id before try block so it's accessible in finally
        node_id = event.node_id

        try:
            node_data = self.node_map.get(node_id)
            if not node_data:
                return

            span_id = node_data["spanId"]
            start_time = node_data["startTime"]
            end_time = _utc_now()
            duration = (end_time - start_time).total_seconds()

            # Check for errors - AfterNodeCallEvent might have error info
            status = "success"
            is_error = False
            error_message = None
            # Note: AfterNodeCallEvent doesn't have exception field in Strands
            # but we track errors from child agents
            if self._has_errors:
                status = "error"
                is_error = True
                if self._error_messages:
                    error_message = "; ".join(self._error_messages[:2])

            # Calculate duration in ms
            duration_ms = duration * 1000

            update_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": status,
                "is_error": is_error,
                "metadata": {
                    "node_id": node_id,
                    "duration_ms": duration_ms,
                    "status": status,
                },
            }

            if error_message:
                update_data["error"] = error_message
                update_data["error_message"] = error_message
                update_data["metadata"]["error"] = error_message

            await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)

            logger.debug(f"[AIGIE] Node span completed: {node_id} (id={span_id}, status={status})")

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_after_node_call: {e}")
        finally:
            # Always cleanup map entry to prevent memory leaks
            if node_id in self.node_map:
                del self.node_map[node_id]

    # Helper methods

    def _get_current_accumulated_usage(self, agent: Any) -> Dict[str, int]:
        """
        Get the current accumulated token usage from the agent's event loop metrics.

        This is used to capture the "before" state so we can calculate per-call
        tokens by computing the delta after the model call completes.

        Args:
            agent: The Strands agent instance

        Returns:
            Dictionary with inputTokens and outputTokens (defaults to 0 if not available)
        """
        result = {"inputTokens": 0, "outputTokens": 0}
        try:
            if hasattr(agent, "event_loop_metrics"):
                metrics = agent.event_loop_metrics
                if hasattr(metrics, "accumulated_usage") and metrics.accumulated_usage:
                    usage = metrics.accumulated_usage
                    result["inputTokens"] = usage.get("inputTokens", 0) or 0
                    result["outputTokens"] = usage.get("outputTokens", 0) or 0
                    logger.debug(f"[AIGIE] Got accumulated usage from metrics: {result}")
                else:
                    logger.debug("[AIGIE] No accumulated_usage in metrics")
            else:
                logger.debug("[AIGIE] Agent has no event_loop_metrics attribute")
        except Exception as e:
            logger.debug(f"[AIGIE] Error getting accumulated usage: {e}")
        return result

    def _extract_usage_from_stop_response(self, stop_response: Any) -> Dict[str, int] | None:
        """
        Extract token usage directly from the stop_response.

        Many model providers include usage data in the response, which is more
        reliable than trying to compute deltas from accumulated metrics.

        Args:
            stop_response: The StopResponse from AfterModelCallEvent

        Returns:
            Dictionary with inputTokens and outputTokens, or None if not available
        """
        if not stop_response:
            return None

        try:
            # Try different common locations for usage data
            usage = None

            # Log stop_response structure for debugging
            logger.debug(f"[AIGIE] stop_response type: {type(stop_response)}")
            logger.debug(
                f"[AIGIE] stop_response attrs: {[a for a in dir(stop_response) if not a.startswith('_')]}"
            )

            # 1. Direct usage attribute
            if hasattr(stop_response, "usage") and stop_response.usage:
                usage = stop_response.usage
                logger.debug(f"[AIGIE] Found usage in stop_response.usage: {usage}")
            # 2. Usage in message
            elif hasattr(stop_response, "message"):
                msg = stop_response.message
                if hasattr(msg, "usage") and msg.usage:
                    usage = msg.usage
                    logger.debug(f"[AIGIE] Found usage in stop_response.message.usage: {usage}")
                # Gemini: check message content for usage_metadata
                elif hasattr(msg, "content") and msg.content:
                    content = msg.content
                    if isinstance(content, list) and len(content) > 0:
                        first_content = content[0]
                        if hasattr(first_content, "usage_metadata"):
                            usage = first_content.usage_metadata
                            logger.debug(
                                f"[AIGIE] Found usage in message.content[0].usage_metadata: {usage}"
                            )
            # 3. Response metadata
            elif hasattr(stop_response, "response") and stop_response.response:
                resp = stop_response.response
                if hasattr(resp, "usage") and resp.usage:
                    usage = resp.usage
                    logger.debug(f"[AIGIE] Found usage in stop_response.response.usage: {usage}")
                elif hasattr(resp, "usage_metadata") and resp.usage_metadata:
                    usage = resp.usage_metadata
                    logger.debug(
                        f"[AIGIE] Found usage in stop_response.response.usage_metadata: {usage}"
                    )

            # 4. Gemini-specific: check for usage in different locations
            # Strands may store usage data in stop_response.metrics
            if not usage and hasattr(stop_response, "metrics") and stop_response.metrics:
                metrics = stop_response.metrics
                if hasattr(metrics, "usage") and metrics.usage:
                    usage = metrics.usage
                    logger.debug(f"[AIGIE] Found usage in stop_response.metrics.usage: {usage}")
                elif isinstance(metrics, dict):
                    if "usage" in metrics:
                        usage = metrics["usage"]
                        logger.debug(
                            f"[AIGIE] Found usage in stop_response.metrics['usage']: {usage}"
                        )
                    elif "inputTokens" in metrics or "input_tokens" in metrics:
                        usage = metrics
                        logger.debug(
                            f"[AIGIE] Found usage directly in stop_response.metrics: {usage}"
                        )

            if usage:
                # Handle different usage formats
                if isinstance(usage, dict):
                    input_tokens = (
                        usage.get("inputTokens", 0)
                        or usage.get("input_tokens", 0)
                        or usage.get("prompt_tokens", 0)
                        or usage.get("prompt_token_count", 0)  # Gemini
                        or 0
                    )
                    output_tokens = (
                        usage.get("outputTokens", 0)
                        or usage.get("output_tokens", 0)
                        or usage.get("completion_tokens", 0)
                        or usage.get("candidates_token_count", 0)  # Gemini
                        or 0
                    )
                    if input_tokens > 0 or output_tokens > 0:
                        logger.debug(
                            f"[AIGIE] Extracted from dict: input={input_tokens}, output={output_tokens}"
                        )
                        return {"inputTokens": input_tokens, "outputTokens": output_tokens}
                else:
                    # Usage is an object with attributes
                    logger.debug(
                        f"[AIGIE] usage type: {type(usage)}, attrs: {[a for a in dir(usage) if not a.startswith('_')]}"
                    )
                    input_tokens = (
                        getattr(usage, "inputTokens", 0)
                        or getattr(usage, "input_tokens", 0)
                        or getattr(usage, "prompt_tokens", 0)
                        or getattr(usage, "prompt_token_count", 0)  # Gemini specific
                        or 0
                    )
                    output_tokens = (
                        getattr(usage, "outputTokens", 0)
                        or getattr(usage, "output_tokens", 0)
                        or getattr(usage, "completion_tokens", 0)
                        or getattr(usage, "candidates_token_count", 0)  # Gemini specific
                        or 0
                    )
                    if input_tokens > 0 or output_tokens > 0:
                        logger.debug(
                            f"[AIGIE] Extracted from object: input={input_tokens}, output={output_tokens}"
                        )
                        return {"inputTokens": input_tokens, "outputTokens": output_tokens}

        except Exception as e:
            logger.debug(f"[AIGIE] Could not extract usage from stop_response: {e}")

        return None

    def _extract_usage_from_event(self, event: Any) -> Dict[str, int] | None:
        """
        Extract token usage directly from the AfterModelCallEvent.

        Strands events may include usage data in various locations depending
        on the model provider.

        Args:
            event: The AfterModelCallEvent

        Returns:
            Dictionary with inputTokens and outputTokens, or None if not available
        """
        if not event:
            return None

        try:
            # Log event structure for debugging
            logger.debug(f"[AIGIE] Event type: {type(event)}")
            logger.debug(f"[AIGIE] Event attrs: {[a for a in dir(event) if not a.startswith('_')]}")

            # 1. Check for usage directly on event
            if hasattr(event, "usage") and event.usage:
                usage = event.usage
                logger.debug(f"[AIGIE] Found usage in event.usage: {usage}")
                if isinstance(usage, dict):
                    input_tokens = usage.get("inputTokens", 0) or usage.get("input_tokens", 0) or 0
                    output_tokens = (
                        usage.get("outputTokens", 0) or usage.get("output_tokens", 0) or 0
                    )
                    if input_tokens > 0 or output_tokens > 0:
                        return {"inputTokens": input_tokens, "outputTokens": output_tokens}

            # 2. Check for metrics on event
            if hasattr(event, "metrics") and event.metrics:
                metrics = event.metrics
                logger.debug(f"[AIGIE] Found metrics in event.metrics: {metrics}")
                if isinstance(metrics, dict):
                    if "usage" in metrics:
                        usage = metrics["usage"]
                        if isinstance(usage, dict):
                            input_tokens = (
                                usage.get("inputTokens", 0) or usage.get("input_tokens", 0) or 0
                            )
                            output_tokens = (
                                usage.get("outputTokens", 0) or usage.get("output_tokens", 0) or 0
                            )
                            if input_tokens > 0 or output_tokens > 0:
                                return {"inputTokens": input_tokens, "outputTokens": output_tokens}

            # 3. Check for token counts directly on event
            if hasattr(event, "input_tokens") or hasattr(event, "inputTokens"):
                input_tokens = (
                    getattr(event, "inputTokens", 0) or getattr(event, "input_tokens", 0) or 0
                )
                output_tokens = (
                    getattr(event, "outputTokens", 0) or getattr(event, "output_tokens", 0) or 0
                )
                if input_tokens > 0 or output_tokens > 0:
                    logger.debug(
                        f"[AIGIE] Found tokens on event: input={input_tokens}, output={output_tokens}"
                    )
                    return {"inputTokens": input_tokens, "outputTokens": output_tokens}

            # 4. Check agent on event for latest usage
            if hasattr(event, "agent") and event.agent:
                agent = event.agent
                # Try to get the most recent cycle's usage
                if hasattr(agent, "last_response_usage"):
                    usage = agent.last_response_usage
                    if usage:
                        logger.debug(f"[AIGIE] Found usage in agent.last_response_usage: {usage}")
                        if isinstance(usage, dict):
                            input_tokens = (
                                usage.get("inputTokens", 0) or usage.get("input_tokens", 0) or 0
                            )
                            output_tokens = (
                                usage.get("outputTokens", 0) or usage.get("output_tokens", 0) or 0
                            )
                            if input_tokens > 0 or output_tokens > 0:
                                return {"inputTokens": input_tokens, "outputTokens": output_tokens}

        except Exception as e:
            logger.debug(f"[AIGIE] Could not extract usage from event: {e}")

        return None

    def _get_latest_cycle_usage(self, agent: Any) -> Dict[str, int] | None:
        """
        Try to extract usage from the most recent event loop cycle.

        Strands tracks usage at the cycle level in:
        agent.event_loop_metrics.agent_invocations[-1].cycles[-1].usage

        This provides per-cycle (approximately per-model-call) token usage.

        Args:
            agent: The Strands agent instance

        Returns:
            Dictionary with per-cycle tokens, or None if not available
        """
        try:
            if not hasattr(agent, "event_loop_metrics"):
                return None

            metrics = agent.event_loop_metrics
            if not hasattr(metrics, "agent_invocations") or not metrics.agent_invocations:
                return None

            current_invocation = metrics.agent_invocations[-1]
            if not hasattr(current_invocation, "cycles") or not current_invocation.cycles:
                return None

            current_cycle = current_invocation.cycles[-1]
            if not hasattr(current_cycle, "usage") or not current_cycle.usage:
                return None

            usage = current_cycle.usage
            return {
                "inputTokens": usage.get("inputTokens", 0) or 0,
                "outputTokens": usage.get("outputTokens", 0) or 0,
                "totalTokens": usage.get("totalTokens", 0) or 0,
            }
        except Exception:
            return None

    def _calculate_per_call_tokens(
        self, agent: Any, start_tokens: Dict[str, int] | None
    ) -> Dict[str, int]:
        """
        Calculate per-call (delta) tokens by comparing current accumulated usage
        with the usage captured at the start of the model call.

        This approach handles the case where AfterModelCallEvent doesn't include
        per-call usage directly - we compute it from the accumulated totals.

        Falls back to cycle-level usage if delta calculation isn't possible.

        Args:
            agent: The Strands agent instance
            start_tokens: The accumulated usage captured at the start of the model call

        Returns:
            Dictionary with per-call inputTokens and outputTokens
        """
        current_usage = self._get_current_accumulated_usage(agent)

        # First, try the delta approach (most accurate for per-call)
        if start_tokens:
            # Calculate delta (per-call tokens)
            per_call_input = max(0, current_usage["inputTokens"] - start_tokens["inputTokens"])
            per_call_output = max(0, current_usage["outputTokens"] - start_tokens["outputTokens"])

            # If we got valid delta tokens, return them
            if per_call_input > 0 or per_call_output > 0:
                return {
                    "inputTokens": per_call_input,
                    "outputTokens": per_call_output,
                    "totalTokens": per_call_input + per_call_output,
                }

        # Fallback: Try to get cycle-level usage
        cycle_usage = self._get_latest_cycle_usage(agent)
        if cycle_usage:
            return cycle_usage

        # Last resort: Use current accumulated (may be inaccurate for per-call)
        return {
            "inputTokens": current_usage["inputTokens"],
            "outputTokens": current_usage["outputTokens"],
            "totalTokens": current_usage["inputTokens"] + current_usage["outputTokens"],
        }

    def _extract_model_id(self, model: Any) -> str | None:
        """Extract model ID from various model types."""
        if not model:
            return None

        # Try different ways to get model_id
        if hasattr(model, "model_id"):
            return model.model_id
        elif hasattr(model, "_model_id"):
            return model._model_id
        elif hasattr(model, "config"):
            if isinstance(model.config, dict):
                return model.config.get("model_id")
            elif hasattr(model.config, "model_id"):
                return model.config.model_id
        # For GeminiModel and similar, check client_args
        elif hasattr(model, "client_args") and isinstance(model.client_args, dict):
            # Fallback to model type name
            return type(model).__name__.replace("Model", "").lower()

        return None

    async def _complete_pending_spans(self) -> None:
        """Complete any pending spans that weren't explicitly closed."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return

        end_time = _utc_now()

        # Complete pending tool spans
        pending_tool_ids = list(self.tool_map.keys())
        for tool_use_id in pending_tool_ids:
            tool_data = self.tool_map.get(tool_use_id)
            if not tool_data:
                continue

            duration = (end_time - tool_data["startTime"]).total_seconds()
            update_data = {
                "id": tool_data["spanId"],
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": "success",  # Assume success if not explicitly failed
                "metadata": {
                    "tool_name": tool_data["toolName"],
                    "pending_cleanup": True,  # Mark as cleanup
                },
            }
            try:
                await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)
            except Exception:
                pass  # Best effort
            del self.tool_map[tool_use_id]

        # Complete pending model spans from map
        pending_model_span_ids = list(self.model_call_map.keys())
        for span_id in pending_model_span_ids:
            model_data = self.model_call_map.get(span_id)
            if not model_data:
                continue

            start_time = model_data.get("startTime")
            if start_time:
                duration = (end_time - start_time).total_seconds()
            else:
                duration = 0.0

            update_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": "success",  # Assume success if not explicitly failed
                "metadata": {
                    "model_id": model_data.get("modelId"),
                    "pending_cleanup": True,  # Mark as cleanup
                },
            }
            try:
                await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)
            except Exception:
                pass  # Best effort
            del self.model_call_map[span_id]

        # Clear backward compat model state
        self.model_span_id = None
        self.model_start_time = None

        # Complete pending multi-agent spans
        pending_orchestrator_ids = list(self.multi_agent_map.keys())
        for orchestrator_id in pending_orchestrator_ids:
            orchestrator_data = self.multi_agent_map.get(orchestrator_id)
            if not orchestrator_data:
                continue

            duration = (end_time - orchestrator_data["startTime"]).total_seconds()
            update_data = {
                "id": orchestrator_data["spanId"],
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": "success",  # Assume success if not explicitly failed
                "metadata": {
                    "orchestrator_type": orchestrator_data["type"],
                    "pending_cleanup": True,  # Mark as cleanup
                },
            }
            try:
                await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)
            except Exception:
                pass  # Best effort
            del self.multi_agent_map[orchestrator_id]

        # Complete pending node spans
        pending_node_ids = list(self.node_map.keys())
        for node_id in pending_node_ids:
            node_data = self.node_map.get(node_id)
            if not node_data:
                continue

            duration = (end_time - node_data["startTime"]).total_seconds()
            update_data = {
                "id": node_data["spanId"],
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": "success",  # Assume success if not explicitly failed
                "metadata": {
                    "node_id": node_id,
                    "pending_cleanup": True,  # Mark as cleanup
                },
            }
            try:
                await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)
            except Exception:
                pass  # Best effort
            del self.node_map[node_id]

    # ========================================================================
    # BidiAgent Streaming Hooks
    # ========================================================================

    def _register_streaming_hooks(self, registry: "HookRegistry") -> None:
        """
        Register BidiAgent streaming hooks with the registry.

        BidiAgent hooks enable real-time tracing of bidirectional streaming
        conversations, including partial results and interruptions.

        Args:
            registry: The hook registry to register callbacks with
        """
        try:
            # Try to import BidiAgent-specific events if available
            # These events may not be available in all Strands versions
            from strands.hooks import (
                BidiAfterInvocationEvent,
                BidiAfterToolCallEvent,
                BidiBeforeInvocationEvent,
                BidiBeforeToolCallEvent,
                BidiInterruptionEvent,
            )

            registry.add_callback(BidiBeforeInvocationEvent, self._on_bidi_before_invocation)
            registry.add_callback(BidiAfterInvocationEvent, self._on_bidi_after_invocation)
            registry.add_callback(BidiBeforeToolCallEvent, self._on_bidi_before_tool_call)
            registry.add_callback(BidiAfterToolCallEvent, self._on_bidi_after_tool_call)
            registry.add_callback(BidiInterruptionEvent, self._on_bidi_interruption)

            logger.debug("[AIGIE] BidiAgent streaming hooks registered")

        except ImportError:
            logger.debug("[AIGIE] BidiAgent events not available - streaming hooks disabled")
        except Exception as e:
            logger.warning(f"[AIGIE] Failed to register BidiAgent hooks: {e}")

    async def _on_bidi_before_invocation(self, event: Any) -> None:
        """Handle BidiBeforeInvocationEvent - start streaming invocation span."""
        if not self.config.enabled or not self.config.trace_streaming:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            # Create streaming invocation span
            span_id = str(uuid.uuid4())
            start_time = _utc_now()

            # Use agent span as parent if available
            parent_id = self.agent_span_id or self._current_parent_span_id

            span_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "parent_id": parent_id,
                "name": "BidiAgent Streaming",
                "type": "llm",
                "start_time": start_time.isoformat(),
                "metadata": {
                    "streaming": True,
                    "bidi_agent": True,
                },
                "status": "running",
            }

            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

            # Store for later completion
            if not hasattr(self, "_bidi_span_id"):
                self._bidi_span_id = None
            if not hasattr(self, "_bidi_start_time"):
                self._bidi_start_time = None

            self._bidi_span_id = span_id
            self._bidi_start_time = start_time

            logger.debug(f"[AIGIE] BidiAgent streaming started (span_id={span_id[:8]})")

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_bidi_before_invocation: {e}")

    async def _on_bidi_after_invocation(self, event: Any) -> None:
        """Handle BidiAfterInvocationEvent - complete streaming invocation span."""
        if not self.config.enabled or not self.config.trace_streaming:
            return

        if not hasattr(self, "_bidi_span_id") or not self._bidi_span_id:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            end_time = _utc_now()
            duration = 0.0
            if hasattr(self, "_bidi_start_time") and self._bidi_start_time:
                duration = (end_time - self._bidi_start_time).total_seconds()

            # Extract streaming result info
            output = None
            status = "success"
            if hasattr(event, "result"):
                if hasattr(event.result, "text"):
                    output = str(event.result.text)[: self.config.max_content_length]
                elif event.result is None:
                    status = "cancelled"

            update_data = {
                "id": self._bidi_span_id,
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": status,
                "output": output,
            }

            await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)

            logger.debug(
                f"[AIGIE] BidiAgent streaming completed (span_id={self._bidi_span_id[:8]})"
            )

            self._bidi_span_id = None
            self._bidi_start_time = None

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_bidi_after_invocation: {e}")

    async def _on_bidi_before_tool_call(self, event: Any) -> None:
        """Handle BidiBeforeToolCallEvent - track streaming tool call."""
        if not self.config.enabled or not self.config.trace_streaming:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            # Extract tool info from event
            tool_name = getattr(event, "tool_name", "unknown_tool")
            tool_use_id = getattr(event, "tool_use_id", str(uuid.uuid4()))

            span_id = str(uuid.uuid4())
            start_time = _utc_now()

            # Use bidi span as parent if available
            parent_id = getattr(self, "_bidi_span_id", None) or self.agent_span_id

            span_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "parent_id": parent_id,
                "name": f"Tool: {tool_name}",
                "type": "tool",
                "start_time": start_time.isoformat(),
                "metadata": {
                    "tool_name": tool_name,
                    "tool_use_id": tool_use_id,
                    "streaming": True,
                },
                "status": "running",
            }

            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

            # Store in tool_map for completion
            self.tool_map[tool_use_id] = {
                "spanId": span_id,
                "startTime": start_time,
                "toolName": tool_name,
                "streaming": True,
            }

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_bidi_before_tool_call: {e}")

    async def _on_bidi_after_tool_call(self, event: Any) -> None:
        """Handle BidiAfterToolCallEvent - complete streaming tool call span."""
        if not self.config.enabled or not self.config.trace_streaming:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            tool_use_id = getattr(event, "tool_use_id", None)
            if not tool_use_id or tool_use_id not in self.tool_map:
                return

            tool_data = self.tool_map[tool_use_id]
            end_time = _utc_now()
            duration = (end_time - tool_data["startTime"]).total_seconds()

            # Extract result
            output = None
            status = "success"
            if hasattr(event, "result"):
                output = str(event.result)[: self.config.max_tool_result_length]
            if hasattr(event, "error") and event.error:
                status = "error"
                output = str(event.error)

            update_data = {
                "id": tool_data["spanId"],
                "trace_id": self.trace_id,
                "end_time": end_time.isoformat(),
                "duration_ns": int(duration * 1_000_000_000),
                "status": status,
                "output": output,
            }

            await aigie._buffer.add(EventType.SPAN_UPDATE, update_data)
            del self.tool_map[tool_use_id]

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_bidi_after_tool_call: {e}")

    async def _on_bidi_interruption(self, event: Any) -> None:
        """Handle BidiInterruptionEvent - track streaming interruptions."""
        if not self.config.enabled or not self.config.trace_streaming:
            return

        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        try:
            # Create interruption span
            span_id = str(uuid.uuid4())
            timestamp = _utc_now()

            parent_id = getattr(self, "_bidi_span_id", None) or self.agent_span_id

            # Extract interruption reason
            reason = getattr(event, "reason", "unknown")
            source = getattr(event, "source", "user")

            span_data = {
                "id": span_id,
                "trace_id": self.trace_id,
                "parent_id": parent_id,
                "name": f"Interruption: {reason}",
                "type": "event",
                "start_time": timestamp.isoformat(),
                "end_time": timestamp.isoformat(),
                "duration_ns": 0,
                "metadata": {
                    "interruption": True,
                    "reason": reason,
                    "source": source,
                },
                "status": "cancelled",
            }

            await aigie._buffer.add(EventType.SPAN_CREATE, span_data)

            logger.debug(f"[AIGIE] BidiAgent interruption recorded (reason={reason})")

        except Exception as e:
            logger.warning(f"[AIGIE] Error in _on_bidi_interruption: {e}")
