"""
PydanticAI agent with tools example — evalkit framework tracker evaluation.

Builds a PydanticAI Agent with a regular tool (get_weather) and a RAG-style
retrieval tool (search_knowledge_base), runs it synchronously on test questions,
and evaluates tool usage + answer quality with evalkit's PydanticAITracker.

The tracker extracts tool calls, RAG context, and the final answer from the
PydanticAI RunResult after each run — no manual extraction or decorators needed.

Pass --server to also upload results to a running evalkit server.

Requirements:
    pip install evalkit-llm[openai,pydanticai]
    export OPENAI_API_KEY=sk-...

Run:
    python examples/pydantic_ai_agent.py
    python examples/pydantic_ai_agent.py --server              # upload to localhost:8000
    python examples/pydantic_ai_agent.py --server http://host:9000  # custom server URL
"""

from __future__ import annotations

from dotenv import load_dotenv

load_dotenv()

from evalkit import ToolCall
from evalkit.contrib.pydanticai import PydanticAITracker
from evalkit.judges.openai import OpenAIJudge

# -- Knowledge base (simulated RAG store) ----------------------------------------

KNOWLEDGE_BASE: dict[str, str] = {
    "python": "Python is a high-level programming language created by Guido van Rossum in 1991.",
    "evalkit": "evalkit is a provider-agnostic LLM/RAG evaluation framework for teams.",
    "fastapi": "FastAPI is a modern web framework for building APIs with Python 3.8+.",
    "pydantic": "Pydantic is a data validation library that uses Python type annotations.",
}


# -- Agent -----------------------------------------------------------------------


def build_agent():
    """Build and return a PydanticAI agent with tools."""
    try:
        from pydantic_ai import Agent
    except ImportError:
        raise ImportError("Install PydanticAI: pip install evalkit-llm[openai,pydanticai]")

    agent = Agent(
        "openai:gpt-4o-mini",
        system_prompt=(
            "You are a helpful assistant with access to tools. "
            "Always use the appropriate tool to answer the user's question. "
            "For factual questions, search the knowledge base first. "
            "Answer concisely based on the tool results."
        ),
    )

    @agent.tool_plain
    def get_weather(city: str, unit: str = "celsius") -> str:
        """Get the current weather for a given city."""
        return f"The weather in {city} is sunny and 22 {unit}."

    @agent.tool_plain
    def search_knowledge_base(query: str) -> str:
        """Search the knowledge base for information on a topic. Use this for any factual question."""
        query_lower = query.lower()
        for key, value in KNOWLEDGE_BASE.items():
            if key in query_lower:
                return value
        return "No relevant information found."

    return agent


# -- Scenarios -------------------------------------------------------------------

SCENARIOS: list[tuple[str, list[ToolCall], str | None]] = [
    (
        "What is the weather like in Paris right now?",
        [ToolCall(name="get_weather", parameters={"city": "Paris"})],
        None,
    ),
    (
        "Who created Python and when was it released?",
        [ToolCall(name="search_knowledge_base", parameters={"query": "Python"})],
        "Python was created by Guido van Rossum in 1991.",
    ),
    (
        "What is the weather in Tokyo in Celsius?",
        [ToolCall(name="get_weather", parameters={"city": "Tokyo", "unit": "celsius"})],
        None,
    ),
    (
        "Tell me about evalkit.",
        [ToolCall(name="search_knowledge_base", parameters={"query": "evalkit"})],
        "evalkit is a provider-agnostic LLM/RAG evaluation framework for teams.",
    ),
    (
        "What is FastAPI?",
        [ToolCall(name="search_knowledge_base", parameters={"query": "FastAPI"})],
        "FastAPI is a modern web framework for building APIs with Python.",
    ),
]


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="PydanticAI agent evaluation with evalkit")
    parser.add_argument(
        "--server",
        nargs="?",
        const="http://localhost:8000",
        default=None,
        metavar="URL",
        help="Upload results to an evalkit server (default: http://localhost:8000)",
    )
    parser.add_argument(
        "--monitor",
        nargs="?",
        const="http://localhost:8000",
        default=None,
        metavar="URL",
        help="Push traces to evalkit monitoring (default: http://localhost:8000)",
    )
    args = parser.parse_args()

    print("Building PydanticAI agent…")
    agent = build_agent()

    judge = None if (args.server or args.monitor) else OpenAIJudge(model="gpt-4o-mini")
    print("Running agent with PydanticAITracker…\n")

    test_cases: list = []
    for question, expected_tools, expected_answer in SCENARIOS:
        result = agent.run_sync(question)
        tracker = PydanticAITracker(result)

        print(f"  Q: {question}")
        print(f"    Captured: {[(tc.name, tc.parameters) for tc in tracker.tool_calls]}")
        print(f"    Expected: {[(tc.name, tc.parameters) for tc in expected_tools]}")
        print(f"    Context:  {tracker.context[:1]}{'…' if len(tracker.context) > 1 else ''}")
        print(f"    Answer:   {tracker.answer[:80]}{'…' if len(tracker.answer) > 80 else ''}")

        tc = tracker.to_test_case(expected_tools, expected_answer)
        if args.server or args.monitor:
            test_cases.append(tc)
        else:
            report = tracker.evaluate(
                judge=judge,
                expected_tool_calls=expected_tools,
                expected_answer=expected_answer,
            )
            report.print_report()

        print()

    if args.server:
        from evalkit.client import EvalKitClient
        from evalkit.models import Suite

        client = EvalKitClient(args.server)
        project = client.create_project("PydanticAI Agent Example")
        evaluation = client.evaluate(project.id, test_cases, suite=Suite.AGENT_TOOL_CALL)
        print(f"Uploaded {len(test_cases)} scenarios, waiting for evaluation…")
        result = client.wait_for_evaluation(evaluation.id)
        print(f"  Score: {result.average_score:.0%}" if result.average_score is not None else "  Evaluation failed")
        print(f"  View:  {args.server}/evaluations/{result.id}")

    if args.monitor:
        from evalkit.client import EvalKitClient
        from evalkit.models import Suite

        client = EvalKitClient(args.monitor)
        project = client.create_project("PydanticAI Agent Example")
        for tc in test_cases:
            client.log_trace(project.id, tc, suite=Suite.AGENT_TOOL_CALL, metadata={"framework": "pydanticai"})
        print(f"Pushed {len(test_cases)} traces → {args.monitor}/projects/{project.id}/monitor")


if __name__ == "__main__":
    main()
