"""PydanticAI tracker — extracts tool calls, context, and answer from a RunResult."""

from __future__ import annotations

from collections.abc import Sequence
from typing import Any

from pydantic_ai import UserContent

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

try:
    from pydantic_ai import AgentRunResult
    from pydantic_ai.messages import (
        ModelRequest,
        ModelResponse,
        TextPart,
        ToolCallPart,
        ToolReturnPart,
        UserPromptPart,
    )
except ImportError:
    raise ImportError(
        "PydanticAI is required for PydanticAITracker. " "Install it with: pip install evalkit-llm[pydanticai]"
    )


class PydanticAITracker(BaseTracker):
    """Extracts tool calls, RAG context, and the answer from a PydanticAI RunResult.

    Unlike the OpenAI and LangChain trackers, this performs post-hoc extraction
    from the result object rather than capturing data during execution::

        from pydantic_ai import Agent
        from evalkit.contrib.pydanticai import PydanticAITracker

        result = agent.run_sync("What's the weather?")
        tracker = PydanticAITracker(result)
        tracker.tool_calls  # [ToolCall(name="get_weather", ...)]
        tracker.answer      # "The weather is sunny."
    """

    def __init__(self, result: AgentRunResult[Any]) -> None:
        self._question: str = ""
        self._tool_calls: list[ToolCall] = []
        self._context: list[str] = []
        self._answer: str = ""
        self._extract(result)

    def _normalize_content(self, content: str | Sequence[UserContent]) -> str:
        if isinstance(content, Sequence):
            return "".join(c if isinstance(c, str) else "" for c in content)
        return content

    def _extract(self, result: AgentRunResult[Any]) -> None:
        messages = result.new_messages()

        for msg in messages:
            if isinstance(msg, ModelRequest):
                for part in msg.parts:
                    if isinstance(part, UserPromptPart) and not self._question:
                        content = self._normalize_content(part.content)
                        self._question = content
                    elif isinstance(part, ToolReturnPart) and part.content:
                        self._context.append(str(part.content))

            elif isinstance(msg, ModelResponse):
                for part in msg.parts:
                    if isinstance(part, ToolCallPart):
                        args: dict[str, Any] = part.args_as_dict()
                        self._tool_calls.append(ToolCall(name=part.tool_name, parameters=args))
                    elif isinstance(part, TextPart) and part.content:
                        self._answer = part.content

        output = result.output
        if output is not None:
            self._answer = str(output)

    @property
    def question(self) -> str:
        return self._question

    @property
    def tool_calls(self) -> list[ToolCall]:
        return list(self._tool_calls)

    @property
    def context(self) -> list[str]:
        return list(self._context)

    @property
    def answer(self) -> str:
        return self._answer
