"""CrewAI integration using Beacon SDK Spans (event listener approach).

This module provides an event listener that traces CrewAI Crew executions
using SDK Span objects, which are exported via OTLP.

Usage:
    from lumenova_beacon import BeaconClient
    from lumenova_beacon.tracing.integrations.crewai import BeaconCrewAIListener
    from crewai import Agent, Crew, Task

    # Initialize BeaconClient first
    client = BeaconClient()

    # Instantiate the listener — it auto-registers with the CrewAI event bus
    listener = BeaconCrewAIListener(
        session_id="my-session",
        crew_name="My Research Crew",
    )

    # Define and run your Crew as normal
    crew = Crew(agents=[...], tasks=[...])
    result = crew.kickoff()
    print(listener.trace_id)  # Link to Beacon trace

Span hierarchy produced:
    crew_span   (SpanType.AGENT)
    └── task_span   (SpanType.TASK)         - one per task
        └── agent_span  (SpanType.AGENT)    - one per agent execution
            ├── llm_span   (SpanType.GENERATION) - each LLM call
            └── tool_span  (SpanType.TOOL)  - each tool usage

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
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:
    from crewai.events import BaseEventListener, crewai_event_bus  # noqa: F401
except ImportError:
    raise ImportError(
        "crewai is required for the CrewAI integration. "
        "Install it with: pip install 'lumenova-beacon[crewai]'"
    )

logger = logging.getLogger(__name__)


class BeaconCrewAIListener(BaseEventListener):
    """Beacon event listener for CrewAI Crews.

    Instantiate this class before running your Crew. It automatically
    registers with the CrewAI event bus and creates Beacon spans covering
    the full crew lifecycle.

    Args:
        session_id: Override session ID for all spans (defaults to
            ``BeaconClient.config.session_id``).
        environment: Deployment environment tag stored as span metadata.
        crew_name: Human-readable name for the root crew span. If not
            provided, falls back to the crew name from the kickoff event.
        metadata: Additional key-value pairs stored as ``beacon.metadata.*``
            attributes on the crew span.
    """

    # OTEL GenAI semantic convention attributes
    ATTR_GEN_AI_AGENT_NAME = "gen_ai.agent.name"
    ATTR_GEN_AI_AGENT_FRAMEWORK = "gen_ai.agent.framework"
    ATTR_GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id"
    ATTR_GEN_AI_REQUEST_MODEL = "gen_ai.request.model"
    ATTR_GEN_AI_RESPONSE_MODEL = "gen_ai.response.model"
    ATTR_GEN_AI_INPUT_MESSAGES = "gen_ai.input.messages"
    ATTR_GEN_AI_OUTPUT_MESSAGES = "gen_ai.output.messages"
    ATTR_GEN_AI_INPUT_TOKENS = "gen_ai.usage.input_tokens"
    ATTR_GEN_AI_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
    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_OPERATION_NAME = "gen_ai.operation.name"
    ATTR_GEN_AI_PROVIDER_NAME = "gen_ai.provider.name"

    # OTEL GenAI agent attributes
    ATTR_GEN_AI_AGENT_DESCRIPTION = "gen_ai.agent.description"

    # CrewAI-specific attributes
    ATTR_CREWAI_TASK_ID = "crewai.task.id"
    ATTR_CREWAI_TASK_DESCRIPTION = "crewai.task.description"
    ATTR_CREWAI_AGENT_BACKSTORY = "crewai.agent.backstory"

    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,
        crew_name: str | None = None,
        metadata: dict[str, Any] | None = None,
    ):
        super().__init__()
        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._crew_name = crew_name
        self._metadata = metadata or {}

        # Root span state
        self._crew_span: Span | None = None
        self._trace_id: str | None = None

        # Span tracking.
        # Agent spans keyed by CrewAI agent ID string (stable across events).
        # LLM spans keyed by resolved agent_id — CrewAI executes LLM calls
        # sequentially per agent, so at most one active LLM call per agent.
        # Tool spans keyed by (agent_id, tool_name) — sequential execution
        # ensures no collisions for the same agent + tool pair.
        self._task_spans: dict[int, Span] = {}
        self._agent_spans: dict[str, Span] = {}
        self._llm_spans: dict[str, Span] = {}
        self._tool_spans: dict[tuple[str, str], Span] = {}

        # Fallback tracker: most recently started agent ID.
        # Used when tool/LLM events carry an agent_id that doesn't
        # match the key we stored in _agent_spans.
        self._current_agent_id: str = ""

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

    # ------------------------------------------------------------------
    # BaseEventListener protocol
    # ------------------------------------------------------------------

    def setup_listeners(self, crewai_event_bus: Any) -> None:  # type: ignore[override]
        """Register all event handlers on the CrewAI event bus."""
        from crewai.events.types.crew_events import (
            CrewKickoffStartedEvent,
            CrewKickoffCompletedEvent,
            CrewKickoffFailedEvent,
        )
        from crewai.events.types.task_events import (
            TaskStartedEvent,
            TaskCompletedEvent,
            TaskFailedEvent,
        )
        from crewai.events.types.agent_events import (
            AgentExecutionStartedEvent,
            AgentExecutionCompletedEvent,
            AgentExecutionErrorEvent,
        )
        from crewai.events.types.tool_usage_events import (
            ToolUsageStartedEvent,
            ToolUsageFinishedEvent,
            ToolUsageErrorEvent,
        )
        from crewai.events.types.llm_events import (
            LLMCallStartedEvent,
            LLMCallCompletedEvent,
            LLMCallFailedEvent,
        )

        @crewai_event_bus.on(CrewKickoffStartedEvent)
        def on_crew_started(source: Any, event: Any) -> None:
            self._on_crew_started(source, event)

        @crewai_event_bus.on(CrewKickoffCompletedEvent)
        def on_crew_completed(source: Any, event: Any) -> None:
            self._on_crew_completed(source, event)

        @crewai_event_bus.on(CrewKickoffFailedEvent)
        def on_crew_failed(source: Any, event: Any) -> None:
            self._on_crew_failed(source, event)

        @crewai_event_bus.on(TaskStartedEvent)
        def on_task_started(source: Any, event: Any) -> None:
            self._on_task_started(source, event)

        @crewai_event_bus.on(TaskCompletedEvent)
        def on_task_completed(source: Any, event: Any) -> None:
            self._on_task_completed(source, event)

        @crewai_event_bus.on(TaskFailedEvent)
        def on_task_failed(source: Any, event: Any) -> None:
            self._on_task_failed(source, event)

        @crewai_event_bus.on(AgentExecutionStartedEvent)
        def on_agent_started(source: Any, event: Any) -> None:
            self._on_agent_started(source, event)

        @crewai_event_bus.on(AgentExecutionCompletedEvent)
        def on_agent_completed(source: Any, event: Any) -> None:
            self._on_agent_completed(source, event)

        @crewai_event_bus.on(AgentExecutionErrorEvent)
        def on_agent_error(source: Any, event: Any) -> None:
            self._on_agent_error(source, event)

        @crewai_event_bus.on(ToolUsageStartedEvent)
        def on_tool_started(source: Any, event: Any) -> None:
            self._on_tool_started(source, event)

        @crewai_event_bus.on(ToolUsageFinishedEvent)
        def on_tool_finished(source: Any, event: Any) -> None:
            self._on_tool_finished(source, event)

        @crewai_event_bus.on(ToolUsageErrorEvent)
        def on_tool_error(source: Any, event: Any) -> None:
            self._on_tool_error(source, event)

        @crewai_event_bus.on(LLMCallStartedEvent)
        def on_llm_started(source: Any, event: Any) -> None:
            self._on_llm_started(source, event)

        @crewai_event_bus.on(LLMCallCompletedEvent)
        def on_llm_completed(source: Any, event: Any) -> None:
            self._on_llm_completed(source, event)

        @crewai_event_bus.on(LLMCallFailedEvent)
        def on_llm_failed(source: Any, event: Any) -> None:
            self._on_llm_failed(source, event)

    # ------------------------------------------------------------------
    # Crew lifecycle
    # ------------------------------------------------------------------

    def _on_crew_started(self, source: Any, event: Any) -> None:
        try:
            name = self._crew_name or getattr(event, "crew_name", None) or "crewai"
            span = Span(
                name=name,
                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._crew_span = span
            self._trace_id = span.trace_id

            span.set_attribute(self.ATTR_GEN_AI_AGENT_NAME, name)
            span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "crewai")
            span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "invoke_agent")
            span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "crewai")
            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))

            # Capture kickoff inputs
            inputs = getattr(event, "inputs", None)
            if inputs:
                span.set_input(inputs)

            set_current_span(span)
            self._client.export_eager_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_crew_started", exc_info=True)

    def _on_crew_completed(self, source: Any, event: Any) -> None:
        try:
            if not self._crew_span:
                return
            output = getattr(event, "output", None)
            if output is not None:
                self._crew_span.set_output(str(output))

            # Aggregate token count from crew
            total_tokens = getattr(event, "total_tokens", 0) or 0
            if total_tokens:
                self._crew_span.set_attribute("gen_ai.usage.total_tokens", total_tokens)

            self._crew_span.set_status(StatusCode.OK)
            self._crew_span.end()
            self._client.export_span(self._crew_span)
            self._crew_span = None
            clear_context()
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_crew_completed", exc_info=True)

    def _on_crew_failed(self, source: Any, event: Any) -> None:
        try:
            if not self._crew_span:
                return
            error = getattr(event, "error", None)
            exc = error if isinstance(error, BaseException) else Exception(str(error) if error else "Crew failed")
            self._crew_span.set_attribute("error.type", type(exc).__qualname__)
            self._crew_span.record_exception(exc)
            self._crew_span.set_status(StatusCode.ERROR, str(error) if error else "Crew failed")
            self._crew_span.end()
            self._client.export_span(self._crew_span)
            self._crew_span = None
            clear_context()
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_crew_failed", exc_info=True)

    # ------------------------------------------------------------------
    # Task lifecycle
    # ------------------------------------------------------------------

    def _on_task_started(self, source: Any, event: Any) -> None:
        try:
            task = getattr(event, "task", None)
            if task is None:
                return

            task_key = id(task)
            task_name = (
                getattr(task, "name", None)
                or (getattr(task, "description", None) or "task")[:60]
            )
            parent_span = self._crew_span
            parent_id = parent_span.span_id if parent_span else None

            span = Span(
                name=task_name,
                trace_id=self._trace_id,
                parent_id=parent_id,
                span_type=SpanType.TASK,
                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()

            span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "crewai")
            span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "invoke_agent")
            span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "crewai")
            task_id = str(getattr(task, "id", ""))
            if task_id:
                span.set_attribute(self.ATTR_CREWAI_TASK_ID, task_id)
            description = getattr(task, "description", None)
            if description:
                span.set_attribute(self.ATTR_CREWAI_TASK_DESCRIPTION, str(description))
                span.set_input(str(description))

            # Task context (if provided)
            context = getattr(event, "context", None)
            if context:
                span.set_input(str(context))

            # Fallback: if the crew span has no input yet, bubble up
            # the first task description so the root span isn't empty.
            if (
                self._crew_span
                and not self._crew_span.attributes.get("span.input")
                and description
            ):
                self._crew_span.set_input(str(description))

            self._task_spans[task_key] = span
            self._client.export_eager_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_task_started", exc_info=True)

    def _on_task_completed(self, source: Any, event: Any) -> None:
        try:
            task = getattr(event, "task", None)
            span = self._task_spans.pop(id(task), None) if task is not None else None
            if not span:
                return
            output = getattr(event, "output", None)
            if output is not None:
                span.set_output(str(output))
            span.set_status(StatusCode.OK)
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_task_completed", exc_info=True)

    def _on_task_failed(self, source: Any, event: Any) -> None:
        try:
            task = getattr(event, "task", None)
            span = self._task_spans.pop(id(task), None) if task is not None else None
            if not span:
                return
            error = getattr(event, "error", None)
            exc = error if isinstance(error, BaseException) else Exception(str(error) if error else "Task failed")
            span.set_attribute("error.type", type(exc).__qualname__)
            span.record_exception(exc)
            span.set_status(StatusCode.ERROR, str(error) if error else "Task failed")
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_task_failed", exc_info=True)

    # ------------------------------------------------------------------
    # Agent execution lifecycle
    # ------------------------------------------------------------------

    def _on_agent_started(self, source: Any, event: Any) -> None:
        try:
            agent = getattr(event, "agent", None)
            if agent is None:
                return

            agent_id = str(getattr(agent, "id", "") or "")
            role = getattr(agent, "role", None) or "agent"

            # Parent: task_span if we can find it, otherwise crew_span
            task = getattr(event, "task", None)
            parent_span = (
                self._task_spans.get(id(task)) if task is not None else None
            ) or self._crew_span
            parent_id = parent_span.span_id if parent_span else None

            span = Span(
                name=role,
                trace_id=self._trace_id,
                parent_id=parent_id,
                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()
            span.set_attribute(self.ATTR_GEN_AI_AGENT_NAME, role)
            span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "crewai")
            span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "invoke_agent")
            span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "crewai")

            goal = getattr(agent, "goal", None)
            if goal:
                span.set_attribute(self.ATTR_GEN_AI_AGENT_DESCRIPTION, str(goal))
            backstory = getattr(agent, "backstory", None)
            if backstory:
                span.set_attribute(self.ATTR_CREWAI_AGENT_BACKSTORY, str(backstory))

            # Capture the task prompt as the agent's input
            task_prompt = getattr(event, "task_prompt", None)
            if task_prompt:
                span.set_input(str(task_prompt))

            if agent_id:
                self._agent_spans[agent_id] = span
                self._current_agent_id = agent_id
            self._client.export_eager_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_agent_started", exc_info=True)

    def _on_agent_completed(self, source: Any, event: Any) -> None:
        try:
            agent = getattr(event, "agent", None)
            agent_id = str(getattr(agent, "id", "") or "") if agent else ""
            span = self._agent_spans.pop(agent_id, None) if agent_id else None
            self._current_agent_id = ""
            if not span:
                return
            output = getattr(event, "output", None)
            if output is not None:
                span.set_output(str(output))
            span.set_status(StatusCode.OK)
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_agent_completed", exc_info=True)

    def _on_agent_error(self, source: Any, event: Any) -> None:
        try:
            agent = getattr(event, "agent", None)
            agent_id = str(getattr(agent, "id", "") or "") if agent else ""
            span = self._agent_spans.pop(agent_id, None) if agent_id else None
            self._current_agent_id = ""
            if not span:
                return
            error = getattr(event, "error", None)
            exc = error if isinstance(error, BaseException) else Exception(str(error) if error else "Agent error")
            span.set_attribute("error.type", type(exc).__qualname__)
            span.record_exception(exc)
            span.set_status(StatusCode.ERROR, str(error) if error else "Agent error")
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_agent_error", exc_info=True)

    # ------------------------------------------------------------------
    # LLM call lifecycle
    # ------------------------------------------------------------------

    def _on_llm_started(self, source: Any, event: Any) -> None:
        try:
            # Find parent agent span (tries multiple lookup strategies)
            agent_id, parent_span = self._resolve_agent_span(event)
            parent_id = parent_span.span_id if parent_span else None

            model = getattr(event, "model", None) or ""
            span = Span(
                name=f"llm-call{': ' + model if model else ''}",
                trace_id=self._trace_id,
                parent_id=parent_id,
                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_AGENT_FRAMEWORK, "crewai")
            span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "chat")
            span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "crewai")

            if model:
                span.set_attribute(self.ATTR_GEN_AI_REQUEST_MODEL, model)
                span.set_attribute(self.ATTR_GEN_AI_RESPONSE_MODEL, model)

            messages = getattr(event, "messages", None)
            if messages is not None:
                span.set_input(messages)
                try:
                    span.set_attribute(
                        self.ATTR_GEN_AI_INPUT_MESSAGES,
                        json.dumps(messages, default=str),
                    )
                except Exception:
                    logger.debug("Could not serialize input messages", exc_info=True)

            # Key by resolved agent_id — CrewAI LLM events don't carry
            # a call_id, but execution is sequential per agent so agent_id
            # uniquely identifies the in-flight LLM call.
            if agent_id:
                self._llm_spans[agent_id] = span
            self._client.export_eager_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_llm_started", exc_info=True)

    def _on_llm_completed(self, source: Any, event: Any) -> None:
        try:
            agent_id, _ = self._resolve_agent_span(event)
            span = self._llm_spans.pop(agent_id, None) if agent_id else None
            if not span:
                return

            # Extract response text
            response = getattr(event, "response", None)
            if response is not None:
                output_text = self._extract_llm_output(response)
                if output_text:
                    span.set_output(output_text)
                    span.set_attribute(
                        self.ATTR_GEN_AI_OUTPUT_MESSAGES,
                        json.dumps([{"role": "assistant", "content": output_text}], default=str),
                    )

                # Extract token usage from the raw response object
                self._apply_token_usage(span, response)

            # Set response model (from event, same as request for CrewAI)
            model = getattr(event, "model", None)
            if model:
                span.set_attribute(self.ATTR_GEN_AI_RESPONSE_MODEL, str(model))

            span.set_status(StatusCode.OK)
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_llm_completed", exc_info=True)

    def _on_llm_failed(self, source: Any, event: Any) -> None:
        try:
            agent_id, _ = self._resolve_agent_span(event)
            span = self._llm_spans.pop(agent_id, None) if agent_id else None
            if not span:
                return
            error = getattr(event, "error", None)
            exc = error if isinstance(error, BaseException) else Exception(str(error) if error else "LLM call failed")
            span.set_attribute("error.type", type(exc).__qualname__)
            span.record_exception(exc)
            span.set_status(StatusCode.ERROR, str(error) if error else "LLM call failed")
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_llm_failed", exc_info=True)

    # ------------------------------------------------------------------
    # Tool usage lifecycle
    # ------------------------------------------------------------------

    def _on_tool_started(self, source: Any, event: Any) -> None:
        try:
            tool_name = getattr(event, "tool_name", None) or "tool"
            agent_id, parent_span = self._resolve_agent_span(event)
            tool_key = (agent_id, tool_name)
            parent_id = parent_span.span_id if parent_span else None

            span = Span(
                name=tool_name,
                trace_id=self._trace_id,
                parent_id=parent_id,
                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,
            )
            span.start()
            span.set_attribute(self.ATTR_GEN_AI_TOOL_NAME, tool_name)
            span.set_attribute(self.ATTR_GEN_AI_TOOL_TYPE, "function")
            span.set_attribute(self.ATTR_GEN_AI_AGENT_FRAMEWORK, "crewai")
            span.set_attribute(self.ATTR_GEN_AI_OPERATION_NAME, "execute_tool")
            span.set_attribute(self.ATTR_GEN_AI_PROVIDER_NAME, "crewai")

            tool_input = getattr(event, "tool_args", None)
            if tool_input is not None:
                span.set_input(tool_input)
                try:
                    args_json = json.dumps(tool_input, default=str)
                    span.set_attribute(self.ATTR_GEN_AI_TOOL_CALL_ARGUMENTS, args_json)
                except Exception:
                    logger.debug("Could not serialize tool arguments", exc_info=True)

            self._tool_spans[tool_key] = span
            self._client.export_eager_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_tool_started", exc_info=True)

    def _on_tool_finished(self, source: Any, event: Any) -> None:
        try:
            tool_name = getattr(event, "tool_name", None) or "tool"
            agent_id, _ = self._resolve_agent_span(event)
            span = self._tool_spans.pop((agent_id, tool_name), None)
            if not span:
                return
            output = getattr(event, "output", None)
            if output is not None:
                output_str = str(output)
                span.set_output(output_str)
                span.set_attribute(self.ATTR_GEN_AI_TOOL_CALL_RESULT, output_str)
            span.set_status(StatusCode.OK)
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_tool_finished", exc_info=True)

    def _on_tool_error(self, source: Any, event: Any) -> None:
        try:
            tool_name = getattr(event, "tool_name", None) or "tool"
            agent_id, _ = self._resolve_agent_span(event)
            span = self._tool_spans.pop((agent_id, tool_name), None)
            if not span:
                return
            error = getattr(event, "error", None)
            exc = error if isinstance(error, BaseException) else Exception(str(error) if error else "Tool error")
            span.set_attribute("error.type", type(exc).__qualname__)
            span.record_exception(exc)
            span.set_status(StatusCode.ERROR, str(error) if error else "Tool error")
            span.end()
            self._client.export_span(span)
        except Exception:
            logger.warning("BeaconCrewAIListener: error in _on_tool_error", exc_info=True)

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

    def _resolve_agent_span(self, event: Any) -> tuple[str, Span | None]:
        """Resolve the agent span for an LLM or tool event.

        Tries multiple lookup strategies in order:
        1. ``event.agent_id`` — direct field on the event
        2. ``event.agent.id`` — agent object attached to the event
        3. ``self._current_agent_id`` — most recently started agent

        Returns ``(agent_id, parent_span)`` where *parent_span* falls
        back to the crew span when no agent span can be found.
        """
        # Strategy 1: direct agent_id field
        agent_id = getattr(event, "agent_id", None) or ""
        if agent_id and agent_id in self._agent_spans:
            return agent_id, self._agent_spans[agent_id]

        # Strategy 2: agent object on the event
        agent_obj = getattr(event, "agent", None)
        if agent_obj is not None:
            obj_id = str(getattr(agent_obj, "id", "") or "")
            if obj_id and obj_id in self._agent_spans:
                return obj_id, self._agent_spans[obj_id]

        # Strategy 3: fallback to most recently started agent
        if self._current_agent_id and self._current_agent_id in self._agent_spans:
            return self._current_agent_id, self._agent_spans[self._current_agent_id]

        # Last resort: crew span
        return agent_id or self._current_agent_id, self._crew_span

    def _extract_llm_output(self, response: Any) -> str | None:
        """Extract the assistant's text content from a LiteLLM/OpenAI response."""
        try:
            # LiteLLM ModelResponse / OpenAI ChatCompletion
            choices = getattr(response, "choices", None)
            if choices and len(choices) > 0:
                message = getattr(choices[0], "message", None)
                if message:
                    content = getattr(message, "content", None)
                    if content:
                        return str(content)
        except Exception:
            logger.debug("Could not extract structured LLM output", exc_info=True)
        # Fallback to string representation
        if response is not None:
            return str(response)
        return None

    def _apply_token_usage(self, span: Span, response: Any) -> None:
        """Extract token usage from a LiteLLM/OpenAI response and set span attributes."""
        try:
            usage = getattr(response, "usage", None)
            if usage is None:
                return

            # LiteLLM / OpenAI style
            input_tokens = (
                getattr(usage, "prompt_tokens", None)
                or getattr(usage, "input_tokens", None)
                or 0
            )
            output_tokens = (
                getattr(usage, "completion_tokens", None)
                or getattr(usage, "output_tokens", None)
                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)
        except Exception:
            logger.debug("Could not extract token usage from LLM response", exc_info=True)
