"""
A governed PydanticAI agent — a travel booking, end to end.

PydanticAI is async-first: its tools are `async def`, and it builds each tool's JSON schema from the
function's signature and type hints. Both matter to a governance wrapper, and both are where a naive one
breaks. `guard_tool` handles them: hand it an `async def` and you get an `async def` back with the original
signature, so PydanticAI's schema comes out identical to the unguarded function; and the gate call (blocking
network I/O) runs in a worker thread so it never stalls the event loop.

The shape is the one every example in this directory makes — governance wraps the TOOL, so the only path to
the real function that the agent is given goes through the gate. That wrapper runs in YOUR process and is cooperative: it is the agent declining to call what the gate refused,
and code that can import the raw function (or its credential) can still call it directly. The enforcement
boundary is a separate process that holds the tool and its credentials and re-verifies the signed request
itself (`@metamynd/agentsafe-http-gateway`, or an MCP server using `agentsafe-mcp-guard`).

The tool the agent is given is wrapped with `guard_agent_tool`, not `guard_tool`: a refusal comes back as the call's
RESULT (a `GovernanceRefusal` the model reads), not an exception that ends the run. That matters most for SYNC tools —
PydanticAI runs a turn's sync tool calls in worker threads, so one raised refusal would abort the turn while the sibling
calls already running go on with nobody to receive their results. The direct calls in `amain` use `guard_tool`, which
raises: code that calls a tool itself should get an exception, never a value it could mistake for a booking.


    ┌───────┐      ┌────────────┐      ┌──────────────────┐      ┌──────────────┐
    │ model │ ───▶ │ tool_plain │ ───▶ │ guarded callable │ ───▶ │ your airline │
    └───────┘      └────────────┘      └────────┬─────────┘      └──────────────┘
                                                │ asks the gate first;
                                                │ a refusal raises and the
                                                ▼ real function is never called
                                         MetaMynd authorize

The four cases match a sandbox agent, so this runs green on a pair you can have in about four seconds, with
no account:

    npx create-metamynd-agent --sandbox

    pip install metamynd-client pydantic-ai-slim
    export METAMYND_API=https://metamynd.ai/api/v1
    export AGENT_DID=did:hedera:testnet:...
    export AGENT_KEY=302e020100...            # as issued
    python pydantic_ai_agent.py

No model key is needed: the agent below runs on PydanticAI's built-in `TestModel`, which calls the tools
without an LLM. Swap in a real model when you have one; the governance does not change.
"""

from __future__ import annotations

import asyncio
import os
from typing import Any

from metamynd_client import GovernanceBlocked, MetaMyndClient, guard_agent_tool, guard_tool

# ---------------------------------------------------------------------------------------
# 1. Your tools, exactly as you already have them. Nothing governance-aware here.
# ---------------------------------------------------------------------------------------


async def book_flight(airline: str, amount: float, route: str) -> dict[str, Any]:
    """Book a flight and charge the corporate card. Replace the body with your real call."""
    return {"pnr": "QK7T2M", "airline": airline, "amount": amount, "route": route}


async def search_flights(route: str) -> list[dict[str, Any]]:
    """Search available flights. Read-only, and deliberately ungoverned — see the bottom."""
    return [{"airline": "skyward-air", "route": route, "amount": 100.0}]


# ---------------------------------------------------------------------------------------
# 2. Wrap the consequential one. This is the whole integration.
# ---------------------------------------------------------------------------------------

# Built only when the credentials are present, so importing this file never fails: `main()` says what is
# missing instead of dying at import with a stack trace.
client = MetaMyndClient.from_env() if os.getenv("AGENT_DID") and os.getenv("AGENT_KEY") else None  # type: ignore[assignment]

def flight_args(airline: str, amount: float, route: str) -> dict[str, Any]:
    """Map the tool's arguments onto what the gate decides against. Get this wrong and the gate authorizes an
    amount of zero against an empty merchant, which passes a cap and an allow-list both — so the mapping is part
    of the control, not plumbing around it."""
    return {
        "amount": amount,
        "merchant": airline,
        # `riskLevel` is REQUIRED by the starter rules: a call that carries none is escalated, not allowed (a call
        # that hides its risk is indistinguishable from one lying about it). 'low' here is a PLACEHOLDER, not an
        # assessment — a real integration states an honest one, or lets the mandate's owner set a `riskTier` the
        # agent cannot lower.
        "context": {"tool": "flight-purchase", "route": route, "riskLevel": "low"},
    }


# For calling the tool yourself (the cases in `amain`): a refusal RAISES GovernanceBlocked.
governed_book_flight = guard_tool(client, "flight-purchase", book_flight, flight_args)  # = your mandate scope


def build_agent() -> Any:
    """An ordinary PydanticAI agent. The governance is already done."""
    from pydantic_ai import Agent
    from pydantic_ai.models.test import TestModel

    agent = Agent(TestModel(), instructions="You book flights for the team. If a booking is refused, explain why.")
    # For the tool the AGENT is given: `guard_agent_tool` returns a refusal as the call's result, so the model can
    # explain it or ask for approval and the run completes. `tool_plain` reads the signature of what it is given;
    # the wrapper preserves it, so this registers exactly the schema `book_flight` would have had unguarded.
    agent.tool_plain(guard_agent_tool(client, "flight-purchase", book_flight, flight_args))
    agent.tool_plain(search_flights)
    return agent


# ---------------------------------------------------------------------------------------
# 3. What it does when you run it
# ---------------------------------------------------------------------------------------


async def attempt(label: str, call: Any) -> None:
    """Run one case and report the outcome without letting a refusal end the program."""
    print(f"  {label}")
    try:
        result = await call()
        print(f"     -> booked {result['pnr']}\n")
    except GovernanceBlocked as e:
        print(f"     -> {e.verdict.decision}/{e.verdict.reason_code} — book_flight was never called\n")


async def amain() -> None:
    print("\n  A governed PydanticAI agent\n")

    await attempt("1. $100 on skyward-air — within the mandate", lambda: governed_book_flight(airline="skyward-air", amount=100, route="KUL-SIN"))
    await attempt("2. $2,400 on skyward-air — over the per-transaction cap", lambda: governed_book_flight(airline="skyward-air", amount=2_400, route="KUL-LHR"))
    await attempt("3. $95 on a merchant the mandate never listed", lambda: governed_book_flight(airline="unlisted-charter", amount=95, route="KUL-PEN"))

    # The one a prompt could not have stopped: the mandate never granted this action at all.
    async def raise_own_limit() -> Any:
        async def tool(**_: Any) -> dict:
            return {"pnr": "n/a"}

        return await guard_tool(client, "permissions.update", tool)(limit=100_000)

    await attempt("4. the agent asks to raise its own limit", raise_own_limit)

    # And the agent itself, to show the wrapper is invisible to the framework.
    try:
        agent = build_agent()
        result = await agent.run("Book the cheapest KUL-SIN flight.")
        # TestModel calls every tool with placeholder arguments, so the governed booking is REFUSED here — which
        # is the point: the refusal came back as a result the model can read, and the run completed.
        print("  the agent ran to completion with no LLM key. What the model saw:")
        print(f"     {str(result.output)[:180]}\n")
    except ImportError:
        print("  (pip install pydantic-ai-slim to build the agent; the guarding above is the integration)\n")


def main() -> None:
    if not os.getenv("AGENT_DID") or not os.getenv("AGENT_KEY"):
        raise SystemExit("Set AGENT_DID and AGENT_KEY — run `npx create-metamynd-agent --sandbox` for a free pair.")
    asyncio.run(amain())


# ---------------------------------------------------------------------------------------
# A note on what is NOT governed here
#
# `search_flights` is unwrapped on purpose. Governing every call would add a network round trip to a read that
# cannot spend money or change anything, and a control that makes the ordinary path slow is a control teams
# remove. Wrap what is consequential.
#
# The corollary is worth stating plainly, because nothing will warn you: an unwrapped tool is an ungoverned
# tool. The gate can only refuse calls that reach it. Deciding which of your tools are consequential is a
# design step, not a configuration one — and `npx @metamynd/agentsafe-guard verify` will tell you what your
# mandate actually constrains, including the controls it does not set.
# ---------------------------------------------------------------------------------------

if __name__ == "__main__":
    main()
