"""Example of extending the router with dynamic domain-specific tools.

This example shows how to:
- Create domain-specific tools using FunctionToolset
- Pass them to the router via deps.domain_toolset
- The router will dynamically include these tools alongside its default delegation tools

This is the recommended pattern for adding custom functionality to the router
without modifying the framework code.
"""

import asyncio

from pydantic_ai import RunContext
from pydantic_ai.toolsets import FunctionToolset

from pydanticai_multiagent import router_agent, create_mock_search_deps, SearchDeps


# 1. Create a domain-specific toolset
inventory_toolset: FunctionToolset[SearchDeps] = FunctionToolset()


@inventory_toolset.tool
async def check_inventory(
    ctx: RunContext[SearchDeps],
    product_id: str,
) -> str:
    """Check the current inventory level for a product.

    Args:
        ctx: The run context with dependencies.
        product_id: The product ID to check.
    """
    # In a real app, this would query your database
    return f"Product {product_id}: 42 units in stock, 10 reserved, 32 available"


@inventory_toolset.tool
async def reserve_inventory(
    ctx: RunContext[SearchDeps],
    product_id: str,
    quantity: int,
) -> str:
    """Reserve inventory for a pending order.

    Args:
        ctx: The run context with dependencies.
        product_id: The product ID to reserve.
        quantity: Number of units to reserve.
    """
    return f"Reserved {quantity} units of {product_id} for user {ctx.deps.user_id}"


@inventory_toolset.tool
async def list_low_stock(
    ctx: RunContext[SearchDeps],
    threshold: int = 10,
) -> str:
    """List products with inventory below the threshold.

    Args:
        ctx: The run context with dependencies.
        threshold: Minimum stock level (default: 10).
    """
    # Mock response
    return f"Low stock items (below {threshold}): SKU-001 (5), SKU-042 (3), SKU-099 (8)"


async def main():
    # 2. Create deps with your domain toolset
    deps = create_mock_search_deps()
    deps.domain_toolset = inventory_toolset  # Inject domain tools

    # 3. Run the router - it now has access to both:
    #    - Default delegation tools (research, analyst, code, writer, support)
    #    - Your custom inventory tools
    print("Testing router with dynamic inventory tools...\n")

    # This will use the custom inventory tools
    result = await router_agent.run(
        "Check the inventory for product SKU-001 and list any low stock items",
        deps=deps,
    )
    print(f"Response:\n{result.output}\n")

    # This will delegate to a specialist (no inventory tools needed)
    result = await router_agent.run(
        "What is the capital of France?",
        deps=deps,
    )
    print(f"Response:\n{result.output}")


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