"""Agent-native tool registration.

Shows @agent.tool, @agent.tool_plain, and tools=[...] working under RegulatedAgent
without explicit RegulatedToolset wrapping.
"""

import asyncio
from enum import Enum

from pydantic_ai import Agent
from pydantic_ai._run_context import RunContext
from pydantic_ai.models.test import TestModel

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})

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

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

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

    @agent.tool_plain
    async def noop() -> str:
        return "ok"

    regulated = RegulatedAgent(agent, agent_id="native")

    first = await regulated.run("first")
    await regulated.run("second", message_history=first.all_messages())

    tools = _tool_names(model)
    print("tools after tags restored:", tools)


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