"""Builtin tool governance (real LLM).

Demonstrates governing a real builtin tool (web search) with rgltr. This
requires a model/provider that supports PydanticAI's WebSearchTool.
The second run activates a sensitive tag via rgltr.context() to hide the tool.
"""

import asyncio
import os
from enum import Enum

from pydantic_ai import Agent
from pydantic_ai._run_context import RunContext
from pydantic_ai.builtin_tools import WebSearchTool

from rgltr import boundary, configure, context
from rgltr.audit import AuditEvent
from rgltr.pydantic_ai import RegulatedAgent


class Tag(Enum):
    SECRET = "secret"


class Boundary(Enum):
    EXTERNAL = "external"


def _hidden_builtin(events: list[AuditEvent], run_id: str) -> bool:
    return any(
        event.run_id == run_id
        and event.event_type == "tool_hidden"
        and event.tool_kind == "builtin"
        for event in events
    )


async def main() -> None:
    events: list[AuditEvent] = []

    def _audit_sink(event: AuditEvent) -> None:
        events.append(event)

    configure(boundaries={Boundary.EXTERNAL: [Tag.SECRET]}, audit_sink=_audit_sink)

    model = os.environ.get("RGLTR_MODEL", "openai:gpt-4.1-mini")
    agent = Agent(model)
    regulated = RegulatedAgent(agent, agent_id="builtin", govern_builtin_tools=True)

    @boundary(Boundary.EXTERNAL)
    async def web_search(_: RunContext[None]) -> WebSearchTool:
        return WebSearchTool()

    # Builtin tools must be passed explicitly so rgltr can wrap them.
    first = await regulated.run(
        "Whats the latest version of Pydantic AI?",
        run_id="builtin-run-1",
        builtin_tools=[web_search],
    )
    print("first output:", first.output)
    print("builtin hidden (first run):", _hidden_builtin(events, "builtin-run-1"))

    async with context() as ctx:
        ctx.activate([Tag.SECRET])
        second = await regulated.run(
            "Whats the latest version of Pydantic AI?",
            run_id="builtin-run-2",
            builtin_tools=[web_search],
        )

    print("second output:", second.output)
    print("builtin hidden (second run):", _hidden_builtin(events, "builtin-run-2"))


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