"""Example of creating a custom agent with pydanticai-multiagent.

This example shows how to:
- Create a custom agent with typed dependencies
- Define tools for the agent
- Use structured outputs
"""

import asyncio
from dataclasses import dataclass

from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
from pydantic_ai.toolsets import FunctionToolset

from pydanticai_multiagent import BaseDeps, create_mock_base_deps


# 1. Define your output model
class RecipeResult(BaseModel):
    """Structured output for recipe queries."""

    name: str
    ingredients: list[str]
    instructions: list[str]
    prep_time_minutes: int
    difficulty: str


# 2. Create a toolset for your domain
recipe_toolset: FunctionToolset[BaseDeps] = FunctionToolset()


@recipe_toolset.tool
async def search_recipes(
    ctx: RunContext[BaseDeps],
    query: str,
    max_results: int = 5,
) -> str:
    """Search for recipes matching the query.

    Args:
        ctx: The run context with dependencies.
        query: Search terms for recipes.
        max_results: Maximum number of results.
    """
    # In a real app, this would query a database or API
    # Here we return mock data
    return f"Found {max_results} recipes matching '{query}': Pasta Carbonara, Greek Salad, ..."


@recipe_toolset.tool
async def get_recipe_details(
    ctx: RunContext[BaseDeps],
    recipe_name: str,
) -> str:
    """Get detailed information about a specific recipe.

    Args:
        ctx: The run context with dependencies.
        recipe_name: Name of the recipe.
    """
    # Mock implementation
    return f"""
    Recipe: {recipe_name}
    Ingredients: pasta, eggs, bacon, parmesan, black pepper
    Instructions: 1. Cook pasta 2. Fry bacon 3. Mix eggs and cheese 4. Combine
    Prep time: 20 minutes
    Difficulty: Medium
    """


# 3. Create your custom agent
recipe_agent: Agent[BaseDeps, RecipeResult] = Agent(
    "openai:gpt-4o-mini",
    deps_type=BaseDeps,
    output_type=RecipeResult,
    toolsets=[recipe_toolset],
    instructions="""You are a helpful cooking assistant.

When asked about recipes:
1. Use search_recipes to find matching recipes
2. Use get_recipe_details to get full information
3. Return a structured RecipeResult with all details
""",
)


async def main():
    # Use mock dependencies for development
    deps = create_mock_base_deps()

    # Run the agent
    result = await recipe_agent.run(
        "I want to make pasta carbonara",
        deps=deps,
    )

    # Access structured output
    recipe = result.output
    print(f"Recipe: {recipe.name}")
    print(f"Difficulty: {recipe.difficulty}")
    print(f"Prep time: {recipe.prep_time_minutes} minutes")
    print(f"Ingredients: {', '.join(recipe.ingredients)}")
    print("Instructions:")
    for i, step in enumerate(recipe.instructions, 1):
        print(f"  {i}. {step}")


if __name__ == "__main__":
    asyncio.run(main())
