"""Tests for the Strands Agents callback handler integration."""

import json
from unittest.mock import MagicMock, patch

import pytest

from lumenova_beacon import BeaconClient
from lumenova_beacon.tracing.integrations.strands import BeaconStrandsHandler
from lumenova_beacon.types import SpanType, SpanKind, StatusCode

PATCH_GET_CLIENT = "lumenova_beacon.tracing.integrations.strands.get_client"


@pytest.fixture
def mock_client():
    """A mock BeaconClient with a stubbed config."""
    client = MagicMock(spec=BeaconClient)
    client.config = MagicMock()
    client.config.session_id = None
    return client


@pytest.fixture
def handler(mock_client):
    """BeaconStrandsHandler with injected mock client."""
    with patch(PATCH_GET_CLIENT, return_value=mock_client):
        return BeaconStrandsHandler()


@pytest.fixture
def named_handler(mock_client):
    """BeaconStrandsHandler with all optional params set."""
    with patch(PATCH_GET_CLIENT, return_value=mock_client):
        return BeaconStrandsHandler(
            session_id="sess-1",
            environment="test",
            agent_name="Test Agent",
            metadata={"user": "alice"},
        )


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_result(text: str = "Hello", input_tokens: int = 10, output_tokens: int = 5) -> MagicMock:
    result = MagicMock()
    # Strands content blocks use {"text": "..."} without a "type" field
    result.message = {
        "role": "assistant",
        "content": [{"text": text}],
    }
    result.metrics = MagicMock()
    result.metrics.accumulated_usage = {
        "inputTokens": input_tokens,
        "outputTokens": output_tokens,
    }
    return result


def _tool_use_event(name: str, tool_use_id: str, input_data: dict | None = None) -> dict:
    """Build a current_tool_use dict as Strands sends it (toolUseId key, JSON string input)."""
    return {
        "toolUseId": tool_use_id,
        "name": name,
        "input": json.dumps(input_data) if input_data else "",
    }


def _assistant_message_with_tools(*tool_specs: tuple) -> dict:
    """Build a message(role=assistant) event with toolUse blocks.

    Each element of tool_specs is a (name, tool_use_id, input_dict_or_None) tuple.
    """
    content = []
    for name, tool_use_id, input_data in tool_specs:
        content.append({
            "toolUse": {
                "toolUseId": tool_use_id,
                "name": name,
                "input": input_data or {},
            }
        })
    return {"role": "assistant", "content": content}


def _assistant_message_text(text: str) -> dict:
    """Build a plain text message(role=assistant) event."""
    return {"role": "assistant", "content": [{"text": text}]}


def _tool_result_message(tool_use_id: str, output: str) -> dict:
    """Build a message event dict that signals a tool result."""
    return {
        "role": "user",
        "content": [
            {
                "toolResult": {
                    "toolUseId": tool_use_id,
                    "content": [{"text": output}],
                }
            }
        ],
    }


# ---------------------------------------------------------------------------
# Initialization
# ---------------------------------------------------------------------------

class TestBeaconStrandsHandlerInit:
    def test_defaults(self, mock_client):
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            handler = BeaconStrandsHandler()
        assert handler._client is mock_client
        assert handler._session_id is None
        assert handler._agent_name is None
        assert handler._metadata == {}
        assert handler._cycle_count == 0
        assert handler._agent_span is None

    def test_session_id_from_client_config(self, mock_client):
        mock_client.config.session_id = "from-config"
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            handler = BeaconStrandsHandler()
        assert handler._session_id == "from-config"

    def test_explicit_session_id_overrides_config(self, mock_client):
        mock_client.config.session_id = "from-config"
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            handler = BeaconStrandsHandler(session_id="explicit")
        assert handler._session_id == "explicit"

    def test_all_params(self, mock_client):
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            handler = BeaconStrandsHandler(
                session_id="s",
                environment="prod",
                agent_name="My Agent",
                metadata={"k": "v"},
            )
        assert handler._session_id == "s"
        assert handler._environment == "prod"
        assert handler._agent_name == "My Agent"
        assert handler._metadata == {"k": "v"}

    def test_trace_id_initially_none(self, handler):
        assert handler.trace_id is None


# ---------------------------------------------------------------------------
# Agent span creation (init_event_loop)
# ---------------------------------------------------------------------------

class TestInitEventLoop:
    def test_creates_agent_span(self, handler):
        handler(init_event_loop=True)
        assert handler._agent_span is not None
        assert handler._agent_span.span_type == SpanType.AGENT
        assert handler._agent_span.kind == SpanKind.INTERNAL

    def test_trace_id_populated(self, handler):
        handler(init_event_loop=True)
        assert handler.trace_id is not None
        assert handler._agent_span.trace_id == handler.trace_id

    def test_agent_span_has_no_parent(self, handler):
        handler(init_event_loop=True)
        assert handler._agent_span.parent_id is None

    def test_agent_name_attribute(self, named_handler):
        named_handler(init_event_loop=True)
        attrs = named_handler._agent_span.attributes
        assert attrs.get("gen_ai.agent.name") == "Test Agent"
        assert named_handler._agent_span.name == "Test Agent"

    def test_default_agent_name(self, handler):
        handler(init_event_loop=True)
        assert handler._agent_span.name == "strands-agent"

    def test_session_id_attribute(self, mock_client):
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            handler = BeaconStrandsHandler(session_id="my-session")
        handler(init_event_loop=True)
        assert handler._agent_span.attributes.get("gen_ai.conversation.id") == "my-session"

    def test_metadata_attributes(self, named_handler):
        named_handler(init_event_loop=True)
        assert named_handler._agent_span.attributes.get("beacon.metadata.user") == "alice"

    def test_environment_attribute(self, named_handler):
        named_handler(init_event_loop=True)
        assert named_handler._agent_span.attributes.get("deployment.environment.name") == "test"

    def test_eager_export_called(self, handler, mock_client):
        handler(init_event_loop=True)
        mock_client.export_eager_span.assert_called_once_with(handler._agent_span)


# ---------------------------------------------------------------------------
# Model call spans (start_event_loop / complete / message(role=assistant))
# ---------------------------------------------------------------------------

class TestModelCallSpans:
    def test_creates_model_span(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        assert handler._model_span is not None
        assert handler._model_span.span_type == SpanType.GENERATION
        assert handler._model_span.kind == SpanKind.CLIENT

    def test_model_span_parent_is_agent_span(self, handler):
        handler(init_event_loop=True)
        agent_span_id = handler._agent_span.span_id
        handler(start_event_loop=True)
        assert handler._model_span.parent_id == agent_span_id

    def test_model_span_shares_trace_id(self, handler):
        handler(init_event_loop=True)
        trace_id = handler.trace_id
        handler(start_event_loop=True)
        assert handler._model_span.trace_id == trace_id

    def test_cycle_attribute_increments(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        assert handler._model_span.attributes.get("strands.event_loop.cycle") == 1

    def test_complete_closes_model_span(self, handler, mock_client):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(complete=True)
        assert handler._model_span is None
        mock_client.export_span.assert_called()

    def test_complete_sets_model_output(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="Hello ")
        handler(data="world")
        model_span = handler._model_span
        handler(complete=True)
        assert model_span.attributes.get("span.output") == "Hello world"

    def test_complete_ok_status(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        model_span = handler._model_span
        handler(complete=True)
        assert model_span.status_code == StatusCode.OK

    def test_complete_with_data_in_same_call(self, handler):
        """complete and data may arrive in the same kwargs call."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        model_span = handler._model_span
        handler(data="last chunk", complete=True)
        assert model_span.status_code == StatusCode.OK
        assert model_span.attributes.get("span.output") == "last chunk"
        assert handler._model_span is None

    def test_assistant_message_closes_model_span(self, handler, mock_client):
        """message(role=assistant) finalizes the model span (primary path — complete doesn't fire)."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="The answer is 2.")
        model_span = handler._model_span
        handler(message=_assistant_message_text("The answer is 2."))
        assert handler._model_span is None
        assert model_span.status_code == StatusCode.OK
        mock_client.export_span.assert_called()

    def test_assistant_message_sets_model_output(self, handler):
        """Text in streaming buffer is set as model span output on assistant message."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="Paris")
        model_span = handler._model_span
        handler(message=_assistant_message_text("Paris"))
        assert model_span.attributes.get("span.output") == "Paris"

    def test_model_span_closed_idempotent(self, handler):
        """Calling finalize twice (complete + assistant message) must not crash."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="hi")
        handler(complete=True)
        # complete already closed the model span; assistant message should be a no-op
        handler(message=_assistant_message_text("hi"))  # must not raise

    def test_multiple_cycles(self, handler, mock_client):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("cycle 1"))
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("cycle 2"))
        handler(result=_make_result())
        # 2 model spans + 1 agent span = 3 export_span calls minimum
        assert mock_client.export_span.call_count >= 3

    def test_safety_net_closes_model_span_at_next_cycle(self, handler, mock_client):
        """If no assistant message fires, start_event_loop should close the previous model span."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        prev_model_span = handler._model_span
        # No complete or assistant message fires — next cycle starts
        handler(start_event_loop=True)
        assert prev_model_span.end_time is not None
        assert prev_model_span.status_code == StatusCode.OK

    def test_model_name_set_from_data_event(self, handler):
        """model name is captured from the 'model' key on data events."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="hi", model="us.anthropic.claude-3-7-sonnet-20250219-v1:0")
        assert handler._model_span.attributes.get("gen_ai.request.model") == (
            "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
        )

    def test_model_name_captured_only_once_per_cycle(self, handler):
        """model info capture flag is reset between cycles."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="a", model="model-1")
        handler(data="b", model="model-2")  # second data — should not override
        assert handler._model_span.attributes.get("gen_ai.request.model") == "model-1"
        # After next cycle, the flag resets
        handler(start_event_loop=True)
        assert not handler._model_info_captured


# ---------------------------------------------------------------------------
# Streaming (data / delta / reasoningText)
# ---------------------------------------------------------------------------

class TestStreaming:
    def test_data_accumulates(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="foo")
        handler(data="bar")
        assert handler._streaming_buffer == ["foo", "bar"]

    def test_delta_treated_same_as_data(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(delta="baz")
        assert handler._streaming_buffer == ["baz"]

    def test_reasoning_text_accumulated_separately(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(reasoningText="thinking...")
        assert handler._reasoning_buffer == ["thinking..."]
        assert handler._streaming_buffer == []

    def test_reasoning_text_stored_on_model_span(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(reasoningText="thought 1")
        handler(reasoningText=" thought 2")
        model_span = handler._model_span
        handler(complete=True)
        assert model_span.attributes.get("gen_ai.reasoning_content") == "thought 1 thought 2"

    def test_ttft_calculated_in_ms(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="first token")
        model_span = handler._model_span
        handler(complete=True)
        assert "gen_ai.response.time_to_first_token_ms" in model_span.attributes

    def test_ttft_not_set_when_no_data(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        model_span = handler._model_span
        handler(complete=True)
        assert "gen_ai.response.time_to_first_token_ms" not in model_span.attributes

    def test_delta_not_double_counted_when_data_present(self, handler):
        """If both data and delta arrive in the same call, only data is used."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="hello", delta="hello")
        assert handler._streaming_buffer == ["hello"]

    def test_delta_used_when_no_data(self, handler):
        """delta is processed when data is absent."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(delta="fallback")
        assert handler._streaming_buffer == ["fallback"]

    def test_buffer_reset_between_cycles(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="cycle 1")
        handler(complete=True)
        handler(start_event_loop=True)
        assert handler._streaming_buffer == []

    def test_gen_ai_output_messages_set(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data="Paris is the capital.")
        model_span = handler._model_span
        handler(complete=True)
        assert "gen_ai.output.messages" in model_span.attributes

    def test_gen_ai_operation_name_set(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        assert handler._model_span.attributes.get("gen_ai.operation.name") == "chat"


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

class TestToolSpans:
    def test_pending_tool_call_accumulated_with_tooluseid(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(current_tool_use=_tool_use_event("calculator", "t-1", {"x": 5}))
        assert "t-1" in handler._pending_tool_calls
        assert handler._pending_tool_calls["t-1"]["name"] == "calculator"

    def test_pending_tool_id_key_is_tooluseid(self, handler):
        """Strands uses 'toolUseId', not 'id'."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(current_tool_use={"toolUseId": "abc", "name": "search", "input": ""})
        assert "abc" in handler._pending_tool_calls

    def test_tool_input_parsed_from_json_string(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(current_tool_use=_tool_use_event("calc", "t-2", {"a": 1, "b": 2}))
        assert handler._pending_tool_calls["t-2"]["input"] == {"a": 1, "b": 2}

    def test_partial_json_not_parsed(self, handler):
        """Incomplete JSON string should not crash and input stays None."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(current_tool_use={"toolUseId": "t-x", "name": "tool", "input": '{"a":'})
        assert handler._pending_tool_calls["t-x"]["input"] is None

    def test_tool_span_created_from_assistant_message(self, handler, mock_client):
        """Tool spans are created when message(role=assistant) contains toolUse blocks."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("calculator", "t-3", {"x": 5})))
        assert "t-3" in handler._active_tool_spans
        tool_span = handler._active_tool_spans["t-3"]
        assert tool_span.span_type == SpanType.TOOL
        assert tool_span.name == "calculator"

    def test_tool_span_also_created_from_complete_fallback(self, handler):
        """complete event still creates tool spans as a fallback (for Strands versions that fire it)."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(current_tool_use=_tool_use_event("calculator", "t-3b", {"x": 5}))
        handler(complete=True)
        assert "t-3b" in handler._active_tool_spans

    def test_tool_span_attributes(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("search", "t-4", {"q": "test"})))
        tool_span = handler._active_tool_spans["t-4"]
        assert tool_span.attributes.get("gen_ai.tool.name") == "search"
        assert tool_span.attributes.get("gen_ai.tool.call.id") == "t-4"

    def test_tool_span_has_input(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("calc", "t-5", {"val": 42})))
        tool_span = handler._active_tool_spans["t-5"]
        assert "span.input" in tool_span.attributes

    def test_tool_span_parent_is_agent_span(self, handler):
        handler(init_event_loop=True)
        agent_id = handler._agent_span.span_id
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("tool", "t-6", {})))
        tool_span = handler._active_tool_spans["t-6"]
        assert tool_span.parent_id == agent_id

    def test_tool_span_closed_by_tool_result_message(self, handler, mock_client):
        """Tool span closes and captures output when toolResult user message arrives."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("calculator", "t-7", {"x": 5})))
        tool_span = handler._active_tool_spans["t-7"]

        handler(message=_tool_result_message("t-7", "42"))

        assert "t-7" not in handler._active_tool_spans
        assert tool_span.end_time is not None
        assert tool_span.status_code == StatusCode.OK
        assert tool_span.attributes.get("span.output") == "42"

    def test_tool_span_output_set_from_tool_result(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("divide", "t-8", {"a": 10, "b": 2})))
        tool_span = handler._active_tool_spans["t-8"]
        handler(message=_tool_result_message("t-8", "5.0"))
        assert tool_span.attributes.get("span.output") == "5.0"

    def test_tool_span_safety_net_at_next_start_event_loop(self, handler):
        """If message event never arrives, start_event_loop should close tool spans."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("tool", "t-9", {})))
        tool_span = handler._active_tool_spans["t-9"]
        handler(start_event_loop=True)
        assert "t-9" not in handler._active_tool_spans
        assert tool_span.end_time is not None

    def test_multiple_tool_spans(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(
            ("tool_a", "ta", {"n": 1}),
            ("tool_b", "tb", {"n": 2}),
        ))
        assert "ta" in handler._active_tool_spans
        assert "tb" in handler._active_tool_spans

    def test_tool_span_with_empty_tool_id_ignored(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        msg = {"role": "assistant", "content": [{"toolUse": {"toolUseId": "", "name": "tool", "input": {}}}]}
        handler(message=msg)
        assert len(handler._active_tool_spans) == 0

    def test_assistant_message_text_only_does_not_create_tool_spans(self, handler):
        """Plain text assistant message should not create any tool spans."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("The answer is 42."))
        assert len(handler._active_tool_spans) == 0

    def test_tool_span_not_duplicated_on_second_assistant_message(self, handler):
        """Same toolUseId in a second assistant message must not create a duplicate span."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("tool", "t-dup", {})))
        assert len(handler._active_tool_spans) == 1
        handler(message=_assistant_message_with_tools(("tool", "t-dup", {})))
        assert len(handler._active_tool_spans) == 1

    def test_tool_span_has_gen_ai_tool_call_arguments(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("add", "t-arg", {"a": 1, "b": 2})))
        tool_span = handler._active_tool_spans["t-arg"]
        assert "gen_ai.tool.call.arguments" in tool_span.attributes

    def test_tool_span_has_gen_ai_tool_call_result(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("add", "t-res", {"a": 1, "b": 2})))
        tool_span = handler._active_tool_spans["t-res"]
        handler(message=_tool_result_message("t-res", "3"))
        assert tool_span.attributes.get("gen_ai.tool.call.result") == "3"

    def test_eager_export_called_for_tool_span(self, handler, mock_client):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("tool", "t-eager", {})))
        # export_eager_span should have been called for init + start + tool
        assert mock_client.export_eager_span.call_count == 3


# ---------------------------------------------------------------------------
# Input messages on model spans
# ---------------------------------------------------------------------------

class TestInputMessages:
    def test_user_message_queued_before_start_event_loop(self, handler):
        """Plain user messages queued before start_event_loop appear as model span input."""
        handler(init_event_loop=True)
        handler(message={"role": "user", "content": "What is 1+1?"})
        handler(start_event_loop=True)
        attrs = handler._model_span.attributes
        assert "gen_ai.input.messages" in attrs
        assert "span.input" in attrs

    def test_user_message_content_in_input(self, handler):
        handler(init_event_loop=True)
        handler(message={"role": "user", "content": "What is 1+1?"})
        handler(start_event_loop=True)
        raw = handler._model_span.attributes["gen_ai.input.messages"]
        parsed = json.loads(raw)
        assert parsed[0]["role"] == "user"
        assert "1+1" in parsed[0]["content"]

    def test_messages_from_data_event_override_pending(self, handler):
        """The 'messages' key on data events provides the full conversation and overrides queued messages."""
        handler(init_event_loop=True)
        handler(message={"role": "user", "content": "Queued message"})
        handler(start_event_loop=True)
        # data event carries the actual conversation history
        full_messages = [
            {"role": "user", "content": [{"text": "What is 1+1?"}]},
        ]
        handler(data="2", messages=full_messages)
        raw = handler._model_span.attributes.get("gen_ai.input.messages", "[]")
        parsed = json.loads(raw)
        # messages from data event should win (they're the full conversation)
        assert parsed == full_messages

    def test_tool_result_message_not_queued_as_input(self, handler):
        """Tool result messages must not appear as model span inputs."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        # Create a tool span via assistant message
        handler(message=_assistant_message_with_tools(("add", "t-q", {})))
        # Close it with a tool result — this should NOT be queued as input
        handler(message=_tool_result_message("t-q", "42"))
        handler(start_event_loop=True)
        # Second model span should have no input messages from tool result
        assert "gen_ai.input.messages" not in handler._model_span.attributes

    def test_input_messages_cleared_after_use(self, handler):
        handler(init_event_loop=True)
        handler(message={"role": "user", "content": "Hello"})
        handler(start_event_loop=True)
        assert handler._pending_input_messages == []


# ---------------------------------------------------------------------------
# Result (agent completion)
# ---------------------------------------------------------------------------

class TestResult:
    def test_agent_span_ended(self, handler, mock_client):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("done"))
        handler(result=_make_result())
        assert handler._agent_span is None
        assert mock_client.export_span.call_count >= 2

    def test_agent_span_ok_status(self, handler):
        handler(init_event_loop=True)
        agent_span = handler._agent_span
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("done"))
        handler(result=_make_result())
        assert agent_span.status_code == StatusCode.OK

    def test_output_extracted_from_result(self, handler):
        handler(init_event_loop=True)
        agent_span = handler._agent_span
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("Paris is the capital of France."))
        handler(result=_make_result("Paris is the capital of France."))
        assert agent_span.attributes.get("span.output") == "Paris is the capital of France."

    def test_token_metrics_applied(self, handler):
        handler(init_event_loop=True)
        agent_span = handler._agent_span
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("done"))
        handler(result=_make_result(input_tokens=100, output_tokens=50))
        assert agent_span.attributes.get("gen_ai.usage.input_tokens") == 100
        assert agent_span.attributes.get("gen_ai.usage.output_tokens") == 50

    def test_cache_token_metrics(self, handler):
        handler(init_event_loop=True)
        agent_span = handler._agent_span
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("done"))
        result = _make_result()
        result.metrics.accumulated_usage = {
            "inputTokens": 10,
            "outputTokens": 5,
            "cacheReadInputTokens": 20,
            "cacheWriteInputTokens": 30,
        }
        handler(result=result)
        assert agent_span.attributes.get("gen_ai.usage.cache_read_input_tokens") == 20
        assert agent_span.attributes.get("gen_ai.usage.cache_creation_input_tokens") == 30

    def test_open_tool_spans_closed_on_result(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("tool", "t-r1", {})))
        tool_span = handler._active_tool_spans["t-r1"]
        handler(result=_make_result())
        assert tool_span.end_time is not None

    def test_result_with_no_metrics(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("ok"))
        result = MagicMock()
        result.message = {"role": "assistant", "content": [{"text": "ok"}]}
        del result.metrics
        handler(result=result)  # must not raise

    def test_result_with_empty_content(self, handler):
        handler(init_event_loop=True)
        agent_span = handler._agent_span
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("done"))
        result = _make_result()
        result.message = {"role": "assistant", "content": []}
        handler(result=result)
        assert "span.output" not in agent_span.attributes

    def test_result_output_from_text_blocks_without_type_field(self, handler):
        """Strands content blocks have no 'type' field — just {'text': '...'}."""
        handler(init_event_loop=True)
        agent_span = handler._agent_span
        handler(start_event_loop=True)
        handler(message=_assistant_message_text("done"))
        result = _make_result()
        # Strands actual format: {"text": "..."} without "type"
        result.message = {"role": "assistant", "content": [{"text": "The answer is 42."}]}
        handler(result=result)
        assert agent_span.attributes.get("span.output") == "The answer is 42."


# ---------------------------------------------------------------------------
# Force stop
# ---------------------------------------------------------------------------

class TestForceStop:
    def test_agent_span_error_status(self, handler):
        handler(init_event_loop=True)
        agent_span = handler._agent_span
        handler(start_event_loop=True)
        handler(force_stop=True, force_stop_reason="Max iterations")
        assert agent_span.status_code == StatusCode.ERROR

    def test_model_span_error_status(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        model_span = handler._model_span
        handler(force_stop=True)
        assert model_span.status_code == StatusCode.ERROR

    def test_spans_exported_on_force_stop(self, handler, mock_client):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(force_stop=True, force_stop_reason="Stopped")
        assert mock_client.export_span.call_count >= 2

    def test_tool_spans_closed_on_force_stop(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_assistant_message_with_tools(("tool", "t-fs", {})))
        tool_span = handler._active_tool_spans["t-fs"]
        handler(force_stop=True)
        assert tool_span.end_time is not None

    def test_force_stop_without_reason(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(force_stop=True)  # must not raise
        assert handler._agent_span is None

    def test_state_cleared_after_force_stop(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(force_stop=True)
        assert handler._agent_span is None
        assert handler._model_span is None
        assert handler._active_tool_spans == {}


# ---------------------------------------------------------------------------
# Error resilience
# ---------------------------------------------------------------------------

class TestErrorResilience:
    def test_handler_does_not_propagate_internal_errors(self, mock_client):
        """Errors inside the handler must be swallowed so the agent continues."""
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            handler = BeaconStrandsHandler()
        mock_client.export_eager_span.side_effect = RuntimeError("network down")
        handler(init_event_loop=True)  # must not raise

    def test_unknown_event_ignored(self, handler):
        handler(init_event_loop=True)
        handler(some_future_event=True, extra_data="whatever")  # must not raise

    def test_complete_without_prior_start_event_loop(self, handler):
        handler(init_event_loop=True)
        handler(complete=True)  # must not raise

    def test_result_without_prior_start(self, handler):
        handler(init_event_loop=True)
        handler(result=_make_result())  # must not raise

    def test_emergency_close_on_handler_error(self, handler, mock_client):
        """When a handler method raises, all open spans must be exported with ERROR status."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        # Keep the model span open (no assistant message) so _finalize_model_span runs in _on_result.
        # Inject bad state — "".join([dict]) raises TypeError inside _finalize_model_span.
        handler._streaming_buffer = [{"not": "a string"}]

        agent_span = handler._agent_span

        handler(result=_make_result())  # triggers TypeError → emergency_close

        # No spans should remain open
        assert handler._agent_span is None
        assert handler._model_span is None
        assert handler._active_tool_spans == {}

        # The agent span must have been exported with ERROR status
        assert agent_span.status_code == StatusCode.ERROR

    def test_data_as_dict_does_not_crash(self, handler):
        """data kwarg may be a dict with extended thinking — must not crash."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data={"text": "chunk from thinking model"})
        assert handler._streaming_buffer == ["chunk from thinking model"]

    def test_data_as_empty_dict_skipped(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(data={})
        assert handler._streaming_buffer == []

    def test_tool_result_for_unknown_tool_id_ignored(self, handler):
        """Tool result for an ID we don't have a span for must not crash."""
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message=_tool_result_message("unknown-id", "42"))  # must not raise

    def test_assistant_message_without_content_ignored(self, handler):
        handler(init_event_loop=True)
        handler(start_event_loop=True)
        handler(message={"role": "assistant", "content": []})  # must not raise
        # model span finalized (empty content = no tools, that's fine)


# ---------------------------------------------------------------------------
# Lazy import
# ---------------------------------------------------------------------------

class TestLazyImport:
    def test_importable_from_main_package(self, mock_client):
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            from lumenova_beacon import BeaconStrandsHandler as H
            assert H is BeaconStrandsHandler
