"""Policy enforcement outcomes.

Demonstrates enforce mode denials and approval-required policies using
DeferredToolRequests output.
"""

import asyncio
from dataclasses import dataclass
from pydantic_ai import Agent
from pydantic_ai.exceptions import UnexpectedModelBehavior
from pydantic_ai.tools import DeferredToolRequests
from pydantic_ai.models.test import TestModel


from rgltr import PolicyDecision, PolicyRequest, configure, policy
from rgltr.pydantic_ai import RegulatedAgent


@dataclass
class User:
    name: str
    is_admin: bool


async def require_admin(request: PolicyRequest) -> PolicyDecision:
    user = request.run_context.deps
    if isinstance(user, User) and user.is_admin:
        return PolicyDecision.allow("admin allowed")
    return PolicyDecision.deny("admin required")


async def always_require_approval(_: PolicyRequest) -> PolicyDecision:
    return PolicyDecision.requires_approval("approval required")


async def main() -> None:
    configure(mode="enforce")
    monitor_model = TestModel(call_tools=["delete_account"])
    monitor_agent = Agent(monitor_model, deps_type=User)

    @monitor_agent.tool_plain
    @policy(require=[require_admin])
    async def delete_account(account_id: int) -> str:
        return f"deleted:{account_id}"

    monitor_regulated = RegulatedAgent(monitor_agent, agent_id="policy-monitor")
    try:
        result = await monitor_regulated.run("enforce mode", deps=User("alice", is_admin=False))
        print("enforce output (non-admin):", result.output)
    except UnexpectedModelBehavior as exc:
        print("enforce output (non-admin): blocked:", exc)

    result = await monitor_regulated.run("enforce mode", deps=User("bob", is_admin=True))
    print("enforce output (admin):", result.output)

    approval_model = TestModel(call_tools=["transfer_funds"])
    approval_agent = Agent(approval_model, deps_type=User)

    @approval_agent.tool_plain
    @policy(require=[always_require_approval])
    async def transfer_funds(amount: int) -> str:
        return f"transferred:{amount}"

    approval_regulated = RegulatedAgent(approval_agent, agent_id="policy-approval")
    approval_result = await approval_regulated.run(
        "approval required",
        deps=User("bob", is_admin=True),
        output_type=(str, DeferredToolRequests),
    )
    print("approval requests:", len(approval_result.output.approvals))


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