"""
Real SDK integration tests for Claude Agent SDK + Cascade tracing.

Runs actual Claude Agent SDK queries with real Anthropic API calls and
verifies span attributes using an in-memory exporter.

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

import asyncio
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.tracing as _tracing_mod

load_dotenv()

pytestmark = pytest.mark.skipif(
    not os.getenv("ANTHROPIC_API_KEY"),
    reason="ANTHROPIC_API_KEY required for real Claude Agent SDK integration tests",
)


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

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


async def _consume(traced_query, **kwargs):
    """Run a traced_query to completion, collecting all messages."""
    messages = []
    async for msg in traced_query(**kwargs):
        messages.append(msg)
    return messages


# ── module-scoped OTEL provider ───────────────────────────────────────────

_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():
    """Clear exporter, wire Cascade tracer, yield exporter for assertions."""
    _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. wrap_claude_query creates agent + LLM spans ────────────────────────

class TestWrapClaudeQuery:

    def test_creates_agent_and_llm_spans(self, span_capture):
        from claude_agent_sdk import query, ClaudeAgentOptions
        from cascade.integrations.claude_agents import wrap_claude_query

        options = ClaudeAgentOptions(model="claude-sonnet-4-5-20250929", max_turns=1)
        traced = wrap_claude_query(query, name="TestAgent")

        asyncio.run(_consume(traced, prompt="Reply with exactly: pong", options=options))

        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_agent_span_has_correct_name(self, span_capture):
        from claude_agent_sdk import query, ClaudeAgentOptions
        from cascade.integrations.claude_agents import wrap_claude_query

        options = ClaudeAgentOptions(model="claude-sonnet-4-5-20250929", max_turns=1)
        traced = wrap_claude_query(query, name="MathAgent")

        asyncio.run(_consume(traced, prompt="What is 2+2? Answer with just the number.", options=options))

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


# ── 2. Span attributes ────────────────────────────────────────────────────

class TestSpanAttributes:

    def test_llm_span_has_provider_and_completion(self, span_capture):
        from claude_agent_sdk import query, ClaudeAgentOptions
        from cascade.integrations.claude_agents import wrap_claude_query

        options = ClaudeAgentOptions(model="claude-sonnet-4-5-20250929", max_turns=1)
        traced = wrap_claude_query(query, name="AttrTest")

        asyncio.run(_consume(traced, prompt="Say hello in one word.", options=options))

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

        attrs = dict(llm_spans[0].attributes)
        assert attrs.get("llm.provider") == "anthropic"
        assert "llm.completion" in attrs

    def test_agent_span_has_prompt_and_completion(self, span_capture):
        from claude_agent_sdk import query, ClaudeAgentOptions
        from cascade.integrations.claude_agents import wrap_claude_query

        options = ClaudeAgentOptions(model="claude-sonnet-4-5-20250929", max_turns=1)
        traced = wrap_claude_query(query, name="IOTest")

        asyncio.run(_consume(traced, prompt="What color is the sky? One word.", options=options))

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

        attrs = dict(agent_spans[0].attributes)
        assert "llm.prompt" in attrs
        assert "sky" in attrs["llm.prompt"].lower()
        assert "llm.completion" in attrs


# ── 3. MCP tool spans ─────────────────────────────────────────────────────

class TestToolSpans:

    def test_mcp_tool_creates_tool_span(self, span_capture):
        """MCP tool calls should produce tool spans.

        Note: The Claude Agent SDK's MCP subprocess transport can fail on
        some platforms (especially Windows in pytest). If no LLM spans are
        produced, the MCP transport itself crashed before any tool call
        happened -- that's not a Cascade integration issue, so we skip.
        """
        from claude_agent_sdk import query, ClaudeAgentOptions, create_sdk_mcp_server
        from claude_agent_sdk import tool as sdk_tool
        from cascade.integrations.claude_agents import wrap_claude_query

        @sdk_tool("calculator", "Perform arithmetic", {"operation": str, "a": float, "b": float})
        async def calculator(args: dict) -> dict:
            op, a, b = args["operation"], float(args["a"]), float(args["b"])
            result = {"add": a + b, "multiply": a * b}.get(op, 0)
            return {"content": [{"type": "text", "text": f"{result}"}]}

        server = create_sdk_mcp_server(name="calc", version="1.0.0", tools=[calculator])
        options = ClaudeAgentOptions(
            model="claude-sonnet-4-5-20250929",
            permission_mode="bypassPermissions",
            mcp_servers={"calc": server},
            allowed_tools=["mcp__calc__calculator"],
            max_turns=3,
        )
        traced = wrap_claude_query(query, name="CalcAgent")

        try:
            asyncio.run(_consume(
                traced,
                prompt="Use the calculator tool to multiply 7 by 8. You MUST call the calculator tool.",
                options=options,
            ))
        except BaseException:
            pass

        spans = span_capture.get_finished_spans()
        llm_spans = [s for s in spans if dict(s.attributes).get("cascade.span_type") == "llm"]
        if not llm_spans:
            pytest.skip("MCP transport crashed before tool calls (platform-specific, not a Cascade issue)")

        tool_spans = [s for s in spans if dict(s.attributes).get("cascade.span_type") == "tool"]
        assert tool_spans, f"No tool span found. Spans: {[(s.name, dict(s.attributes).get('cascade.span_type')) for s in spans]}"
        attrs = dict(tool_spans[0].attributes)
        assert attrs["tool.name"] == "calculator"
        assert "tool.input" in attrs


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

class TestErrorHandling:

    def test_integration_does_not_crash(self, span_capture):
        from claude_agent_sdk import query, ClaudeAgentOptions
        from cascade.integrations.claude_agents import wrap_claude_query

        options = ClaudeAgentOptions(model="claude-sonnet-4-5-20250929", max_turns=1)
        traced = wrap_claude_query(query, name="ErrorTest")

        try:
            asyncio.run(_consume(traced, prompt="Just say OK.", options=options))
        except Exception:
            pass

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