import pytest

from bandito.mappers import extract


pydantic_ai = pytest.importorskip("pydantic_ai")
anthropic = pytest.importorskip("anthropic")
openai = pytest.importorskip("openai")


class TestPydanticAIMapper:
    def test_extracts_llm_and_tool_spans(self):
        """Build a minimal AgentRunResult-like flow and verify extraction."""
        # We can't easily construct an AgentRunResult without running an agent,
        # so test the mapper directly.
        from pydantic_ai.messages import (
            ModelRequest,
            ModelResponse,
            TextPart,
            ToolCallPart,
            ToolReturnPart,
            UserPromptPart,
        )
        from pydantic_ai.usage import RequestUsage

        from bandito.mappers.pydantic_ai import (
            _response_text,
            _tool_calls_by_id,
        )

        response = ModelResponse(
            parts=[
                TextPart(content="Let me check."),
                ToolCallPart(
                    tool_name="get_weather",
                    args='{"city": "Tokyo"}',
                    tool_call_id="call_1",
                ),
            ],
            usage=RequestUsage(input_tokens=100, output_tokens=20),
            model_name="claude-haiku-4-5-20251001",
            timestamp=None,
        )
        assert _response_text(response) == "Let me check."
        calls = _tool_calls_by_id(response)
        assert "call_1" in calls
        assert calls["call_1"].tool_name == "get_weather"


class TestAnthropicMapper:
    def test_simple_message(self):
        from anthropic.types import Message, TextBlock, Usage

        msg = Message(
            id="msg_123",
            content=[TextBlock(text="Hello world", type="text")],
            model="claude-haiku-4-5-20251001",
            role="assistant",
            stop_reason="end_turn",
            type="message",
            usage=Usage(input_tokens=10, output_tokens=5),
        )
        output, spans = extract(msg)
        assert output == "Hello world"
        assert len(spans) == 1
        assert spans[0].kind == "llm"
        assert spans[0].model == "claude-haiku-4-5-20251001"
        assert spans[0].provider == "anthropic"
        assert spans[0].input_tokens == 10
        assert spans[0].output_tokens == 5
        assert spans[0].finish_reason == "end_turn"
        assert spans[0].response_id == "msg_123"

    def test_tool_use(self):
        from anthropic.types import Message, TextBlock, ToolUseBlock, Usage

        msg = Message(
            id="msg_456",
            content=[
                TextBlock(text="I'll check.", type="text"),
                ToolUseBlock(
                    id="tu_1",
                    name="get_weather",
                    input={"city": "Tokyo"},
                    type="tool_use",
                ),
            ],
            model="claude-haiku-4-5-20251001",
            role="assistant",
            stop_reason="tool_use",
            type="message",
            usage=Usage(input_tokens=50, output_tokens=30),
        )
        output, spans = extract(msg)
        assert output == "I'll check."
        assert len(spans) == 2
        assert spans[0].kind == "llm"
        assert spans[0].finish_reason == "tool_use"
        assert spans[1].kind == "tool"
        assert spans[1].name == "get_weather"
        assert spans[1].input == {"city": "Tokyo"}
        assert spans[1].parent_id == spans[0].span_id


class TestOpenAIMapper:
    def test_simple_completion(self):
        from openai.types import CompletionUsage
        from openai.types.chat import ChatCompletion
        from openai.types.chat.chat_completion import Choice
        from openai.types.chat.chat_completion_message import ChatCompletionMessage

        comp = ChatCompletion(
            id="chatcmpl-abc",
            choices=[
                Choice(
                    finish_reason="stop",
                    index=0,
                    message=ChatCompletionMessage(
                        content="Hi there", role="assistant"
                    ),
                )
            ],
            created=1234567890,
            model="gpt-4o-mini",
            object="chat.completion",
            usage=CompletionUsage(
                prompt_tokens=15, completion_tokens=8, total_tokens=23
            ),
        )
        output, spans = extract(comp)
        assert output == "Hi there"
        assert len(spans) == 1
        assert spans[0].kind == "llm"
        assert spans[0].model == "gpt-4o-mini"
        assert spans[0].provider == "openai"
        assert spans[0].input_tokens == 15
        assert spans[0].output_tokens == 8
        assert spans[0].finish_reason == "stop"
        assert spans[0].response_id == "chatcmpl-abc"

    def test_tool_calls(self):
        from openai.types import CompletionUsage
        from openai.types.chat import ChatCompletion
        from openai.types.chat.chat_completion import Choice
        from openai.types.chat.chat_completion_message import ChatCompletionMessage
        from openai.types.chat.chat_completion_message_tool_call import (
            ChatCompletionMessageToolCall,
            Function,
        )

        comp = ChatCompletion(
            id="chatcmpl-def",
            choices=[
                Choice(
                    finish_reason="tool_calls",
                    index=0,
                    message=ChatCompletionMessage(
                        content=None,
                        role="assistant",
                        tool_calls=[
                            ChatCompletionMessageToolCall(
                                id="call_1",
                                type="function",
                                function=Function(
                                    name="get_weather",
                                    arguments='{"city": "Tokyo"}',
                                ),
                            )
                        ],
                    ),
                )
            ],
            created=1234567890,
            model="gpt-4o",
            object="chat.completion",
            usage=CompletionUsage(
                prompt_tokens=50, completion_tokens=20, total_tokens=70
            ),
        )
        output, spans = extract(comp)
        assert output is None
        assert len(spans) == 2
        assert spans[0].kind == "llm"
        assert spans[0].finish_reason == "tool_calls"
        assert spans[1].kind == "tool"
        assert spans[1].name == "get_weather"
        assert spans[1].input == {"city": "Tokyo"}
        assert spans[1].parent_id == spans[0].span_id

    def test_system_fingerprint(self):
        from openai.types import CompletionUsage
        from openai.types.chat import ChatCompletion
        from openai.types.chat.chat_completion import Choice
        from openai.types.chat.chat_completion_message import ChatCompletionMessage

        comp = ChatCompletion(
            id="chatcmpl-ghi",
            choices=[
                Choice(
                    finish_reason="stop",
                    index=0,
                    message=ChatCompletionMessage(content="ok", role="assistant"),
                )
            ],
            created=1234567890,
            model="gpt-4o",
            object="chat.completion",
            usage=CompletionUsage(
                prompt_tokens=10, completion_tokens=5, total_tokens=15
            ),
            system_fingerprint="fp_abc123",
        )
        _, spans = extract(comp)
        assert spans[0].provider_data == {"system_fingerprint": "fp_abc123"}


class TestTraceIdShortcut:
    def test_trace_id_property(self):
        from bandito.context import TraceContext

        ctx = TraceContext("proj")
        assert ctx.trace_id == ctx.trace.trace_id


class TestUnknownResult:
    def test_returns_none(self):
        assert extract("just a string") is None
        assert extract(42) is None
        assert extract({"key": "value"}) is None
