"""Tests for evalkit.contrib.pydanticai.PydanticAITracker."""

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock

import pytest

from evalkit.contrib.tracker import BaseTracker
from evalkit.models import ToolCall


@pytest.fixture()
def _import_guard() -> None:
    """Skip if pydantic-ai is not installed."""
    pytest.importorskip("pydantic_ai")


def _make_result(
    messages: list[Any],
    output: Any = None,
) -> MagicMock:
    """Build a mock AgentRunResult with the given messages and output."""
    result = MagicMock()
    result.new_messages.return_value = messages
    result.output = output
    result.data = output  # older pydantic-ai versions used .data
    return result


@pytest.mark.usefixtures("_import_guard")
class TestPydanticAITracker:
    def _make_tracker(self, result: MagicMock) -> Any:
        from evalkit.contrib.pydanticai import PydanticAITracker

        return PydanticAITracker(result)

    def _make_user_prompt(self, content: str) -> Any:
        from pydantic_ai.messages import UserPromptPart

        return UserPromptPart(content=content)

    def _make_tool_call(self, tool_name: str, args: dict[str, Any], tool_call_id: str = "tc-1") -> Any:
        from pydantic_ai.messages import ToolCallPart

        return ToolCallPart(tool_name=tool_name, args=args, tool_call_id=tool_call_id)

    def _make_tool_return(self, tool_name: str, content: Any, tool_call_id: str = "tc-1") -> Any:
        from pydantic_ai.messages import ToolReturnPart

        return ToolReturnPart(tool_name=tool_name, content=content, tool_call_id=tool_call_id)

    def _make_text(self, content: str) -> Any:
        from pydantic_ai.messages import TextPart

        return TextPart(content=content)

    def _make_request(self, parts: list[Any]) -> Any:
        from pydantic_ai.messages import ModelRequest

        return ModelRequest(parts=parts)

    def _make_response(self, parts: list[Any]) -> Any:
        from pydantic_ai.messages import ModelResponse

        return ModelResponse(parts=parts)

    def test_is_base_tracker(self) -> None:
        result = _make_result(messages=[], output="hello")
        tracker = self._make_tracker(result)
        assert isinstance(tracker, BaseTracker)

    def test_extracts_question_from_user_prompt(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("What is Python?")]),
        ]
        tracker = self._make_tracker(_make_result(messages, output="A language."))
        assert tracker.question == "What is Python?"

    def test_extracts_tool_calls(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("Q")]),
            self._make_response(
                [
                    self._make_tool_call("get_weather", {"city": "Paris"}, "tc-1"),
                ]
            ),
        ]
        tracker = self._make_tracker(_make_result(messages))
        assert tracker.tool_calls == [
            ToolCall(name="get_weather", parameters={"city": "Paris"}),
        ]

    def test_extracts_context_from_tool_return(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("Q")]),
            self._make_response(
                [
                    self._make_tool_call("search", {"q": "python"}, "tc-1"),
                ]
            ),
            self._make_request(
                [
                    self._make_tool_return("search", "Python is a language.", "tc-1"),
                ]
            ),
        ]
        tracker = self._make_tracker(_make_result(messages))
        assert tracker.context == ["Python is a language."]

    def test_extracts_answer_from_text_part(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("Q")]),
            self._make_response(
                [
                    self._make_text("The answer is 42."),
                ]
            ),
        ]
        tracker = self._make_tracker(_make_result(messages, output="The answer is 42."))
        assert tracker.answer == "The answer is 42."

    def test_extracts_answer_from_output(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("Q")]),
        ]
        tracker = self._make_tracker(_make_result(messages, output="Final output"))
        assert tracker.answer == "Final output"

    def test_output_overrides_text_part(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("Q")]),
            self._make_response([self._make_text("From text part")]),
        ]
        tracker = self._make_tracker(_make_result(messages, output="From output"))
        assert tracker.answer == "From output"

    def test_empty_result(self) -> None:
        tracker = self._make_tracker(_make_result(messages=[], output=None))
        assert tracker.question == ""
        assert tracker.tool_calls == []
        assert tracker.context == []
        assert tracker.answer == ""

    def test_multiple_tool_calls(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("Q")]),
            self._make_response(
                [
                    self._make_tool_call("weather", {"city": "Paris"}, "tc-1"),
                    self._make_tool_call("search", {"q": "python"}, "tc-2"),
                ]
            ),
        ]
        tracker = self._make_tracker(_make_result(messages))
        assert len(tracker.tool_calls) == 2
        assert tracker.tool_calls[0].name == "weather"
        assert tracker.tool_calls[1].name == "search"

    def test_non_string_tool_return_content(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("Q")]),
            self._make_response(
                [
                    self._make_tool_call("calc", {"expr": "1+1"}, "tc-1"),
                ]
            ),
            self._make_request(
                [
                    self._make_tool_return("calc", 42, "tc-1"),
                ]
            ),
        ]
        tracker = self._make_tracker(_make_result(messages))
        assert tracker.context == ["42"]

    def test_full_flow(self) -> None:
        messages = [
            self._make_request(
                [
                    self._make_user_prompt("Tell me about Python"),
                ]
            ),
            self._make_response(
                [
                    self._make_tool_call("search_kb", {"query": "python"}, "tc-1"),
                ]
            ),
            self._make_request(
                [
                    self._make_tool_return("search_kb", "Python is a programming language.", "tc-1"),
                ]
            ),
            self._make_response(
                [
                    self._make_text("Python is a high-level programming language."),
                ]
            ),
        ]
        tracker = self._make_tracker(_make_result(messages, output="Python is a high-level programming language."))

        assert tracker.question == "Tell me about Python"
        assert tracker.tool_calls == [
            ToolCall(name="search_kb", parameters={"query": "python"}),
        ]
        assert tracker.context == ["Python is a programming language."]
        assert tracker.answer == "Python is a high-level programming language."

    def test_to_test_case(self) -> None:
        messages = [
            self._make_request([self._make_user_prompt("What is 2+2?")]),
            self._make_response(
                [
                    self._make_tool_call("calc", {"expr": "2+2"}, "tc-1"),
                ]
            ),
            self._make_request(
                [
                    self._make_tool_return("calc", "4", "tc-1"),
                ]
            ),
            self._make_response([self._make_text("The answer is 4.")]),
        ]
        tracker = self._make_tracker(_make_result(messages, output="The answer is 4."))
        tc = tracker.to_test_case(
            expected_tool_calls=[ToolCall(name="calc", parameters={"expr": "2+2"})],
            expected_answer="4",
        )
        assert tc.question == "What is 2+2?"
        assert tc.answer == "The answer is 4."
        assert tc.context == ["4"]
        assert tc.expected_answer == "4"
        assert tc.actual_tool_calls == [ToolCall(name="calc", parameters={"expr": "2+2"})]
