"""OpenAI Agents SDK tracker — captures tool calls, context, and answer via RunHooks."""

from __future__ import annotations

import json
from typing import Any

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

try:
    from agents import RunContextWrapper, Tool
    from agents.lifecycle import AgentHookContext, RunHooksBase
except ImportError:
    raise ImportError(
        "OpenAI Agents SDK is required for OpenAIAgentTracker. "
        "Install it with: pip install evalkit-llm[openai-agents]"
    )


class OpenAIAgentTracker(BaseTracker, RunHooksBase):  # type: ignore[type-arg]
    """Captures tool calls, RAG context, and the final answer from an OpenAI Agents run.

    Pass as ``hooks=`` to ``Runner.run_sync()`` or ``Runner.run()``::

        from agents import Agent, Runner
        from evalkit.contrib.openai import OpenAIAgentTracker

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

    def __init__(self) -> None:
        self._question: str = ""
        self._tool_calls: list[ToolCall] = []
        self._context: list[str] = []
        self._answer: str = ""
        self._pending_tool_name: str = ""

    @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

    async def on_agent_start(
        self,
        context: AgentHookContext[Any],
        agent: Any,
    ) -> None:
        if self._question:
            return
        turn_input = context.turn_input
        if isinstance(turn_input, str):
            self._question = turn_input
        elif isinstance(turn_input, list):
            for item in turn_input:
                if isinstance(item, str):
                    self._question = item
                    break
                if isinstance(item, dict):
                    content = item.get("content", "")
                    role = item.get("role", "")
                    if role == "user" and isinstance(content, str):
                        self._question = content
                        break

    async def on_tool_start(
        self,
        context: RunContextWrapper[Any],
        agent: Any,
        tool: Tool,
    ) -> None:
        self._pending_tool_name = getattr(tool, "name", "")

    async def on_tool_end(
        self,
        context: RunContextWrapper[Any],
        agent: Any,
        tool: Tool,
        result: str,
    ) -> None:
        tool_name = getattr(tool, "name", self._pending_tool_name)
        args: dict[str, Any] = {}
        # Extract arguments from context if available
        tool_arguments = getattr(context, "tool_arguments", None)
        if tool_arguments is not None:
            if isinstance(tool_arguments, str):
                try:
                    args = json.loads(tool_arguments)
                except (json.JSONDecodeError, TypeError):
                    pass
            elif isinstance(tool_arguments, dict):
                args = tool_arguments

        self._tool_calls.append(ToolCall(name=tool_name, parameters=args))

        if isinstance(result, str) and result:
            self._context.append(result)

    async def on_agent_end(
        self,
        context: AgentHookContext[Any],
        agent: Any,
        output: Any,
    ) -> None:
        if isinstance(output, str):
            self._answer = output
        elif output is not None:
            self._answer = str(output)
