# Copyright 2025 Tenro.ai
# SPDX-License-Identifier: Apache-2.0

"""Utilities for filtering span objects by various criteria.

This module provides generic filtering functions that work with all span types
(ToolCall, LLMCall, AgentRun) since they all inherit from BaseSpan.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, TypeVar

if TYPE_CHECKING:
    from tenro.core.spans import AgentRun, BaseSpan

T = TypeVar("T", bound="BaseSpan")


def filter_by_agent(agent: str | None, spans: list[T], agent_runs: list[AgentRun]) -> list[T]:
    """Filter spans by agent name using agent_id relationship.

    Works with any span type (ToolCall, LLMCall, AgentRun):
    - ToolCall/LLMCall: Filters by agent_id (the agent that made the call)
    - AgentRun: Filters by parent_agent_id (child agents spawned by parent)

    Args:
        agent: Agent name to filter by. If `None`, returns all spans.
        spans: List of spans to filter (generic type T).
        agent_runs: List of agent runs to search for matching agent.

    Returns:
        Filtered list of spans created by the specified agent.
        Returns empty list if agent not found.
        Returns all spans if agent is `None`.

    Examples:
        >>> # Filter tool calls by agent
        >>> manager_tools = filter_by_agent("Manager", tool_calls, agent_runs)
        >>> # Filter LLM calls by agent
        >>> researcher_llms = filter_by_agent("Researcher", llm_calls, agent_runs)
        >>> # Filter child agents by parent
        >>> manager_children = filter_by_agent("Manager", agent_runs, agent_runs)
        >>> # No filter - returns all
        >>> all_spans = filter_by_agent(None, spans, agent_runs)
    """
    if not agent:
        return spans

    from tenro.core.spans import AgentRun as AgentRunClass

    matching_agents = [a for a in agent_runs if a.name == agent]
    if not matching_agents:
        return []
    matching_ids = {a.id for a in matching_agents}

    result: list[T] = []
    for span in spans:
        if isinstance(span, AgentRunClass):
            if span.parent_agent_id in matching_ids:
                result.append(span)
        else:
            if span.agent_id in matching_ids:
                result.append(span)
    return result


__all__ = ["filter_by_agent"]
