"""Strands Agents integration using Beacon SDK Spans (callback handler approach).

This module provides a callback handler that traces Strands Agent executions
using SDK Span objects, which are exported via OTLP.

Usage:
    from lumenova_beacon import BeaconClient
    from lumenova_beacon.tracing.integrations.strands import BeaconStrandsHandler
    from strands import Agent

    # Initialize BeaconClient first
    client = BeaconClient()

    # Create callback handler
    handler = BeaconStrandsHandler(
        session_id="my-session",
        agent_name="My Agent",
    )

    # Use with Strands Agent
    agent = Agent(model=model, callback_handler=handler)
    result = agent("Hello, world!")
    print(handler.trace_id)  # Link to Beacon trace

Span hierarchy produced:
    agent_span  (SpanType.AGENT)
    ├── model_span  (SpanType.GENERATION) - each event loop cycle
    ├── tool_span A  (SpanType.TOOL) - message(assistant,toolUse) → message(user,toolResult)
    └── tool_span B  (SpanType.TOOL)

Event flow:
    init_event_loop  → create agent span
    start_event_loop → create model span (+ close previous tool spans as safety net)
    data             → accumulate streaming text; capture model name + input messages once
    reasoningText    → accumulate reasoning text
    current_tool_use → accumulate streaming tool call info (partial JSON input)
    message(assistant) → finalize model span; create tool spans (input already parsed)
    message(user,toolResult) → close matching tool spans with output
    result           → finalize agent span with tokens + output
    force_stop       → close all spans with ERROR status

Enable debug logging to see every raw event Strands sends:
    import logging
    logging.getLogger("lumenova_beacon.tracing.integrations.strands").setLevel(logging.DEBUG)

Environment Variables:
    BEACON_ENDPOINT - Beacon API endpoint (for OTLP export)
    BEACON_API_KEY  - Beacon API key (for OTLP export)
    BEACON_SESSION_ID - Default session ID (optional)
"""

import json
import logging
import time
from typing import Any

from lumenova_beacon.core.client import get_client
from lumenova_beacon.tracing.span import Span
from lumenova_beacon.tracing.trace import set_current_span, clear_context
from lumenova_beacon.types import SpanType, SpanKind, StatusCode

try:
    import strands  # noqa: F401 - validates strands-agents is installed
except ImportError:
    raise ImportError(
        "strands-agents is required for Strands Agents integration. "
        "Install it with: pip install 'lumenova-beacon[strands]'"
    )

logger = logging.getLogger(__name__)


class BeaconStrandsHandler:
    """Beacon callback handler for Strands Agents.

    Pass an instance of this class as the ``callback_handler`` argument when
    constructing a Strands ``Agent``. The handler automatically creates and
    exports Beacon spans covering the full agent lifecycle.

    Args:
        session_id: Override session ID for all spans (defaults to
            ``BeaconClient.config.session_id``).
        environment: Deployment environment tag stored as span metadata.
        agent_name: Human-readable name for the agent span and
            ``gen_ai.agent.name`` attribute.
        metadata: Additional key-value pairs stored as ``beacon.metadata.*``
            attributes on the agent span.
    """

    # OTEL GenAI semantic convention attributes (aligned with LangChain handler)
    ATTR_GEN_AI_REQUEST_MODEL = "gen_ai.request.model"
    ATTR_GEN_AI_RESPONSE_MODEL = "gen_ai.response.model"
    ATTR_GEN_AI_INPUT_TOKENS = "gen_ai.usage.input_tokens"
    ATTR_GEN_AI_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
    ATTR_GEN_AI_CACHE_READ_TOKENS = "gen_ai.usage.cache_read_input_tokens"
    ATTR_GEN_AI_CACHE_CREATION_TOKENS = "gen_ai.usage.cache_creation_input_tokens"
    ATTR_GEN_AI_REASONING_TOKENS = "gen_ai.usage.reasoning_tokens"
    ATTR_GEN_AI_INPUT_MESSAGES = "gen_ai.input.messages"
    ATTR_GEN_AI_OUTPUT_MESSAGES = "gen_ai.output.messages"
    ATTR_GEN_AI_TOOL_NAME = "gen_ai.tool.name"
    ATTR_GEN_AI_TOOL_TYPE = "gen_ai.tool.type"
    ATTR_GEN_AI_TOOL_CALL_ARGUMENTS = "gen_ai.tool.call.arguments"
    ATTR_GEN_AI_TOOL_CALL_RESULT = "gen_ai.tool.call.result"
    ATTR_GEN_AI_TOOL_CALL_ID = "gen_ai.tool.call.id"
    ATTR_GEN_AI_AGENT_NAME = "gen_ai.agent.name"
    ATTR_GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id"
    ATTR_GEN_AI_TTFT_MS = "gen_ai.response.time_to_first_token_ms"
    ATTR_GEN_AI_REASONING_CONTENT = "gen_ai.reasoning_content"
    ATTR_GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
    ATTR_GEN_AI_AGENT_FRAMEWORK = "gen_ai.agent.framework"
    ATTR_GEN_AI_PROVIDER_NAME = "gen_ai.provider.name"
    ATTR_STRANDS_CYCLE = "strands.event_loop.cycle"

    def __init__(
        self,
        session_id: str | None = None,
        user_email: str | None = None,
        user_name: str | None = None,
        user_id: str | None = None,
        environment: str | None = None,
        agent_name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ):
        self._client = get_client()
        self._session_id = session_id or self._client.config.session_id
        self._user_email = user_email if user_email is not None else self._client.config.user_email
        self._user_name = user_name if user_name is not None else self._client.config.user_name
        self._user_id = user_id if user_id is not None else self._client.config.user_id
        self._environment = environment if environment is not None else self._client.config.environment
        self._record_stacktrace = self._client.config.record_stacktrace
        self._agent_name = agent_name
        self._metadata = metadata or {}

        # Span state
        self._agent_span: Span | None = None
        self._model_span: Span | None = None
        self._trace_id: str | None = None

        # Tool tracking
        # toolUseId → {name, id, input (parsed dict or None)}
        self._pending_tool_calls: dict[str, dict[str, Any]] = {}
        # toolUseId → Span, created after complete, closed on tool result message
        self._active_tool_spans: dict[str, Span] = {}

        # Streaming state per cycle
        self._streaming_buffer: list[str] = []
        self._reasoning_buffer: list[str] = []
        self._first_token_time: float | None = None
        self._loop_start_time: float | None = None
        self._cycle_count: int = 0

        # Message tracking for input/output on model spans
        # Stores the last user-initiated message (not a tool result)
        self._pending_input_messages: list[dict[str, Any]] = []

        # Full assistant message content blocks (text + toolUse) captured from
        # message(assistant) events — used to build structured gen_ai.output.messages
        self._assistant_message_content: list[dict[str, Any]] | None = None

        # Flag: have we already captured model name + input messages from a data event this cycle?
        self._model_info_captured: bool = False

        # Flag: have we already set the agent span input?
        self._agent_input_set: bool = False

    @property
    def trace_id(self) -> str | None:
        """The trace ID for the current (or most recent) agent invocation."""
        return self._trace_id

    def __call__(self, **kwargs: Any) -> None:
        """Handle a Strands callback event.

        Strands passes all events as keyword arguments to this callable.
        Events are identified by the presence of specific keys.

        Lifecycle events (mutually exclusive with each other) are handled
        first. Content events are checked independently because the Strands
        SDK may send ``complete=True`` together with the last ``data`` chunk.

        Note: ``data`` and ``delta`` represent the same streamed text and are
        treated as mutually exclusive — ``delta`` is only processed when
        ``data`` is absent.
        """
        logger.debug("BeaconStrandsHandler event keys: %s", list(kwargs.keys()))

        try:
            # Exclusive lifecycle events
            if "init_event_loop" in kwargs:
                self._on_init_event_loop(**kwargs)
                return
            if "start_event_loop" in kwargs:
                self._on_start_event_loop(**kwargs)
                return
            if "result" in kwargs:
                self._on_result(kwargs["result"])
                return
            if "force_stop" in kwargs:
                self._on_force_stop(**kwargs)
                return

            # Content events — ``data`` and ``delta`` are mutually exclusive
            # (Strands sometimes sends both for the same chunk; we take ``data`` first)
            # Pass full kwargs to ``data`` events so we can extract model/messages metadata.
            if "data" in kwargs:
                logger.debug("BeaconStrandsHandler data: %r", kwargs["data"])
                self._on_data(kwargs["data"], kwargs)
            elif "delta" in kwargs:
                logger.debug("BeaconStrandsHandler delta: %r", kwargs["delta"])
                self._on_data(kwargs["delta"], None)

            if "reasoningText" in kwargs:
                logger.debug("BeaconStrandsHandler reasoningText: %r", kwargs["reasoningText"])
                self._on_reasoning_text(kwargs["reasoningText"])

            if "current_tool_use" in kwargs:
                logger.debug("BeaconStrandsHandler current_tool_use: %r", kwargs["current_tool_use"])
                self._on_current_tool_use(kwargs["current_tool_use"])

            if "message" in kwargs:
                logger.debug("BeaconStrandsHandler message: %r", kwargs["message"])
                self._on_message(kwargs["message"])

            if "complete" in kwargs:
                logger.debug("BeaconStrandsHandler complete: %r", kwargs["complete"])
                self._on_complete(**kwargs)

        except Exception as exc:
            logger.warning("BeaconStrandsHandler encountered an error", exc_info=True)
            self._emergency_close(exc)

    # ------------------------------------------------------------------
    # Event handlers
    # ------------------------------------------------------------------

    def _on_init_event_loop(self, **kwargs: Any) -> None:
        """Agent execution has started — create the root agent span."""
        logger.debug("BeaconStrandsHandler _on_init_event_loop")
        span = Span(
            name=self._agent_name or "strands-agent",
            span_type=SpanType.AGENT,
            kind=SpanKind.INTERNAL,
            session_id=self._session_id,
            user_email=self._user_email,
            user_name=self._user_name,
            user_id=self._user_id,
            record_stacktrace=self._record_stacktrace,
        )
        span.start()
        self._agent_span = span
        self._trace_id = span.trace_id

        span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "strands")
        span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "invoke_agent")
        span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "strands")
        if self._agent_name:
            span.set_attribute(self.ATTR_GEN_AI_AGENT_NAME, self._agent_name)
        if self._session_id:
            span.set_attribute(self.ATTR_GEN_AI_CONVERSATION_ID, self._session_id)
        if self._environment:
            span.set_attribute("deployment.environment.name", self._environment)
        for key, value in self._metadata.items():
            span.set_attribute(f"beacon.metadata.{key}", str(value))

        set_current_span(span)
        self._client.export_eager_span(span)

    def _on_start_event_loop(self, **kwargs: Any) -> None:
        """A new model call cycle is beginning.

        If any tool spans are still open (i.e. no ``message`` event arrived
        with the tool result), close them here as a safety net.
        """
        logger.debug("BeaconStrandsHandler _on_start_event_loop (cycle %d)", self._cycle_count + 1)

        # Safety-net: close any tool spans not yet closed by a message event
        self._end_active_tool_spans(status=StatusCode.OK)

        # Safety-net: close any model span from the previous cycle
        if self._model_span:
            self._finalize_model_span()

        # Reset per-cycle state
        self._cycle_count += 1
        self._loop_start_time = time.perf_counter()
        self._streaming_buffer = []
        self._reasoning_buffer = []
        self._first_token_time = None
        self._pending_tool_calls = {}
        self._model_info_captured = False
        self._assistant_message_content = None

        # Create model call span
        span = Span(
            name="model-call",
            trace_id=self._trace_id,
            parent_id=self._agent_span.span_id if self._agent_span else None,
            span_type=SpanType.GENERATION,
            kind=SpanKind.CLIENT,
            session_id=self._session_id,
            user_email=self._user_email,
            user_name=self._user_name,
            user_id=self._user_id,
            record_stacktrace=self._record_stacktrace,
        )
        span.start()
        span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "chat")
        span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "strands")
        span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "strands")
        span.set_attribute(self.ATTR_STRANDS_CYCLE, self._cycle_count)

        # Attach any input messages collected before this cycle
        if self._pending_input_messages:
            messages_json = json.dumps(self._pending_input_messages, default=str)
            span.set_attribute(self.ATTR_GEN_AI_INPUT_MESSAGES, messages_json)
            span.set_input(self._pending_input_messages)
            self._pending_input_messages = []

        self._model_span = span
        set_current_span(span)
        self._client.export_eager_span(span)

    def _on_data(self, data: Any, event_kwargs: dict[str, Any] | None = None) -> None:
        """A streaming text chunk arrived from the model.

        Strands normally sends a plain ``str``, but with extended thinking or
        certain model backends it may pass a raw event ``dict`` containing the
        text under a ``"text"`` key.  We normalise both forms here.

        ``event_kwargs`` carries the full callback kwargs. For ``data`` events
        (TextStreamEvent), Strands merges ``invocation_state`` which includes
        the ``agent`` object. We extract model name and conversation messages
        from the agent once per cycle.
        """
        # Capture model name + input messages once per cycle
        if self._model_span and not self._model_info_captured and event_kwargs:
            # Extract model name: try explicit "model" key first, then agent object
            model = event_kwargs.get("model")
            if not model:
                model = self._extract_model_name_from_agent(event_kwargs.get("agent"))
            if model and isinstance(model, str):
                self._model_span.set_attribute(self.ATTR_GEN_AI_REQUEST_MODEL, model)
                self._model_span.set_attribute(self.ATTR_GEN_AI_RESPONSE_MODEL, model)
                if self._agent_span:
                    self._agent_span.set_attribute(self.ATTR_GEN_AI_REQUEST_MODEL, model)

            # Extract messages: try explicit "messages" key first, then agent object
            messages = event_kwargs.get("messages")
            if not messages:
                messages = self._extract_messages_from_agent(event_kwargs.get("agent"))
            if messages:
                try:
                    self._model_span.set_attribute(
                        self.ATTR_GEN_AI_INPUT_MESSAGES,
                        json.dumps(messages, default=str),
                    )
                    self._model_span.set_input(messages)
                    # Don't also apply pending_input_messages — data event wins
                    self._pending_input_messages = []

                    # Set the agent span input once (first cycle only) from the
                    # user's initial message in the conversation history.
                    if not self._agent_input_set and self._agent_span and isinstance(messages, list):
                        user_input = self._extract_user_input(messages)
                        if user_input is not None:
                            self._agent_span.set_input(user_input)
                            self._agent_input_set = True
                except Exception:
                    logger.debug("Could not set input messages from data event", exc_info=True)
            self._model_info_captured = True

        text = self._extract_text(data)
        if not text:
            return
        if self._first_token_time is None:
            self._first_token_time = time.perf_counter()
        self._streaming_buffer.append(text)

    def _on_reasoning_text(self, reasoning_text: Any) -> None:
        """A reasoning/thinking chunk arrived from the model."""
        text = self._extract_text(reasoning_text)
        if text:
            self._reasoning_buffer.append(text)

    def _on_current_tool_use(self, current_tool_use: dict[str, Any]) -> None:
        """The model is streaming a tool call — accumulate tool info.

        Strands uses ``toolUseId`` as the tool identifier. The ``input``
        field is a JSON string that accumulates across multiple callback
        events as the model streams the tool arguments.
        """
        # Strands uses "toolUseId"; fall back to "id" for forward-compat
        tool_id = current_tool_use.get("toolUseId") or current_tool_use.get("id", "")
        if not tool_id:
            return

        tool_name = current_tool_use.get("name", "")
        input_raw = current_tool_use.get("input", "")

        if tool_id not in self._pending_tool_calls:
            self._pending_tool_calls[tool_id] = {
                "name": tool_name,
                "id": tool_id,
                "input": None,
            }

        # input may be a partial/complete JSON string or already a dict
        if isinstance(input_raw, dict):
            self._pending_tool_calls[tool_id]["input"] = input_raw
        elif isinstance(input_raw, str) and input_raw:
            try:
                self._pending_tool_calls[tool_id]["input"] = json.loads(input_raw)
            except json.JSONDecodeError:
                pass  # Input still streaming — try again on next event

    def _on_complete(self, **kwargs: Any) -> None:
        """The model has finished streaming its response for this cycle.

        Tool calls (if any) will execute after this point. We close the
        model span here and create tool spans for any pending tool calls.
        Tool spans will be closed when the corresponding ``message`` event
        arrives with the tool result, or at the next ``start_event_loop``
        as a fallback.
        """
        logger.debug(
            "BeaconStrandsHandler _on_complete: buffer_len=%d pending_tools=%s",
            len(self._streaming_buffer),
            list(self._pending_tool_calls.keys()),
        )
        self._finalize_model_span()

        # Create spans for tool calls — they execute between now and the
        # next start_event_loop / message(toolResult) event.
        for tool_id, tool_info in self._pending_tool_calls.items():
            if tool_id in self._active_tool_spans:
                continue  # Already created
            tool_span = Span(
                name=tool_info["name"] or "tool",
                trace_id=self._trace_id,
                parent_id=self._agent_span.span_id if self._agent_span else None,
                span_type=SpanType.TOOL,
                kind=SpanKind.INTERNAL,
                session_id=self._session_id,
                user_email=self._user_email,
                user_name=self._user_name,
                user_id=self._user_id,
                record_stacktrace=self._record_stacktrace,
            )
            tool_span.start()
            tool_span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "execute_tool")
            tool_span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "strands")
            tool_span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "strands")
            tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_NAME, tool_info["name"])
            tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_TYPE, "function")
            tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_CALL_ID, tool_id)
            if tool_info.get("input") is not None:
                args_json = json.dumps(tool_info["input"], default=str)
                tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_CALL_ARGUMENTS, args_json)
                tool_span.set_input(tool_info["input"])
            self._active_tool_spans[tool_id] = tool_span
            self._client.export_eager_span(tool_span)

        self._pending_tool_calls = {}

    def _on_message(self, message: dict[str, Any]) -> None:
        """Handle message events.

        ``role="assistant"``:
            Finalize the model span (``complete`` never reliably fires, so the
            assistant message is our signal that the model is done). Also create
            tool spans for any ``toolUse`` blocks — the input dict is already
            fully parsed here, so we don't need the partial accumulation from
            ``current_tool_use`` events.

        ``role="user"`` with ``toolResult`` blocks:
            A tool has finished; close the matching tool span and capture output.

        ``role="user"`` without ``toolResult`` blocks:
            The user's query is being added to conversation; queue as fallback
            input for the next model span (overridden by data-event messages).
        """
        role = message.get("role")
        content = message.get("content", [])

        logger.debug("BeaconStrandsHandler _on_message role=%s content_len=%d", role, len(content))

        if role == "assistant":
            # Capture the full content blocks (text + toolUse) before finalizing
            # so _finalize_model_span can build structured gen_ai.output.messages
            if content and isinstance(content, list):
                self._assistant_message_content = content

            # Finalize the model span — complete event doesn't reliably fire
            self._finalize_model_span()

            # Create tool spans from toolUse blocks (input is already fully parsed)
            for block in content:
                if not isinstance(block, dict) or "toolUse" not in block:
                    continue
                tool_use = block["toolUse"]
                tool_id = tool_use.get("toolUseId") or tool_use.get("id", "")
                if not tool_id or tool_id in self._active_tool_spans:
                    continue
                tool_name = tool_use.get("name") or "tool"
                tool_input = tool_use.get("input")

                logger.debug(
                    "BeaconStrandsHandler creating tool span: id=%s name=%s",
                    tool_id,
                    tool_name,
                )
                tool_span = Span(
                    name=tool_name,
                    trace_id=self._trace_id,
                    parent_id=self._agent_span.span_id if self._agent_span else None,
                    span_type=SpanType.TOOL,
                    kind=SpanKind.INTERNAL,
                    session_id=self._session_id,
                    user_email=self._user_email,
                    user_name=self._user_name,
                    user_id=self._user_id,
                    record_stacktrace=self._record_stacktrace,
                )
                tool_span.start()
                tool_span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "execute_tool")
                tool_span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "strands")
                tool_span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "strands")
                tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_NAME, tool_name)
                tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_TYPE, "function")
                tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_CALL_ID, tool_id)
                if tool_input is not None:
                    args_json = json.dumps(tool_input, default=str)
                    tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_CALL_ARGUMENTS, args_json)
                    tool_span.set_input(tool_input)
                self._active_tool_spans[tool_id] = tool_span
                self._client.export_eager_span(tool_span)

        elif role == "user":
            has_tool_result = any("toolResult" in block for block in content if isinstance(block, dict))

            if has_tool_result:
                # Tool result(s) — close matching tool spans
                for block in content:
                    if not isinstance(block, dict) or "toolResult" not in block:
                        continue
                    tool_result = block["toolResult"]
                    tool_use_id = tool_result.get("toolUseId", "")

                    output_text = None
                    for item in tool_result.get("content", []):
                        if isinstance(item, dict) and "text" in item:
                            output_text = item["text"]
                            break

                    logger.debug(
                        "BeaconStrandsHandler tool result: id=%s output=%r",
                        tool_use_id,
                        output_text[:80] if output_text else None,
                    )

                    tool_span = self._active_tool_spans.pop(tool_use_id, None)
                    if tool_span:
                        if output_text:
                            tool_span.set_attribute(self.ATTR_GEN_AI_TOOL_CALL_RESULT, output_text)
                            tool_span.set_output(output_text)
                        tool_span.set_status(StatusCode.OK)
                        tool_span.end()
                        self._client.export_span(tool_span)
            else:
                # Plain user message — queue as fallback input for the next model span
                # (will be overridden by the ``messages`` key on data events if present)
                serialized = self._serialize_message(message)
                self._pending_input_messages.append(serialized)

    def _on_result(self, result: Any) -> None:
        """The agent has completed successfully."""
        logger.debug("BeaconStrandsHandler _on_result: %r", type(result).__name__)

        # Close any spans still open (safety nets)
        self._end_active_tool_spans(status=StatusCode.OK)
        if self._model_span:
            self._finalize_model_span()

        if self._agent_span:
            output = self._extract_result_output(result)
            if output:
                self._agent_span.set_output(output)
                self._agent_span.set_attribute(
                    self.ATTR_GEN_AI_OUTPUT_MESSAGES,
                    json.dumps([{"role": "assistant", "content": output}], default=str),
                )

            self._apply_token_metrics(self._agent_span, result)
            self._apply_model_name(self._agent_span, result)

            self._agent_span.set_status(StatusCode.OK)
            self._agent_span.end()
            self._client.export_span(self._agent_span)
            self._agent_span = None

        clear_context()

    def _on_force_stop(self, force_stop: bool, force_stop_reason: str | None = None, **kwargs: Any) -> None:
        """The agent was force-stopped before completing."""
        logger.debug("BeaconStrandsHandler _on_force_stop reason=%r", force_stop_reason)
        self._end_active_tool_spans(status=StatusCode.ERROR)

        if self._model_span:
            self._model_span.set_attribute("error.type", "ForceStop")
            self._model_span.set_status(StatusCode.ERROR, force_stop_reason or "Force stopped")
            self._model_span.end()
            self._client.export_span(self._model_span)
            self._model_span = None

        if self._agent_span:
            reason = force_stop_reason or "Agent force-stopped"
            self._agent_span.set_attribute("error.type", "ForceStop")
            self._agent_span.record_exception(Exception(reason))
            self._agent_span.set_status(StatusCode.ERROR, reason)
            self._agent_span.end()
            self._client.export_span(self._agent_span)
            self._agent_span = None

        clear_context()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _finalize_model_span(self) -> None:
        """Set output/attributes and export the current model span."""
        if not self._model_span:
            return

        # Use full assistant message content when available (includes toolUse blocks),
        # otherwise fall back to the concatenated streaming buffer text.
        if self._assistant_message_content is not None:
            self._model_span.set_output(self._assistant_message_content)
            self._model_span.set_attribute(
                self.ATTR_GEN_AI_OUTPUT_MESSAGES,
                json.dumps(
                    [{"role": "assistant", "content": self._assistant_message_content}],
                    default=str,
                ),
            )
        else:
            output = "".join(self._streaming_buffer)
            if output:
                self._model_span.set_output(output)
                self._model_span.set_attribute(
                    self.ATTR_GEN_AI_OUTPUT_MESSAGES,
                    json.dumps([{"role": "assistant", "content": output}], default=str),
                )

        if self._reasoning_buffer:
            self._model_span.set_attribute(
                self.ATTR_GEN_AI_REASONING_CONTENT,
                "".join(self._reasoning_buffer),
            )

        if self._first_token_time is not None and self._loop_start_time is not None:
            ttft_ms = (self._first_token_time - self._loop_start_time) * 1000
            self._model_span.set_attribute(self.ATTR_GEN_AI_TTFT_MS, ttft_ms)

        self._model_span.set_status(StatusCode.OK)
        self._model_span.end()
        self._client.export_span(self._model_span)
        self._model_span = None

    def _emergency_close(self, exc: Exception) -> None:
        """Close all open spans with ERROR status when the handler itself fails.

        Called whenever an unhandled exception escapes an event handler so
        that traces are always finalised rather than left dangling.
        """
        try:
            error_msg = str(exc)
            error_type = type(exc).__qualname__
            for tool_span in self._active_tool_spans.values():
                tool_span.set_attribute("error.type", error_type)
                tool_span.record_exception(exc)
                tool_span.set_status(StatusCode.ERROR, error_msg)
                tool_span.end()
                self._client.export_span(tool_span)
            self._active_tool_spans = {}

            if self._model_span:
                self._model_span.set_attribute("error.type", error_type)
                self._model_span.record_exception(exc)
                self._model_span.set_status(StatusCode.ERROR, error_msg)
                self._model_span.end()
                self._client.export_span(self._model_span)
                self._model_span = None

            if self._agent_span:
                self._agent_span.set_attribute("error.type", error_type)
                self._agent_span.record_exception(exc)
                self._agent_span.set_status(StatusCode.ERROR, error_msg)
                self._agent_span.end()
                self._client.export_span(self._agent_span)
                self._agent_span = None

            clear_context()
        except Exception:
            logger.debug("Error during emergency close", exc_info=True)

    def _extract_model_name_from_agent(self, agent: Any) -> str | None:
        """Extract the model name/ID from a Strands Agent object.

        Strands merges the ``agent`` object into callback kwargs via
        ``invocation_state``. The model name lives on the agent's model
        provider, typically as ``agent.model.config["model_id"]``.
        """
        if agent is None:
            return None
        try:
            model_provider = getattr(agent, "model", None)
            if model_provider is None:
                return None
            # Most Strands model providers store model_id in config dict
            config = getattr(model_provider, "config", None)
            if isinstance(config, dict):
                model_id = config.get("model_id")
                if model_id:
                    return str(model_id)
            # Fallback: check for model_id attribute directly
            model_id = getattr(model_provider, "model_id", None)
            if model_id:
                return str(model_id)
        except Exception:
            logger.debug("Could not extract model name from agent", exc_info=True)
        return None

    def _extract_messages_from_agent(self, agent: Any) -> list[dict[str, Any]] | None:
        """Extract the conversation messages from a Strands Agent object.

        Strands stores the conversation history on
        ``agent.conversation_manager`` or ``agent.messages``.
        """
        if agent is None:
            return None
        try:
            messages = getattr(agent, "messages", None)
            if isinstance(messages, list) and messages:
                return messages
        except Exception:
            logger.debug("Could not extract messages from agent", exc_info=True)
        return None

    def _extract_text(self, value: Any) -> str:
        """Normalise a data/reasoningText value to a plain string.

        Strands normally passes a ``str``, but with extended thinking or
        certain model backends it may pass a raw event ``dict`` (e.g.
        ``{"text": "..."}``).  Returns an empty string for anything that
        cannot be coerced to meaningful text.
        """
        if isinstance(value, str):
            return value
        if isinstance(value, dict):
            return value.get("text") or value.get("content") or ""
        return ""

    def _extract_user_input(self, messages: list[dict[str, Any]]) -> str | None:
        """Extract the user's initial query from the conversation history.

        Walks the messages list backwards and returns the text content of the
        last ``role="user"`` message that is NOT a tool result.
        """
        for msg in reversed(messages):
            if not isinstance(msg, dict) or msg.get("role") != "user":
                continue
            content = msg.get("content", "")
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                # Skip tool-result messages
                if any(isinstance(b, dict) and "toolResult" in b for b in content):
                    continue
                texts = [
                    b["text"] for b in content
                    if isinstance(b, dict) and "text" in b
                ]
                if texts:
                    return "\n".join(texts)
        return None

    def _end_active_tool_spans(self, status: StatusCode = StatusCode.OK) -> None:
        """Close all active tool spans with the given status (safety net)."""
        for tool_span in self._active_tool_spans.values():
            tool_span.set_status(status)
            tool_span.end()
            self._client.export_span(tool_span)
        self._active_tool_spans = {}

    def _serialize_message(self, message: dict[str, Any]) -> dict[str, Any]:
        """Serialize a Strands message dict for storage as a span attribute."""
        role = message.get("role", "unknown")
        content = message.get("content", "")

        if isinstance(content, str):
            return {"role": role, "content": content}

        if isinstance(content, list):
            parts: list[str] = []
            for block in content:
                if isinstance(block, dict):
                    if "text" in block:
                        parts.append(block["text"])
                    elif "toolUse" in block:
                        tu = block["toolUse"]
                        parts.append(f"[tool:{tu.get('name')}]")
                elif isinstance(block, str):
                    parts.append(block)
            return {"role": role, "content": " ".join(parts)}

        return {"role": role, "content": str(content)}

    def _extract_result_output(self, result: Any) -> str | None:
        """Extract the final text output from an ``AgentResult``.

        Strands content blocks are ``{"text": "..."}`` (no ``"type"`` field),
        so we match any block that has a non-empty ``"text"`` key directly.
        """
        try:
            msg = getattr(result, "message", None)
            if not isinstance(msg, dict):
                return None
            content = msg.get("content", [])
            texts = [
                block["text"]
                for block in content
                if isinstance(block, dict) and block.get("text")
            ]
            return "\n".join(texts) if texts else None
        except Exception:
            logger.debug("Could not extract result output", exc_info=True)
            return None

    def _apply_token_metrics(self, span: Span, result: Any) -> None:
        """Copy accumulated token usage from ``AgentResult.metrics`` onto *span*."""
        try:
            metrics = getattr(result, "metrics", None)
            if metrics is None:
                return
            usage: dict[str, Any] = getattr(metrics, "accumulated_usage", None) or {}
            logger.debug("BeaconStrandsHandler token usage: %r", usage)

            input_tokens = usage.get("inputTokens", 0) or 0
            output_tokens = usage.get("outputTokens", 0) or 0

            if input_tokens:
                span.set_attribute(self.ATTR_GEN_AI_INPUT_TOKENS, input_tokens)
            if output_tokens:
                span.set_attribute(self.ATTR_GEN_AI_OUTPUT_TOKENS, output_tokens)
            if usage.get("cacheReadInputTokens"):
                span.set_attribute(self.ATTR_GEN_AI_CACHE_READ_TOKENS, usage["cacheReadInputTokens"])
            if usage.get("cacheWriteInputTokens"):
                span.set_attribute(self.ATTR_GEN_AI_CACHE_CREATION_TOKENS, usage["cacheWriteInputTokens"])
        except Exception:
            logger.debug("Could not extract token metrics", exc_info=True)

    def _apply_model_name(self, span: Span, result: Any) -> None:
        """Set model name attributes from ``AgentResult`` if not already set.

        The model name is normally captured earlier from the ``data`` event's
        ``model`` key. This method is a fallback for the agent span which
        finalises after streaming has already ended.
        """
        try:
            metrics = getattr(result, "metrics", None)
            model = (
                getattr(metrics, "model", None)
                or getattr(result, "model_id", None)
                or getattr(result, "model", None)
            )
            if model:
                model_str = str(model)
                span.set_attribute(self.ATTR_GEN_AI_REQUEST_MODEL, model_str)
                span.set_attribute(self.ATTR_GEN_AI_RESPONSE_MODEL, model_str)
        except Exception:
            logger.debug("Could not extract model name", exc_info=True)
