from types import SimpleNamespace
from typing import Any, Generator, Optional, cast

import pytest
import vcr  # type: ignore
from agno.agent import Agent
from agno.models.base import Model
from agno.models.openai.chat import OpenAIChat
from agno.run.agent import RunOutput
from agno.team import Team
from agno.tools.duckduckgo import DuckDuckGoTools
from agno.tools.yfinance import YFinanceTools
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.util._importlib_metadata import entry_points

from openinference.instrumentation import OITracer
from openinference.instrumentation.agno import AgnoInstrumentor
from openinference.semconv.trace import SpanAttributes

test_vcr = vcr.VCR(
    serializer="yaml",
    cassette_library_dir="tests/openinference/instrumentation/agno/fixtures/",
    record_mode="none",
    match_on=["uri", "method"],
)


@pytest.fixture()
def in_memory_span_exporter() -> InMemorySpanExporter:
    return InMemorySpanExporter()


@pytest.fixture()
def tracer_provider(in_memory_span_exporter: InMemorySpanExporter) -> TracerProvider:
    resource = Resource(attributes={})
    tracer_provider = TracerProvider(resource=resource)
    tracer_provider.add_span_processor(SimpleSpanProcessor(in_memory_span_exporter))
    return tracer_provider


@pytest.fixture()
def setup_agno_instrumentation(
    tracer_provider: TracerProvider,
) -> Generator[None, None, None]:
    AgnoInstrumentor().instrument(tracer_provider=tracer_provider)
    yield
    AgnoInstrumentor().uninstrument()


class TestInstrumentor:
    def test_entrypoint_for_opentelemetry_instrument(self) -> None:
        (instrumentor_entrypoint,) = entry_points(
            group="opentelemetry_instrumentor",
            name="agno",
        )
        instrumentor = instrumentor_entrypoint.load()()
        assert isinstance(instrumentor, AgnoInstrumentor)

    # Ensure we're using the common OITracer from common openinference-instrumentation pkg
    def test_oitracer(self, setup_agno_instrumentation: Any) -> None:
        assert isinstance(AgnoInstrumentor()._tracer, OITracer)


def test_agno_instrumentation(
    tracer_provider: TracerProvider,
    in_memory_span_exporter: InMemorySpanExporter,
    setup_agno_instrumentation: Any,
) -> None:
    with test_vcr.use_cassette("agent_run.yaml", filter_headers=["authorization", "X-API-KEY"]):
        import os

        os.environ["OPENAI_API_KEY"] = "fake_key"
        agent = Agent(
            name="News Agent",  # For best results, set a name that will be used by the tracer
            model=OpenAIChat(id="gpt-4o-mini"),
            tools=[DuckDuckGoTools()],
            user_id="test_user_123",
        )
        agent.run("What's trending on Twitter?", session_id="test_session")
    spans = in_memory_span_exporter.get_finished_spans()
    assert len(spans) == 4
    checked_spans = 0
    for span in spans:
        attributes = dict(span.attributes or dict())
        if span.name == "News_Agent.run":
            checked_spans += 1
            assert attributes.get("openinference.span.kind") == "AGENT"
            assert attributes.get("output.value")
            assert attributes.get("session.id") == "test_session"
            # assert that there are no tokens on the kickoff chain so that we do not
            # double count token when a user is also instrumenting with another instrumentor
            # that provides token counts via the spans.
            assert attributes.get("llm.token_count.prompt") is None
            assert attributes.get("llm.token_count.completion") is None
            assert attributes.get("llm.token_count.total") is None
            # Validate agent-specific attributes
            assert attributes.get("agno.agent.id") is not None, "Agent ID should be present"
            assert attributes.get("agno.run.id") is not None, "Run ID should be present"
            assert attributes.get("agent.name") == "News Agent"
            assert attributes.get("user.id") == "test_user_123"
            assert span.status.is_ok
        elif span.name == "ToolUsage._use":
            checked_spans += 1
            assert attributes.get("openinference.span.kind") == "TOOL"
            assert attributes.get("output.value")
            assert attributes.get("tool.name") in (
                "Search the internet with Serper",
                "Ask question to coworker",
                "Delegate work to coworker",
            )
            assert span.status.is_ok
        elif span.name == "Task._execute_core":
            checked_spans += 1
            assert attributes.get("output.value")
            assert attributes["openinference.span.kind"] == "AGENT"
            assert attributes.get("input.value")
            assert span.status.is_ok
        elif span.name == "OpenAIChat.invoke":
            checked_spans += 1
            assert attributes.get("openinference.span.kind") == "LLM"
            assert attributes.get("llm.model_name") == "gpt-4o-mini"
            assert attributes.get("llm.provider") == "OpenAI"
            assert attributes.get("llm.system") == "openai"
            assert span.status.is_ok
    assert checked_spans >= 3  # We expect at least agent, tool, and LLM spans


def test_agent_metadata_captured() -> None:
    """Test that Agent.metadata dict is captured as a span attribute."""
    import json

    from openinference.instrumentation.agno._runs_wrapper import _agent_run_attributes

    agent = Agent(
        name="Test Agent",
        metadata={
            "department": "finance",
            "cost_center": "dept_123",
            "environment": "production",
        },
    )
    attributes = dict(_agent_run_attributes(agent))

    raw_metadata = attributes.get("metadata")
    assert raw_metadata is not None, "metadata attribute should be present"
    assert isinstance(raw_metadata, str), "metadata should be a JSON string"

    metadata = json.loads(raw_metadata)
    assert metadata["department"] == "finance"
    assert metadata["cost_center"] == "dept_123"
    assert metadata["environment"] == "production"


def test_agent_no_metadata() -> None:
    """Test that no metadata attribute is set when agent has no metadata."""
    from openinference.instrumentation.agno._runs_wrapper import _agent_run_attributes

    agent = Agent(name="Test Agent")
    attributes = dict(_agent_run_attributes(agent))
    assert "metadata" not in attributes


def test_team_metadata_captured() -> None:
    """Test that Team.metadata dict is captured as a span attribute."""
    import json

    from openinference.instrumentation.agno._runs_wrapper import _agent_run_attributes

    team = Team(
        name="Test Team",
        members=[Agent(name="Member Agent")],
        metadata={
            "project": "alpha",
            "priority": "high",
        },
    )
    attributes = dict(_agent_run_attributes(team))

    raw_metadata = attributes.get("metadata")
    assert raw_metadata is not None, "metadata attribute should be present"
    assert isinstance(raw_metadata, str), "metadata should be a JSON string"

    metadata = json.loads(raw_metadata)
    assert metadata["project"] == "alpha"
    assert metadata["priority"] == "high"


@pytest.mark.parametrize(
    "model_name, expected_system",
    [
        ("gpt-4o-mini", "openai"),
        ("claude-sonnet-4-6", "anthropic"),
        ("command-r", "cohere"),
        ("mistral-large-latest", "mistralai"),
        # Non-inferable model names stay represented by llm.provider only.
        ("llama-3.1-70b-versatile", None),
        ("custom-model", None),
        # Empty / missing model ids yield no llm.system
        (None, None),
        ("", None),
        ("   ", None),
    ],
)
def test_get_llm_system(model_name: Any, expected_system: Any) -> None:
    """llm.system is derived from the model id, not only the provider."""
    from openinference.instrumentation.agno._model_wrapper import _get_llm_system

    model = cast(Model, SimpleNamespace(id=model_name))
    assert _get_llm_system(model) == expected_system


@pytest.mark.parametrize(
    "vertexai, vertexai_env, client_params, expected_system",
    [
        (False, None, None, "google"),
        (True, None, None, "vertexai"),
        (False, "TRUE", None, "vertexai"),
        (False, "1", None, "vertexai"),
        (False, None, {"vertexai": True}, "vertexai"),
        (True, None, {"vertexai": False}, "google"),
        (True, None, {"vertexai": None}, "google"),
        (True, "1", {"vertexai": None}, "vertexai"),
    ],
)
def test_get_llm_system_for_gemini_api_mode(
    vertexai: bool,
    vertexai_env: Optional[str],
    client_params: Optional[dict[str, Any]],
    expected_system: str,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from openinference.instrumentation.agno._model_wrapper import _get_llm_system

    if vertexai_env is None:
        monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
    else:
        monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", vertexai_env)

    model = cast(
        Model,
        SimpleNamespace(
            id="gemini-2.0-flash",
            name="Gemini",
            provider="Google",
            vertexai=vertexai,
            client=None,
            client_params=client_params,
        ),
    )
    assert _get_llm_system(model) == expected_system


@pytest.mark.parametrize(
    "client_vertexai, expected_system",
    [
        (False, "google"),
        (True, "vertexai"),
        (None, None),
    ],
)
def test_get_llm_system_for_prebuilt_gemini_client(
    client_vertexai: Optional[bool],
    expected_system: Optional[str],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    from openinference.instrumentation.agno._model_wrapper import _get_llm_system

    monkeypatch.setenv("GOOGLE_GENAI_USE_VERTEXAI", "true")
    model = cast(
        Model,
        SimpleNamespace(
            id="gemini-2.0-flash",
            name="Gemini",
            provider="Google",
            vertexai=True,
            client=SimpleNamespace(vertexai=client_vertexai),
            client_params={"vertexai": True},
        ),
    )
    assert _get_llm_system(model) == expected_system


def test_get_llm_system_for_google_gemini_without_vertexai_setting(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Google's Gemini Interactions API is Developer API-only."""
    from openinference.instrumentation.agno._model_wrapper import _get_llm_system

    monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False)
    model = cast(
        Model,
        SimpleNamespace(
            id="gemini-3-flash-preview",
            name="GeminiInteractions",
            provider="Google",
            client=None,
            client_params=None,
        ),
    )
    assert _get_llm_system(model) == "google"


def test_agno_team_coordinate_instrumentation(
    tracer_provider: TracerProvider,
    in_memory_span_exporter: InMemorySpanExporter,
    setup_agno_instrumentation: Any,
) -> None:
    with test_vcr.use_cassette(
        "team_coordinate_run.yaml", filter_headers=["authorization", "X-API-KEY"]
    ):
        import os
        import re

        os.environ["OPENAI_API_KEY"] = "fake_key"

        web_agent = Agent(
            name="Web Agent",
            role="Search the web for information",
            model=OpenAIChat(id="gpt-4o-mini"),
            tools=[DuckDuckGoTools()],
            instructions="Always include sources",
        )

        finance_agent = Agent(
            name="Finance Agent",
            role="Get financial data",
            model=OpenAIChat(id="gpt-4o-mini"),
            tools=[YFinanceTools()],
            instructions="Use tables to display data",
        )

        agent_team = Team(
            name="Team",
            members=[web_agent, finance_agent],
            model=OpenAIChat(id="gpt-4o-mini"),
            instructions=["Always include sources", "Use tables to display data"],
            user_id="team_user_999",
        )

        agent_team.run(
            "What's the market outlook and financial performance of NVIDIA?",
            session_id="test_session",
        )

    spans = in_memory_span_exporter.get_finished_spans()
    assert len(spans) >= 2

    # Collect spans by name for graph relationship validation
    team_span = None
    web_agent_span = None
    finance_agent_span = None

    for span in spans:
        attributes = dict(span.attributes or dict())
        if span.name == "Team.run":
            team_span = attributes
        elif span.name == "Web_Agent.run":
            web_agent_span = attributes
        elif span.name == "Finance_Agent.run":
            finance_agent_span = attributes

    # Helper function to validate node ID format (16-character hex string)
    def is_valid_node_id(node_id: str) -> bool:
        return bool(re.match(r"^[0-9a-f]{16}$", node_id))

    # Validate graph attributes for team span
    assert team_span is not None, "Team span should be found"
    team_node_id = team_span.get(SpanAttributes.GRAPH_NODE_ID)
    assert team_node_id is not None, "Team node ID should be present"
    assert isinstance(team_node_id, str), f"Team node ID should be a string: {team_node_id}"
    assert is_valid_node_id(team_node_id), f"Team node ID should be valid hex: {team_node_id}"
    # Team should have no parent (root node)
    assert team_span.get(SpanAttributes.GRAPH_NODE_PARENT_ID) is None
    # Validate team-specific attributes
    assert team_span.get("agno.team.id") is not None, "Team ID should be present"
    assert team_span.get("agno.run.id") is not None, "Team run ID should be present"
    assert team_span.get(SpanAttributes.AGENT_NAME) == "Team"
    assert team_span.get("user.id") == "team_user_999"

    # Validate graph attributes for web agent span
    if web_agent_span is not None:
        web_agent_node_id = web_agent_span.get(SpanAttributes.GRAPH_NODE_ID)
        assert web_agent_node_id is not None, "Web agent node ID should be present"
        assert isinstance(web_agent_node_id, str), (
            f"Web agent node ID should be a string: {web_agent_node_id}"
        )
        assert is_valid_node_id(web_agent_node_id), (
            f"Web agent node ID should be valid hex: {web_agent_node_id}"
        )
        assert web_agent_span.get(SpanAttributes.GRAPH_NODE_NAME) == "Web Agent"
        assert web_agent_span.get(SpanAttributes.AGENT_NAME) == "Web Agent"
        # Web agent should have team as parent
        assert web_agent_span.get(SpanAttributes.GRAPH_NODE_PARENT_ID) == team_node_id
        # Ensure web agent has different node ID than team (uniqueness)
        assert web_agent_node_id != team_node_id, "Web agent should have unique node ID"

    # Validate graph attributes for finance agent span
    if finance_agent_span is not None:
        finance_agent_node_id = finance_agent_span.get(SpanAttributes.GRAPH_NODE_ID)
        assert finance_agent_node_id is not None, "Finance agent node ID should be present"
        assert isinstance(finance_agent_node_id, str), (
            f"Finance agent node ID should be a string: {finance_agent_node_id}"
        )
        assert is_valid_node_id(finance_agent_node_id), (
            f"Finance agent node ID should be valid hex: {finance_agent_node_id}"
        )
        assert finance_agent_span.get(SpanAttributes.GRAPH_NODE_NAME) == "Finance Agent"
        assert finance_agent_span.get(SpanAttributes.AGENT_NAME) == "Finance Agent"
        # Finance agent should have team as parent
        assert finance_agent_span.get(SpanAttributes.GRAPH_NODE_PARENT_ID) == team_node_id
        # Ensure finance agent has different node ID than team (uniqueness)
        assert finance_agent_node_id != team_node_id, "Finance agent should have unique node ID"

    # If both agents are present, ensure they have different node IDs
    if web_agent_span is not None and finance_agent_span is not None:
        assert web_agent_node_id != finance_agent_node_id, "Agents should have unique node IDs"

    # At least one agent span should be present to validate parent-child relationship
    assert web_agent_span is not None or finance_agent_span is not None, (
        "At least one agent span should be found"
    )


def test_extract_run_response_output_str_content() -> None:
    """String content is returned verbatim."""
    from openinference.instrumentation.agno._runs_wrapper import _extract_run_response_output

    run_response = RunOutput(content="hello world")
    assert _extract_run_response_output(run_response) == "hello world"


def test_extract_run_response_output_pydantic_content() -> None:
    """Content exposing model_dump_json is serialized via that method."""
    from types import SimpleNamespace

    from openinference.instrumentation.agno._runs_wrapper import _extract_run_response_output

    content = SimpleNamespace(model_dump_json=lambda: '{"answer": 42}')
    run_response = RunOutput(content=content)
    assert _extract_run_response_output(run_response) == '{"answer": 42}'


def test_extract_run_response_output_dict_content() -> None:
    """Dict content is serialized as valid JSON.

    Standalone ``agent.run`` calls that use ``output_schema``/JSON mode can return
    a plain ``dict`` as ``content``. Previously this raised
    ``'dict' object has no attribute 'model_dump_json'``.
    """
    from openinference.instrumentation.agno._runs_wrapper import _extract_run_response_output

    run_response = RunOutput(content={"answer": True})
    assert _extract_run_response_output(run_response) == '{"answer": true}'
