"""
Real SDK integration tests for Pydantic AI + Cascade tracing.

Runs actual Pydantic AI agents with real OpenAI API calls and verifies
span attributes using an in-memory exporter.

Requires: OPENAI_API_KEY in .env (loaded automatically via dotenv).
"""

import os

import pytest
from dotenv import load_dotenv
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

import cascade.integrations.pydantic_ai as _pai_mod
import cascade.tracing as _tracing_mod

load_dotenv()

pytestmark = pytest.mark.skipif(
    not os.getenv("OPENAI_API_KEY"),
    reason="OPENAI_API_KEY required for real Pydantic AI integration tests",
)


# ── helpers ────────────────────────────────────────────────────────────────

def _unwrap(fn):
    """Walk the functools.wraps __wrapped__ chain back to the original."""
    while hasattr(fn, "__wrapped__"):
        fn = fn.__wrapped__
    return fn


def _spans_by_type(spans):
    """Group finished spans by their cascade.span_type attribute."""
    groups: dict[str, list] = {}
    for s in spans:
        st = dict(s.attributes).get("cascade.span_type")
        groups.setdefault(st, []).append(s)
    return groups


# ── module-scoped OTEL provider (set_tracer_provider can only be called once)

_module_exporter = InMemorySpanExporter()
_module_provider = TracerProvider()
_module_provider.add_span_processor(SimpleSpanProcessor(_module_exporter))
trace.set_tracer_provider(_module_provider)


# ── per-test fixture ──────────────────────────────────────────────────────

@pytest.fixture()
def span_capture():
    """
    Per-test setup: clear the in-memory exporter, restore Agent methods
    to their truly-original (unpatched) state, and wire the Cascade tracer
    to the shared TracerProvider.
    """
    from pydantic_ai import Agent

    Agent.run = _unwrap(Agent.run)
    Agent.run_sync = _unwrap(Agent.run_sync)
    Agent.run_stream = _unwrap(Agent.run_stream)
    if hasattr(Agent, "run_stream_sync"):
        Agent.run_stream_sync = _unwrap(Agent.run_stream_sync)
    _pai_mod._patched = False

    _module_exporter.clear()

    _tracing_mod._tracer = trace.get_tracer("cascade.tracing")
    _tracing_mod._tracer_provider = _module_provider
    _tracing_mod._cascade_endpoint = None
    _tracing_mod._cascade_api_key = None
    _tracing_mod._cascade_project = None

    yield _module_exporter

    _module_exporter.clear()
    _tracing_mod._tracer = None
    _tracing_mod._tracer_provider = None


# ── 1. setup_pydantic_ai() global patching ─────────────────────────────────

class TestSetupPydanticAiGlobal:
    """Verify that setup_pydantic_ai() class-level patching creates spans."""

    def test_creates_agent_and_llm_spans(self, span_capture):
        from pydantic_ai import Agent

        _pai_mod._patched = False
        _pai_mod.setup_pydantic_ai()

        agent = Agent("openai:gpt-4o-mini", system_prompt="You are helpful.")

        @agent.tool_plain
        def greet(name: str) -> str:
            """Return a greeting."""
            return f"Hello, {name}!"

        agent.run_sync("Greet Alice using the greet tool.")

        spans = span_capture.get_finished_spans()
        by_type = _spans_by_type(spans)

        assert "agent" in by_type, f"No agent span. Spans: {[s.name for s in spans]}"
        assert "llm" in by_type, f"No llm span. Spans: {[s.name for s in spans]}"

    def test_multiple_tool_calls_produce_separate_spans(self, span_capture):
        from pydantic_ai import Agent

        _pai_mod._patched = False
        _pai_mod.setup_pydantic_ai()

        agent = Agent(
            "openai:gpt-4o-mini",
            system_prompt="You MUST call ALL available tools before answering.",
        )

        @agent.tool_plain
        def add(a: int, b: int) -> str:
            """Add two numbers."""
            return str(a + b)

        @agent.tool_plain
        def multiply(a: int, b: int) -> str:
            """Multiply two numbers."""
            return str(a * b)

        agent.run_sync("Add 10 and 20, then multiply 5 and 6. Use BOTH tools.")

        spans = span_capture.get_finished_spans()
        tool_spans = [
            s for s in spans
            if dict(s.attributes).get("cascade.span_type") == "tool"
        ]
        tool_names = {dict(s.attributes).get("tool.name") for s in tool_spans}

        assert len(tool_spans) >= 2, (
            f"Expected >=2 tool spans, got {len(tool_spans)}: {tool_names}"
        )
        assert "add" in tool_names
        assert "multiply" in tool_names


# ── 2. instrument_agent() per-instance patching ───────────────────────────

class TestInstrumentAgent:
    """Verify that instrument_agent() instruments a single instance."""

    def test_per_instance_produces_spans(self, span_capture):
        from pydantic_ai import Agent

        agent = Agent("openai:gpt-4o-mini", system_prompt="You are helpful.")

        @agent.tool_plain
        def echo(text: str) -> str:
            """Echo text back."""
            return text

        _pai_mod.instrument_agent(agent)
        agent.run_sync("Echo 'hello' using the echo tool.")

        spans = span_capture.get_finished_spans()
        by_type = _spans_by_type(spans)

        assert "agent" in by_type, f"No agent span. Spans: {[s.name for s in spans]}"
        assert "llm" in by_type, f"No llm span. Spans: {[s.name for s in spans]}"

    def test_uninstrumented_agent_produces_no_spans(self, span_capture):
        from pydantic_ai import Agent

        agent = Agent("openai:gpt-4o-mini", system_prompt="Be concise.")
        agent.run_sync("Say hi.")

        spans = span_capture.get_finished_spans()
        cascade_spans = [
            s for s in spans
            if dict(s.attributes).get("cascade.span_type") is not None
        ]
        assert len(cascade_spans) == 0, (
            f"Uninstrumented agent should produce 0 Cascade spans, got {len(cascade_spans)}"
        )


# ── 3. Span attributes ────────────────────────────────────────────────────

class TestSpanAttributes:
    """Verify individual span attributes match the integration spec."""

    def test_agent_span_has_required_attrs(self, span_capture):
        from pydantic_ai import Agent

        _pai_mod._patched = False
        _pai_mod.setup_pydantic_ai()

        agent = Agent("openai:gpt-4o-mini", system_prompt="Calculator")

        @agent.tool_plain
        def divide(a: int, b: int) -> str:
            """Divide a by b."""
            return str(a / b) if b else "undefined"

        agent.run_sync("Divide 100 by 4. Use the divide tool.")

        agent_spans = [
            s for s in span_capture.get_finished_spans()
            if dict(s.attributes).get("cascade.span_type") == "agent"
        ]
        assert agent_spans, "No agent span found"
        attrs = dict(agent_spans[0].attributes)

        assert attrs["cascade.span_type"] == "agent"
        assert "llm.model" in attrs
        assert "llm.prompt" in attrs

    def test_llm_span_has_token_counts(self, span_capture):
        from pydantic_ai import Agent

        _pai_mod._patched = False
        _pai_mod.setup_pydantic_ai()

        agent = Agent("openai:gpt-4o-mini", system_prompt="Be concise.")
        agent.run_sync("Say hello.")

        llm_spans = [
            s for s in span_capture.get_finished_spans()
            if dict(s.attributes).get("cascade.span_type") == "llm"
        ]
        assert llm_spans, "No LLM span found"
        attrs = dict(llm_spans[0].attributes)

        assert "llm.input_tokens" in attrs
        assert "llm.output_tokens" in attrs
        assert attrs["llm.input_tokens"] > 0
        assert attrs["llm.output_tokens"] > 0

    def test_tool_span_has_name_input_output(self, span_capture):
        from pydantic_ai import Agent

        _pai_mod._patched = False
        _pai_mod.setup_pydantic_ai()

        agent = Agent(
            "openai:gpt-4o-mini",
            system_prompt="You MUST call the square tool.",
        )

        @agent.tool_plain
        def square(n: int) -> str:
            """Return n squared."""
            return str(n * n)

        agent.run_sync("Square the number 7. Use the square tool.")

        tool_spans = [
            s for s in span_capture.get_finished_spans()
            if dict(s.attributes).get("cascade.span_type") == "tool"
        ]
        assert tool_spans, "No tool span found"
        attrs = dict(tool_spans[0].attributes)

        assert attrs["tool.name"] == "square"
        assert "tool.input" in attrs
        assert "tool.output" in attrs


# ── 4. Error handling ──────────────────────────────────────────────────────

class TestErrorHandling:
    """Verify the integration handles errors gracefully."""

    def test_error_does_not_crash_integration(self, span_capture):
        """Even if a tool raises, the integration still produces spans."""
        from pydantic_ai import Agent

        _pai_mod._patched = False
        _pai_mod.setup_pydantic_ai()

        agent = Agent(
            "openai:gpt-4o-mini",
            system_prompt="You MUST call the boom tool immediately.",
        )

        @agent.tool_plain
        def boom() -> str:
            """This tool always fails."""
            raise RuntimeError("Simulated failure")

        try:
            agent.run_sync("Call the boom tool now.")
        except Exception:
            pass

        spans = span_capture.get_finished_spans()
        assert len(spans) >= 1, "Expected at least one span even on error"

        agent_spans = [
            s for s in spans
            if dict(s.attributes).get("cascade.span_type") == "agent"
        ]
        assert agent_spans, f"No agent span found. Spans: {[s.name for s in spans]}"
