"""Toolsets and dynamic toolsets.

Demonstrates FunctionToolset, wrapper toolsets, and a dynamic toolset factory
that RegulatedAgent wraps automatically.
"""

import asyncio
from enum import Enum
from typing import Any

from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from pydantic_ai.toolsets import FilteredToolset
from pydantic_ai.toolsets.function import FunctionToolset

from rgltr import boundary, configure, tag
from rgltr.pydantic_ai import RegulatedAgent


class Tag(Enum):
    CUSTOMER = "customer"


class Boundary(Enum):
    EXTERNAL = "external"


def _tool_names(model: TestModel) -> list[str]:
    params = model.last_model_request_parameters
    if params is None or params.function_tools is None:
        return []
    return [tool.name for tool in params.function_tools]


async def main() -> None:
    configure(boundaries={Boundary.EXTERNAL: True})

    toolset = FunctionToolset()

    @toolset.tool
    @boundary(Boundary.EXTERNAL)
    async def post_to_slack(message: str) -> str:
        return message

    @toolset.tool
    async def internal_note(note: str) -> str:
        return f"note:{note}"

    filtered = FilteredToolset(toolset, lambda _ctx, tool_def: tool_def.name != "internal_note")

    @tag(activates=[Tag.CUSTOMER])
    async def get_customer(customer_id: int) -> str:
        return f"customer:{customer_id}"

    async def dynamic_toolset(_ctx: Any) -> FilteredToolset:
        return filtered

    model = TestModel(call_tools=["get_customer"])
    agent = Agent(model, tools=[get_customer], toolsets=[dynamic_toolset])


    regulated = RegulatedAgent(agent, agent_id="toolsets")
    first = await regulated.run("first")
    await regulated.run("second", message_history=first.all_messages())

    tools = _tool_names(model)
    print("tools after tag restore:", tools)


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