from dataclasses import dataclass
from typing import Any

from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.capabilities.abstract import ValidatedToolArgs, WrapToolExecuteHandler
from pydantic_ai.messages import ToolCallPart
from pydantic_ai.tools import ToolDefinition


@dataclass
class ArgInspector(AbstractCapability[None]):
    """Capability that inspects the type of args in wrap_tool_execute."""

    async def wrap_tool_execute(
        self,
        ctx: RunContext[None],
        *,
        call: ToolCallPart,
        tool_def: ToolDefinition,
        args: ValidatedToolArgs,  # <-- type says dict
        handler: WrapToolExecuteHandler,
    ) -> Any:
        print(f"Type annotation says: dict[str, Any]")
        print(f"Actual type at runtime: {type(args).__name__}")
        print(f"Is dict? {isinstance(args, dict)}")
        print(f"Is BaseModel? {isinstance(args, BaseModel)}")

        return await handler(args)


class MyArgument(BaseModel):
    name: str
    count: int


agent = Agent(
    model="openai:gpt-4o",
    capabilities=[ArgInspector()],
)


@agent.tool_plain
def greet(argument: MyArgument) -> str:
    """Greet someone a number of times."""
    return f"Hello {argument.name}! " * argument.count


def main() -> None:
    result = agent.run_sync("Greet Alice 2 times")
    print(f"\nResult: {result.output}")


if __name__ == "__main__":
    main()