"""Tests for the CrewAI event listener integration."""

import json
from unittest.mock import MagicMock, patch
from typing import Any

import pytest

from lumenova_beacon import BeaconClient
from lumenova_beacon.tracing.integrations.crewai import BeaconCrewAIListener
from lumenova_beacon.types import SpanType, SpanKind, StatusCode

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

# Default IDs used across test helpers
DEFAULT_AGENT_ID = "agent-uuid-1"


# ---------------------------------------------------------------------------
# Fixtures & helpers
# ---------------------------------------------------------------------------


@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 listener(mock_client):
    """BeaconCrewAIListener with injected mock client, event bus bypassed."""
    with patch(PATCH_GET_CLIENT, return_value=mock_client):
        inst = BeaconCrewAIListener.__new__(BeaconCrewAIListener)
        inst._client = mock_client
        inst._session_id = None
        inst._environment = None
        inst._crew_name = None
        inst._metadata = {}
        inst._crew_span = None
        inst._trace_id = None
        inst._task_spans = {}
        inst._agent_spans = {}
        inst._llm_spans = {}
        inst._tool_spans = {}
        inst._current_agent_id = ""
        return inst


def _make_task(
    description: str = "Research the topic",
    task_id: str = "task-uuid-1",
) -> MagicMock:
    task = MagicMock()
    task.id = task_id
    task.name = None
    task.description = description
    return task


def _make_agent(
    role: str = "Researcher",
    goal: str = "Find info",
    backstory: str = "Expert",
    agent_id: str = DEFAULT_AGENT_ID,
) -> MagicMock:
    agent = MagicMock()
    agent.id = agent_id
    agent.role = role
    agent.goal = goal
    agent.backstory = backstory
    return agent


def _crew_started_event(
    crew_name: str = "TestCrew",
    inputs: dict | None = None,
) -> MagicMock:
    event = MagicMock()
    event.crew_name = crew_name
    event.inputs = inputs
    return event


def _crew_completed_event(
    output: str = "Final result",
    crew_name: str = "TestCrew",
    total_tokens: int = 0,
) -> MagicMock:
    event = MagicMock()
    event.crew_name = crew_name
    event.output = output
    event.total_tokens = total_tokens
    return event


def _crew_failed_event(error: str = "Crew crashed") -> MagicMock:
    event = MagicMock()
    event.error = error
    return event


def _task_started_event(
    task: MagicMock,
    context: str | None = None,
) -> MagicMock:
    event = MagicMock()
    event.task = task
    event.context = context
    return event


def _task_completed_event(
    task: MagicMock,
    output: str = "Task done",
) -> MagicMock:
    event = MagicMock()
    event.task = task
    event.output = output
    return event


def _task_failed_event(
    task: MagicMock,
    error: str = "Task failed",
) -> MagicMock:
    event = MagicMock()
    event.task = task
    event.error = error
    return event


def _agent_started_event(
    agent: MagicMock,
    task: MagicMock | None = None,
    task_prompt: str | None = None,
) -> MagicMock:
    event = MagicMock()
    event.agent = agent
    event.task = task
    event.task_prompt = task_prompt
    return event


def _agent_completed_event(
    agent: MagicMock,
    output: str = "Agent done",
) -> MagicMock:
    event = MagicMock()
    event.agent = agent
    event.output = output
    return event


def _agent_error_event(
    agent: MagicMock,
    error: str = "Agent error",
) -> MagicMock:
    event = MagicMock()
    event.agent = agent
    event.error = error
    return event


def _llm_started_event(
    model: str = "gpt-4o",
    messages: list | None = None,
    agent_id: str = DEFAULT_AGENT_ID,
) -> MagicMock:
    event = MagicMock()
    event.model = model
    event.messages = messages or [{"role": "user", "content": "Hello"}]
    event.agent_id = agent_id
    event.agent = None
    return event


def _llm_completed_event(
    response: Any = "LLM response",
    model: str | None = None,
    agent_id: str = DEFAULT_AGENT_ID,
) -> MagicMock:
    event = MagicMock()
    event.response = response
    event.model = model
    event.agent_id = agent_id
    event.agent = None
    return event


def _llm_failed_event(
    error: str = "LLM error",
    agent_id: str = DEFAULT_AGENT_ID,
) -> MagicMock:
    event = MagicMock()
    event.error = error
    event.agent_id = agent_id
    event.agent = None
    return event


def _tool_started_event(
    tool_name: str = "search",
    tool_input: Any = "query",
    agent_id: str = DEFAULT_AGENT_ID,
) -> MagicMock:
    event = MagicMock()
    event.tool_name = tool_name
    event.tool_args = tool_input
    event.agent_id = agent_id
    event.agent = None
    return event


def _tool_finished_event(
    tool_name: str = "search",
    output: str = "results",
    agent_id: str = DEFAULT_AGENT_ID,
) -> MagicMock:
    event = MagicMock()
    event.tool_name = tool_name
    event.output = output
    event.agent_id = agent_id
    event.agent = None
    return event


def _tool_error_event(
    tool_name: str = "search",
    error: str = "Tool failed",
    agent_id: str = DEFAULT_AGENT_ID,
) -> MagicMock:
    event = MagicMock()
    event.tool_name = tool_name
    event.error = error
    event.agent_id = agent_id
    event.agent = None
    return event


def _make_llm_response(
    content: str = "result text",
    prompt_tokens: int = 100,
    completion_tokens: int = 50,
) -> MagicMock:
    """Build a mock LiteLLM/OpenAI response with choices and usage."""
    response = MagicMock()
    choice = MagicMock()
    choice.message.content = content
    response.choices = [choice]
    response.usage.prompt_tokens = prompt_tokens
    response.usage.completion_tokens = completion_tokens
    response.usage.input_tokens = None
    response.usage.output_tokens = None
    return response


# ---------------------------------------------------------------------------
# TestBeaconCrewAIListenerInit
# ---------------------------------------------------------------------------


class TestBeaconCrewAIListenerInit:
    def test_init_with_defaults(self, mock_client):
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            inst = BeaconCrewAIListener()
        assert inst._session_id is None
        assert inst._environment is None
        assert inst._crew_name is None
        assert inst._metadata == {}
        assert inst._crew_span is None
        assert inst._trace_id is None
        assert inst._task_spans == {}
        assert inst._agent_spans == {}
        assert inst._llm_spans == {}
        assert inst._tool_spans == {}
        assert inst._current_agent_id == ""

    def test_init_with_all_params(self, mock_client):
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            inst = BeaconCrewAIListener(
                session_id="sess-1",
                environment="production",
                crew_name="My Crew",
                metadata={"project": "alpha"},
            )
        assert inst._session_id == "sess-1"
        assert inst._environment == "production"
        assert inst._crew_name == "My Crew"
        assert inst._metadata == {"project": "alpha"}

    def test_session_id_falls_back_to_client_config(self, mock_client):
        mock_client.config.session_id = "client-session"
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            inst = BeaconCrewAIListener()
        assert inst._session_id == "client-session"

    def test_explicit_session_id_overrides_client(self, mock_client):
        mock_client.config.session_id = "client-session"
        with patch(PATCH_GET_CLIENT, return_value=mock_client):
            inst = BeaconCrewAIListener(session_id="explicit-session")
        assert inst._session_id == "explicit-session"

    def test_trace_id_property_initially_none(self, listener):
        assert listener.trace_id is None


# ---------------------------------------------------------------------------
# TestCrewLifecycle
# ---------------------------------------------------------------------------


class TestCrewLifecycle:
    def test_crew_started_creates_span(self, listener, mock_client):
        source = MagicMock()
        event = _crew_started_event("MyCrew")
        listener._on_crew_started(source, event)

        assert listener._crew_span is not None
        assert listener._crew_span.name == "MyCrew"
        assert listener._crew_span.span_type == SpanType.AGENT
        assert listener._crew_span.kind == SpanKind.INTERNAL
        mock_client.export_eager_span.assert_called_once_with(listener._crew_span)

    def test_crew_started_sets_trace_id(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        assert listener.trace_id is not None
        assert len(listener.trace_id) == 32  # 32-char hex

    def test_crew_name_from_constructor_takes_priority(self, listener):
        listener._crew_name = "OverrideName"
        listener._on_crew_started(MagicMock(), _crew_started_event("EventName"))
        assert listener._crew_span.name == "OverrideName"

    def test_crew_name_falls_back_to_event(self, listener):
        listener._crew_name = None
        listener._on_crew_started(MagicMock(), _crew_started_event("FromEvent"))
        assert listener._crew_span.name == "FromEvent"

    def test_crew_name_falls_back_to_default(self, listener):
        listener._crew_name = None
        event = MagicMock()
        event.crew_name = None
        listener._on_crew_started(MagicMock(), event)
        assert listener._crew_span.name == "crewai"

    def test_crew_started_sets_attributes(self, listener):
        listener._session_id = "sess-abc"
        listener._environment = "staging"
        listener._metadata = {"team": "backend"}
        listener._on_crew_started(MagicMock(), _crew_started_event("Crew"))
        attrs = listener._crew_span.attributes
        assert attrs.get("gen_ai.agent.name") == "Crew"
        assert attrs.get("gen_ai.agent.framework") == "crewai"
        assert attrs.get("gen_ai.operation.name") == "invoke_agent"
        assert attrs.get("gen_ai.provider.name") == "crewai"
        assert attrs.get("gen_ai.conversation.id") == "sess-abc"
        assert attrs.get("deployment.environment.name") == "staging"
        assert attrs.get("beacon.metadata.team") == "backend"

    def test_crew_started_captures_inputs(self, listener):
        event = _crew_started_event("Crew", inputs={"topic": "AI safety"})
        listener._on_crew_started(MagicMock(), event)
        assert listener._crew_span.attributes.get("span.input") is not None

    def test_crew_input_falls_back_to_first_task_description(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event("Crew", inputs=None))
        assert listener._crew_span.attributes.get("span.input") is None
        task = _make_task(description="Calculate: 5+4")
        listener._on_task_started(MagicMock(), _task_started_event(task))
        # First task description should bubble up to crew span
        assert listener._crew_span.attributes.get("span.input") is not None

    def test_crew_input_not_overwritten_by_task_description(self, listener):
        event = _crew_started_event("Crew", inputs={"topic": "AI safety"})
        listener._on_crew_started(MagicMock(), event)
        assert listener._crew_span.attributes.get("span.input") is not None
        task = _make_task(description="Calculate: 5+4")
        listener._on_task_started(MagicMock(), _task_started_event(task))
        # Crew span should still have the kickoff inputs, not task description
        raw = listener._crew_span.attributes.get("span.input")
        assert "AI safety" in str(raw)

    def test_crew_completed_finalizes_span(self, listener, mock_client):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        crew_span = listener._crew_span

        listener._on_crew_completed(MagicMock(), _crew_completed_event("Done!"))

        assert listener._crew_span is None
        assert crew_span.status_code == StatusCode.OK
        assert crew_span.attributes.get("span.output") == "Done!"
        mock_client.export_span.assert_called_once_with(crew_span)

    def test_crew_completed_sets_total_tokens(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        crew_span = listener._crew_span
        listener._on_crew_completed(
            MagicMock(), _crew_completed_event(total_tokens=1500)
        )
        assert crew_span.attributes.get("gen_ai.usage.total_tokens") == 1500

    def test_crew_completed_calls_clear_context(self, listener):
        with patch("lumenova_beacon.tracing.integrations.crewai.clear_context") as mock_clear:
            listener._on_crew_started(MagicMock(), _crew_started_event())
            listener._on_crew_completed(MagicMock(), _crew_completed_event())
            mock_clear.assert_called_once()

    def test_crew_failed_sets_error_status(self, listener, mock_client):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        crew_span = listener._crew_span

        listener._on_crew_failed(MagicMock(), _crew_failed_event("Crew crashed"))

        assert listener._crew_span is None
        assert crew_span.status_code == StatusCode.ERROR
        assert crew_span.attributes.get("error.type") is not None
        mock_client.export_span.assert_called_once_with(crew_span)

    def test_crew_completed_noop_without_crew_span(self, listener, mock_client):
        listener._on_crew_completed(MagicMock(), _crew_completed_event())
        mock_client.export_span.assert_not_called()

    def test_crew_failed_noop_without_crew_span(self, listener, mock_client):
        listener._on_crew_failed(MagicMock(), _crew_failed_event())
        mock_client.export_span.assert_not_called()


# ---------------------------------------------------------------------------
# TestTaskLifecycle
# ---------------------------------------------------------------------------


class TestTaskLifecycle:
    def _start_crew(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())

    def test_task_started_creates_span(self, listener, mock_client):
        self._start_crew(listener)
        task = _make_task("Do research")
        listener._on_task_started(MagicMock(), _task_started_event(task))

        assert id(task) in listener._task_spans
        task_span = listener._task_spans[id(task)]
        assert task_span.span_type == SpanType.TASK
        assert task_span.kind == SpanKind.INTERNAL
        mock_client.export_eager_span.assert_any_call(task_span)

    def test_task_span_name_uses_description(self, listener):
        self._start_crew(listener)
        task = _make_task(description="Analyze the data", task_id="t-1")
        listener._on_task_started(MagicMock(), _task_started_event(task))
        span = listener._task_spans[id(task)]
        assert "Analyze the data" in span.name

    def test_task_span_name_truncated_to_60_chars(self, listener):
        self._start_crew(listener)
        task = _make_task(description="A" * 100)
        listener._on_task_started(MagicMock(), _task_started_event(task))
        span = listener._task_spans[id(task)]
        assert len(span.name) <= 60

    def test_task_span_is_child_of_crew_span(self, listener):
        self._start_crew(listener)
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        task_span = listener._task_spans[id(task)]
        assert task_span.parent_id == listener._crew_span.span_id
        assert task_span.trace_id == listener.trace_id

    def test_task_span_sets_task_id_attribute(self, listener):
        self._start_crew(listener)
        task = _make_task(task_id="uuid-999")
        listener._on_task_started(MagicMock(), _task_started_event(task))
        span = listener._task_spans[id(task)]
        assert span.attributes.get("crewai.task.id") == "uuid-999"

    def test_task_started_sets_description_as_input(self, listener):
        self._start_crew(listener)
        task = _make_task(description="Summarize the article")
        listener._on_task_started(MagicMock(), _task_started_event(task))
        span = listener._task_spans[id(task)]
        assert span.attributes.get("span.input") is not None

    def test_task_completed_closes_span(self, listener, mock_client):
        self._start_crew(listener)
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        task_span = listener._task_spans[id(task)]

        listener._on_task_completed(MagicMock(), _task_completed_event(task, "Task output"))

        assert id(task) not in listener._task_spans
        assert task_span.status_code == StatusCode.OK
        assert task_span.attributes.get("span.output") == "Task output"
        mock_client.export_span.assert_called_once_with(task_span)

    def test_task_failed_sets_error(self, listener, mock_client):
        self._start_crew(listener)
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        task_span = listener._task_spans[id(task)]

        listener._on_task_failed(MagicMock(), _task_failed_event(task, "timeout"))

        assert id(task) not in listener._task_spans
        assert task_span.status_code == StatusCode.ERROR
        mock_client.export_span.assert_called_once_with(task_span)

    def test_multiple_tasks_tracked_independently(self, listener):
        self._start_crew(listener)
        task_a = _make_task("Task A", "id-a")
        task_b = _make_task("Task B", "id-b")
        listener._on_task_started(MagicMock(), _task_started_event(task_a))
        listener._on_task_started(MagicMock(), _task_started_event(task_b))

        assert id(task_a) in listener._task_spans
        assert id(task_b) in listener._task_spans
        assert listener._task_spans[id(task_a)] is not listener._task_spans[id(task_b)]

    def test_task_started_noop_without_task(self, listener, mock_client):
        self._start_crew(listener)
        event = MagicMock()
        event.task = None
        listener._on_task_started(MagicMock(), event)
        # Only the crew eager export, no task export
        mock_client.export_eager_span.assert_called_once()


# ---------------------------------------------------------------------------
# TestAgentLifecycle
# ---------------------------------------------------------------------------


class TestAgentLifecycle:
    def _start_crew_and_task(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        return task

    def test_agent_started_creates_span(self, listener, mock_client):
        task = self._start_crew_and_task(listener)
        agent = _make_agent("Researcher")
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))

        agent_id = str(agent.id)
        assert agent_id in listener._agent_spans
        span = listener._agent_spans[agent_id]
        assert span.name == "Researcher"
        assert span.span_type == SpanType.AGENT
        mock_client.export_eager_span.assert_any_call(span)

    def test_agent_span_is_child_of_task_span(self, listener):
        task = self._start_crew_and_task(listener)
        task_span = listener._task_spans[id(task)]
        agent = _make_agent("Researcher")
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))

        agent_span = listener._agent_spans[str(agent.id)]
        assert agent_span.parent_id == task_span.span_id
        assert agent_span.trace_id == listener.trace_id

    def test_agent_span_falls_back_to_crew_when_no_task(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        agent = _make_agent()
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task=None))
        span = listener._agent_spans[str(agent.id)]
        assert span.parent_id == listener._crew_span.span_id

    def test_agent_started_sets_attributes(self, listener):
        task = self._start_crew_and_task(listener)
        agent = _make_agent("Researcher", "Find insights", "Expert in analysis")
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        span = listener._agent_spans[str(agent.id)]
        assert span.attributes.get("gen_ai.agent.name") == "Researcher"
        assert span.attributes.get("gen_ai.agent.description") == "Find insights"
        assert span.attributes.get("crewai.agent.backstory") == "Expert in analysis"

    def test_agent_started_captures_task_prompt(self, listener):
        task = self._start_crew_and_task(listener)
        agent = _make_agent()
        listener._on_agent_started(
            MagicMock(),
            _agent_started_event(agent, task, task_prompt="Calculate 5+4"),
        )
        span = listener._agent_spans[str(agent.id)]
        assert span.attributes.get("span.input") is not None

    def test_agent_started_sets_current_agent_id(self, listener):
        task = self._start_crew_and_task(listener)
        agent = _make_agent(agent_id="my-agent-99")
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        assert listener._current_agent_id == "my-agent-99"

    def test_agent_completed_closes_span(self, listener, mock_client):
        task = self._start_crew_and_task(listener)
        agent = _make_agent()
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        agent_id = str(agent.id)
        agent_span = listener._agent_spans[agent_id]

        listener._on_agent_completed(MagicMock(), _agent_completed_event(agent, "Agent result"))

        assert agent_id not in listener._agent_spans
        assert agent_span.status_code == StatusCode.OK
        assert agent_span.attributes.get("span.output") == "Agent result"
        mock_client.export_span.assert_any_call(agent_span)

    def test_agent_completed_clears_current_agent_id(self, listener):
        task = self._start_crew_and_task(listener)
        agent = _make_agent()
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        assert listener._current_agent_id == str(agent.id)

        listener._on_agent_completed(MagicMock(), _agent_completed_event(agent))
        assert listener._current_agent_id == ""

    def test_agent_error_sets_error_status(self, listener, mock_client):
        task = self._start_crew_and_task(listener)
        agent = _make_agent()
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        agent_id = str(agent.id)
        agent_span = listener._agent_spans[agent_id]

        listener._on_agent_error(MagicMock(), _agent_error_event(agent, "something went wrong"))

        assert agent_id not in listener._agent_spans
        assert agent_span.status_code == StatusCode.ERROR
        mock_client.export_span.assert_any_call(agent_span)

    def test_agent_error_clears_current_agent_id(self, listener):
        task = self._start_crew_and_task(listener)
        agent = _make_agent()
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        listener._on_agent_error(MagicMock(), _agent_error_event(agent))
        assert listener._current_agent_id == ""

    def test_agent_started_noop_without_agent(self, listener, mock_client):
        self._start_crew_and_task(listener)
        event = MagicMock()
        event.agent = None
        listener._on_agent_started(MagicMock(), event)
        assert listener._agent_spans == {}


# ---------------------------------------------------------------------------
# TestLLMCallLifecycle
# ---------------------------------------------------------------------------


class TestLLMCallLifecycle:
    def _start_with_agent(self, listener, agent_id: str = DEFAULT_AGENT_ID):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        agent = _make_agent(agent_id=agent_id)
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        return agent

    def test_llm_started_creates_span(self, listener, mock_client):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_llm_started(MagicMock(), _llm_started_event("gpt-4o"))

        assert agent_id in listener._llm_spans
        llm_span = listener._llm_spans[agent_id]
        assert llm_span.span_type == SpanType.GENERATION
        assert llm_span.kind == SpanKind.CLIENT
        mock_client.export_eager_span.assert_any_call(llm_span)

    def test_llm_span_is_child_of_agent_span(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        agent_span = listener._agent_spans[agent_id]
        listener._on_llm_started(MagicMock(), _llm_started_event())
        llm_span = listener._llm_spans[agent_id]
        assert llm_span.parent_id == agent_span.span_id
        assert llm_span.trace_id == listener.trace_id

    def test_llm_started_sets_model_attribute(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_llm_started(
            MagicMock(),
            _llm_started_event("claude-sonnet-4-6"),
        )
        llm_span = listener._llm_spans[agent_id]
        assert llm_span.attributes.get("gen_ai.request.model") == "claude-sonnet-4-6"
        assert llm_span.attributes.get("gen_ai.response.model") == "claude-sonnet-4-6"

    def test_llm_started_sets_messages_as_input(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        messages = [{"role": "user", "content": "Hello!"}]
        listener._on_llm_started(
            MagicMock(),
            _llm_started_event(messages=messages),
        )
        llm_span = listener._llm_spans[agent_id]
        assert llm_span.attributes.get("span.input") is not None

    def test_llm_started_sets_gen_ai_input_messages(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        messages = [{"role": "user", "content": "Hello!"}]
        listener._on_llm_started(
            MagicMock(),
            _llm_started_event(messages=messages),
        )
        llm_span = listener._llm_spans[agent_id]
        raw = llm_span.attributes.get("gen_ai.input.messages")
        assert raw is not None
        parsed = json.loads(raw)
        assert parsed[0]["role"] == "user"

    def test_llm_completed_closes_span(self, listener, mock_client):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_llm_started(MagicMock(), _llm_started_event())
        llm_span = listener._llm_spans[agent_id]

        listener._on_llm_completed(
            MagicMock(),
            _llm_completed_event("The answer"),
        )

        assert agent_id not in listener._llm_spans
        assert llm_span.status_code == StatusCode.OK
        assert llm_span.attributes.get("span.output") == "The answer"
        mock_client.export_span.assert_any_call(llm_span)

    def test_llm_completed_sets_gen_ai_output_messages(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_llm_started(MagicMock(), _llm_started_event())
        llm_span = listener._llm_spans[agent_id]

        listener._on_llm_completed(
            MagicMock(),
            _llm_completed_event("The answer"),
        )

        raw = llm_span.attributes.get("gen_ai.output.messages")
        assert raw is not None
        parsed = json.loads(raw)
        assert parsed[0]["role"] == "assistant"
        assert parsed[0]["content"] == "The answer"

    def test_llm_completed_extracts_token_usage(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_llm_started(MagicMock(), _llm_started_event())
        llm_span = listener._llm_spans[agent_id]

        response = _make_llm_response("result text", prompt_tokens=100, completion_tokens=50)
        listener._on_llm_completed(
            MagicMock(),
            _llm_completed_event(response=response),
        )

        assert llm_span.attributes.get("gen_ai.usage.input_tokens") == 100
        assert llm_span.attributes.get("gen_ai.usage.output_tokens") == 50

    def test_llm_completed_sets_response_model(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_llm_started(MagicMock(), _llm_started_event())
        llm_span = listener._llm_spans[agent_id]

        listener._on_llm_completed(
            MagicMock(),
            _llm_completed_event("answer", model="gpt-4o-2024-08-06"),
        )

        assert llm_span.attributes.get("gen_ai.response.model") == "gpt-4o-2024-08-06"

    def test_llm_failed_sets_error(self, listener, mock_client):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_llm_started(MagicMock(), _llm_started_event())
        llm_span = listener._llm_spans[agent_id]

        listener._on_llm_failed(
            MagicMock(),
            _llm_failed_event("rate limit"),
        )

        assert agent_id not in listener._llm_spans
        assert llm_span.status_code == StatusCode.ERROR
        assert llm_span.attributes.get("error.type") is not None
        mock_client.export_span.assert_any_call(llm_span)

    def test_llm_completed_noop_without_prior_start(self, listener, mock_client):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        listener._on_llm_completed(
            MagicMock(),
            _llm_completed_event(agent_id="no-match"),
        )
        mock_client.export_span.assert_not_called()


# ---------------------------------------------------------------------------
# TestToolLifecycle
# ---------------------------------------------------------------------------


class TestToolLifecycle:
    def _start_with_agent(self, listener, agent_id: str = DEFAULT_AGENT_ID):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        agent = _make_agent(agent_id=agent_id)
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        return agent

    def test_tool_started_creates_span(self, listener, mock_client):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("web_search", agent_id=agent_id),
        )

        tool_key = (agent_id, "web_search")
        assert tool_key in listener._tool_spans
        span = listener._tool_spans[tool_key]
        assert span.span_type == SpanType.TOOL
        assert span.kind == SpanKind.INTERNAL
        mock_client.export_eager_span.assert_any_call(span)

    def test_tool_span_is_child_of_agent_span(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        agent_span = listener._agent_spans[agent_id]
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("web_search", agent_id=agent_id),
        )
        tool_span = listener._tool_spans[(agent_id, "web_search")]
        assert tool_span.parent_id == agent_span.span_id
        assert tool_span.trace_id == listener.trace_id

    def test_tool_started_sets_tool_name_attribute(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("calculator", agent_id=agent_id),
        )
        span = listener._tool_spans[(agent_id, "calculator")]
        assert span.attributes.get("gen_ai.tool.name") == "calculator"

    def test_tool_started_sets_tool_type(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("search", agent_id=agent_id),
        )
        span = listener._tool_spans[(agent_id, "search")]
        assert span.attributes.get("gen_ai.tool.type") == "function"

    def test_tool_started_sets_input_and_arguments(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("search", tool_input={"q": "python docs"}, agent_id=agent_id),
        )
        span = listener._tool_spans[(agent_id, "search")]
        assert span.attributes.get("span.input") is not None
        raw = span.attributes.get("gen_ai.tool.call.arguments")
        assert raw is not None
        parsed = json.loads(raw)
        assert parsed["q"] == "python docs"

    def test_tool_finished_closes_span(self, listener, mock_client):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("web_search", agent_id=agent_id),
        )
        tool_span = listener._tool_spans[(agent_id, "web_search")]

        listener._on_tool_finished(
            MagicMock(),
            _tool_finished_event("web_search", "Results found", agent_id=agent_id),
        )

        assert (agent_id, "web_search") not in listener._tool_spans
        assert tool_span.status_code == StatusCode.OK
        assert tool_span.attributes.get("span.output") == "Results found"
        mock_client.export_span.assert_any_call(tool_span)

    def test_tool_finished_sets_result_attribute(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("search", agent_id=agent_id),
        )
        tool_span = listener._tool_spans[(agent_id, "search")]

        listener._on_tool_finished(
            MagicMock(),
            _tool_finished_event("search", "42", agent_id=agent_id),
        )
        assert tool_span.attributes.get("gen_ai.tool.call.result") == "42"

    def test_tool_error_sets_error_status(self, listener, mock_client):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("web_search", agent_id=agent_id),
        )
        tool_span = listener._tool_spans[(agent_id, "web_search")]

        listener._on_tool_error(
            MagicMock(),
            _tool_error_event("web_search", "Connection timeout", agent_id=agent_id),
        )

        assert (agent_id, "web_search") not in listener._tool_spans
        assert tool_span.status_code == StatusCode.ERROR
        mock_client.export_span.assert_any_call(tool_span)

    def test_multiple_tools_tracked_independently(self, listener):
        agent = self._start_with_agent(listener)
        agent_id = str(agent.id)
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("search", agent_id=agent_id),
        )
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("calculator", agent_id=agent_id),
        )

        assert (agent_id, "search") in listener._tool_spans
        assert (agent_id, "calculator") in listener._tool_spans

    def test_tool_finished_noop_without_prior_start(self, listener, mock_client):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        listener._on_tool_finished(MagicMock(), _tool_finished_event("search"))
        mock_client.export_span.assert_not_called()


# ---------------------------------------------------------------------------
# TestResolveAgentSpan
# ---------------------------------------------------------------------------


class TestResolveAgentSpan:
    """Test the _resolve_agent_span fallback logic."""

    def _setup_agent(self, listener, agent_id: str = DEFAULT_AGENT_ID):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        agent = _make_agent(agent_id=agent_id)
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        return agent

    def test_resolves_via_event_agent_id(self, listener):
        """Strategy 1: event.agent_id matches stored key."""
        self._setup_agent(listener, "agent-A")
        agent_span = listener._agent_spans["agent-A"]

        event = MagicMock()
        event.agent_id = "agent-A"
        event.agent = None

        resolved_id, resolved_span = listener._resolve_agent_span(event)
        assert resolved_id == "agent-A"
        assert resolved_span is agent_span

    def test_resolves_via_event_agent_object(self, listener):
        """Strategy 2: event.agent.id matches stored key."""
        self._setup_agent(listener, "agent-B")
        agent_span = listener._agent_spans["agent-B"]

        event = MagicMock()
        event.agent_id = "wrong-id"
        event.agent = MagicMock()
        event.agent.id = "agent-B"

        resolved_id, resolved_span = listener._resolve_agent_span(event)
        assert resolved_id == "agent-B"
        assert resolved_span is agent_span

    def test_resolves_via_current_agent_id_fallback(self, listener):
        """Strategy 3: _current_agent_id fallback."""
        self._setup_agent(listener, "agent-C")
        agent_span = listener._agent_spans["agent-C"]
        assert listener._current_agent_id == "agent-C"

        event = MagicMock()
        event.agent_id = ""
        event.agent = None

        resolved_id, resolved_span = listener._resolve_agent_span(event)
        assert resolved_id == "agent-C"
        assert resolved_span is agent_span

    def test_falls_back_to_crew_span_when_no_match(self, listener):
        """All strategies fail — falls back to crew span."""
        listener._on_crew_started(MagicMock(), _crew_started_event())
        crew_span = listener._crew_span

        event = MagicMock()
        event.agent_id = "nonexistent"
        event.agent = None

        resolved_id, resolved_span = listener._resolve_agent_span(event)
        assert resolved_span is crew_span

    def test_tool_span_parents_correctly_via_fallback(self, listener):
        """End-to-end: tool event with empty agent_id still parents to agent."""
        self._setup_agent(listener, "agent-D")
        agent_span = listener._agent_spans["agent-D"]

        # Tool event with empty agent_id — resolves via _current_agent_id
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("add", agent_id=""),
        )
        tool_span = listener._tool_spans[("agent-D", "add")]
        assert tool_span.parent_id == agent_span.span_id


# ---------------------------------------------------------------------------
# TestFullHierarchy
# ---------------------------------------------------------------------------


class TestFullHierarchy:
    def test_full_span_hierarchy(self, listener, mock_client):
        """Verify trace_id and parent_id chain through the full hierarchy."""
        listener._on_crew_started(MagicMock(), _crew_started_event("MyCrew"))
        crew_span = listener._crew_span
        trace_id = listener.trace_id

        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        task_span = listener._task_spans[id(task)]

        agent = _make_agent(agent_id="agent-full")
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        agent_span = listener._agent_spans["agent-full"]

        listener._on_llm_started(
            MagicMock(),
            _llm_started_event(agent_id="agent-full"),
        )
        llm_span = listener._llm_spans["agent-full"]

        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("search", agent_id="agent-full"),
        )
        tool_span = listener._tool_spans[("agent-full", "search")]

        # Verify all spans share the same trace_id
        for span in [crew_span, task_span, agent_span, llm_span, tool_span]:
            assert span.trace_id == trace_id

        # Verify parent chain
        assert task_span.parent_id == crew_span.span_id
        assert agent_span.parent_id == task_span.span_id
        assert llm_span.parent_id == agent_span.span_id
        assert tool_span.parent_id == agent_span.span_id

    def test_export_order(self, listener, mock_client):
        """eager exports on start, final exports on completion."""
        listener._on_crew_started(MagicMock(), _crew_started_event())
        crew_span = listener._crew_span

        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        task_span = listener._task_spans[id(task)]

        listener._on_task_completed(MagicMock(), _task_completed_event(task))
        listener._on_crew_completed(MagicMock(), _crew_completed_event())

        eager_calls = [c.args[0] for c in mock_client.export_eager_span.call_args_list]
        final_calls = [c.args[0] for c in mock_client.export_span.call_args_list]
        assert crew_span in eager_calls
        assert task_span in eager_calls
        assert task_span in final_calls
        assert crew_span in final_calls


# ---------------------------------------------------------------------------
# TestErrorResilience
# ---------------------------------------------------------------------------


class TestErrorResilience:
    def test_internal_exception_swallowed_in_crew_started(self, listener):
        """Handler exceptions do not propagate to caller."""
        listener._client.export_eager_span.side_effect = RuntimeError("export broke")
        listener._on_crew_started(MagicMock(), _crew_started_event())

    def test_internal_exception_swallowed_in_task_completed(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        listener._client.export_span.side_effect = RuntimeError("export broke")
        listener._on_task_completed(MagicMock(), _task_completed_event(task))

    def test_missing_event_field_handled_gracefully(self, listener, mock_client):
        """getattr with default handles missing event attributes."""
        listener._on_crew_started(MagicMock(), _crew_started_event())
        event = MagicMock(spec=[])  # no attributes defined
        listener._on_crew_completed(MagicMock(), event)
        assert listener._crew_span is None  # span still finalized

    def test_crew_completed_noop_when_no_crew_span(self, listener, mock_client):
        listener._on_crew_completed(MagicMock(), _crew_completed_event())
        mock_client.export_span.assert_not_called()

    def test_task_completed_noop_when_task_not_tracked(self, listener, mock_client):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_completed(MagicMock(), _task_completed_event(task))
        mock_client.export_span.assert_not_called()

    def test_agent_completed_noop_when_not_tracked(self, listener, mock_client):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        agent = _make_agent()
        listener._on_agent_completed(MagicMock(), _agent_completed_event(agent))
        mock_client.export_span.assert_not_called()


# ---------------------------------------------------------------------------
# TestAttributes
# ---------------------------------------------------------------------------


class TestAttributes:
    def test_gen_ai_agent_name_set_on_crew_span(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event("ResearchCrew"))
        assert listener._crew_span.attributes.get("gen_ai.agent.name") == "ResearchCrew"

    def test_gen_ai_agent_name_set_on_agent_span(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        agent = _make_agent(role="Writer")
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        span = listener._agent_spans[str(agent.id)]
        assert span.attributes.get("gen_ai.agent.name") == "Writer"

    def test_crewai_task_description_attribute(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task(description="Summarize the document")
        listener._on_task_started(MagicMock(), _task_started_event(task))
        span = listener._task_spans[id(task)]
        assert span.attributes.get("crewai.task.description") == "Summarize the document"

    def test_gen_ai_request_model_on_llm_span(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        agent = _make_agent()
        agent_id = str(agent.id)
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        listener._on_llm_started(
            MagicMock(),
            _llm_started_event("gpt-4o-mini"),
        )
        llm_span = listener._llm_spans[agent_id]
        assert llm_span.attributes.get("gen_ai.request.model") == "gpt-4o-mini"

    def test_gen_ai_tool_name_on_tool_span(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        task = _make_task()
        listener._on_task_started(MagicMock(), _task_started_event(task))
        agent = _make_agent()
        agent_id = str(agent.id)
        listener._on_agent_started(MagicMock(), _agent_started_event(agent, task))
        listener._on_tool_started(
            MagicMock(),
            _tool_started_event("DuckDuckGo", agent_id=agent_id),
        )
        tool_span = listener._tool_spans[(agent_id, "DuckDuckGo")]
        assert tool_span.attributes.get("gen_ai.tool.name") == "DuckDuckGo"

    def test_metadata_prefixed_with_beacon(self, listener):
        listener._metadata = {"version": "2", "env": "prod"}
        listener._on_crew_started(MagicMock(), _crew_started_event())
        attrs = listener._crew_span.attributes
        assert attrs.get("beacon.metadata.version") == "2"
        assert attrs.get("beacon.metadata.env") == "prod"

    def test_gen_ai_operation_name_on_crew_span(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        assert listener._crew_span.attributes.get("gen_ai.operation.name") == "invoke_agent"

    def test_gen_ai_provider_name_on_crew_span(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        assert listener._crew_span.attributes.get("gen_ai.provider.name") == "crewai"

    def test_error_type_set_on_crew_failure(self, listener):
        listener._on_crew_started(MagicMock(), _crew_started_event())
        crew_span = listener._crew_span
        listener._on_crew_failed(MagicMock(), _crew_failed_event("boom"))
        assert crew_span.attributes.get("error.type") is not None
