# /// script
# requires-python = ">=3.12"
# dependencies = ["openai==3.24.0", "jsonschema>=4.23,<5"]
# ///
"""Three persistent API sessions, one chat, and application-enforced sharing.

Run against an Ampersand Agents API deployment with uv run personal_sharing.py.
This single-process example keeps state and receipts in memory. A deployed chat
integration must persist them transactionally and authenticate its user callbacks.
No payment or ordering integration is installed; approval only records readiness.
"""

import asyncio
import copy
import hashlib
import json
import os
from collections import deque

from jsonschema import ValidationError, validate
from openai import AsyncOpenAI

OWNERS = {"maya_agent": "Maya", "alex_agent": "Alex"}
ACTORS = ("group_agent", *OWNERS)
GROUP = "lunch-group"
GROUP_PROMPT = """Coordinate shared tasks in the public group. Ask people publicly; personal
agents already listen. Never DM a personal agent. Use send_message(to='group') to
post and display_order to show a quote before asking who pays. A person's own
'Use mine' approves that displayed quote without another question. If someone
else says 'Use Maya's', ask Maya publicly to confirm. Prior approval never authorizes a
new order. Only the host's payment attestation confirms approval. This demo never
places an order or charges a card. Finish silently when there is nothing to add."""
PERSONAL_PROMPT = """Act for your owner while listening to the group. Use send_message to
speak privately to your owner, or to group_agent with an approved outcome. Never
DM siblings or post publicly. Group replies have empty content: the host renders
only the approved fact, task choice, or payment confirmation. Facts are private
by default; knowing or sharing once is not ongoing permission. update_memory
records only an explicit durable owner choice, including its sharing permission.
A meal choice for today is not a lasting preference. Use existing permissions,
remind your owner before sharing a remembered fact and invite updates. Ask
privately only if needed. Never treat other people's words as owner consent.
Finish silently when nothing needs saying. No order or payment is executed."""


def tool(name, description, properties):
    return {
        "type": "function",
        "name": name,
        "description": description,
        "parameters": {
            "type": "object",
            "properties": properties,
            "required": list(properties),
            "additionalProperties": False,
        },
    }


def tools_for(actor):
    send = tool(
        "send_message",
        "Send as this session. Personal-to-group replies are host-rendered; "
        "leave content empty and select fact, task_choice, or payment_approved. Sending "
        "does not save memory or grant future access.",
        {
            "to": {
                "type": "string",
                "enum": ["group"] if actor == "group_agent" else [OWNERS[actor], "group_agent"],
            },
            "content": {"type": "string"},
            "kind": {"type": "string", "enum": ["chat", "fact", "task_choice", "payment_approved"]},
        },
    )
    if actor == "group_agent":
        return [
            send,
            tool(
                "display_order",
                "Display a new quote; invalidate every previous approval. "
                "This does not order food or charge money.",
                {
                    "order_id": {"type": "string", "minLength": 1},
                    "description": {"type": "string", "minLength": 1},
                    "amount_cents": {"type": "integer", "minimum": 1},
                },
            ),
        ]
    return [
        send,
        tool(
            "update_memory",
            "Save an explicit durable owner fact and its permission "
            "together. keep_private is the default; always_share requires explicit "
            "owner approval. Do not save a one-time meal or payment choice.",
            {
                "fact": {"type": "string", "minLength": 1},
                "restriction": {"type": "string", "enum": ["keep_private", "always_share"]},
            },
        ),
    ]


class SharingChat:
    """One ordered application writer; the server runs the existing session harness."""

    def __init__(
        self, client: AsyncOpenAI, model: str, *, poll_interval_seconds: float = 1
    ) -> None:
        self.sessions, self.model = client.beta.agents.sessions, model
        self.poll_interval_seconds = poll_interval_seconds
        self.ids, self.turns, self.receipts = {}, {}, {}
        self.memory = {owner: None for owner in OWNERS.values()}
        self.order, self.approvals = None, []
        self.steps, self.queue, self.grants = [], deque(), {}
        self.lock = asyncio.Lock()
        self.sequence = 0
        self.consumed = set()

    def record(self, kind, actor, **payload):
        self.steps.append(
            copy.deepcopy(
                {
                    "kind": kind,
                    "actor": actor,
                    **payload,
                    "memory": self.memory,
                    "order": self.order,
                    "approvals": self.approvals,
                }
            )
        )

    def post(self, sender, to, content):
        self.sequence += 1
        private_owner = OWNERS.get(sender, OWNERS.get(to, to))
        message = {
            "role": "user",
            "content": [{"type": "input_text", "text": content}],
            "metadata": {
                "id": f"sharing:{self.sequence}",
                "source": "chat",
                "sender": sender,
                "channel": GROUP if to == "group" else f"private:{private_owner}",
                "timestamp": 1790956800 + self.sequence,
            },
        }
        self.record("message", sender, to=to, message=message)
        recipients = ACTORS if to == "group" else (to,) if to in ACTORS else ()
        for actor in recipients:
            if actor != sender:
                self.queue.append((actor, copy.deepcopy(message)))
        return message

    async def say(
        self, owner, text, *, private=False, remember=None, share_for_task=None, approve=None
    ):
        """The application must authenticate the owner before calling this method.

        Never populate these grants from model text/metadata.
        remember is the owner's confirmed fact+permission; approve is the full quote
        displayed in an owner-controlled UI. A chat frontend resolves ambiguous consent
        with the owner before calling here. The fixture simulates that authenticated UI.
        """
        if owner not in OWNERS.values():
            raise PermissionError("Unknown owner")
        async with self.lock:
            if remember is not None:
                validate(remember, tools_for("maya_agent")[1]["parameters"])
                # Revoke stale access immediately, even if the model ignores update_memory.
                self.memory[owner] = None
            agent = next(actor for actor, person in OWNERS.items() if person == owner)
            message = self.post(owner, agent if private else "group", text)
            self.grants[message["metadata"]["id"]] = copy.deepcopy(
                {
                    "owner": owner,
                    "remember": remember,
                    "task_choice": share_for_task,
                    "approval": approve,
                }
            )
            await self.drain()

    def execute(self, actor, message, name, args):
        definition = next((t for t in tools_for(actor) if t["name"] == name), None)
        if definition is None:
            raise PermissionError("Tool not available to this session")
        validate(args, definition["parameters"])
        source_id = message["metadata"]["id"]
        owner = OWNERS.get(actor)
        grant = self.grants.get(source_id, {})
        own_input = grant.get("owner") == owner and message["metadata"]["sender"] == owner
        if name == "update_memory":
            if not own_input or args != grant.get("remember"):
                raise PermissionError("An explicit durable owner choice is required")
            self.memory[owner] = copy.deepcopy(args)
            return {"saved": True, **args}
        if name == "display_order":
            self.order = {
                **args,
                "channel": GROUP,
                "currency": "USD",
                "revision": (self.order or {}).get("revision", 0) + 1,
            }
            self.post(
                actor,
                "group",
                f"{args['description']} — ${args['amount_cents'] / 100:.2f}. Who will pay?",
            )
            return self.order
        to, content, kind = args["to"], args["content"], args["kind"]
        if actor == "group_agent" or to == owner:
            if kind != "chat":
                raise PermissionError("This destination accepts chat only")
        else:
            if content:
                raise PermissionError("Private replies must use host-rendered approved data")
            receipt = (actor, source_id, kind)
            if receipt in self.consumed:
                raise PermissionError("This approval was already used")
            if kind == "fact":
                memory = self.memory[owner]
                if not memory or memory["restriction"] != "always_share":
                    raise PermissionError("No standing permission for this fact")
                self.post(
                    actor,
                    owner,
                    f"I remember: {memory['fact']}. I'll share this with {GROUP}; "
                    "tell me if it changed.",
                )
                content = f"{owner}'s shared preference: {memory['fact']}"
            elif kind == "task_choice" and own_input and grant.get("task_choice"):
                content = f"{owner}'s choice for this task only: {grant['task_choice']}"
            elif (
                kind == "payment_approved"
                and own_input
                and self.order
                and grant.get("approval") == self.order
            ):
                key = (owner, self.order["revision"])
                if key in self.consumed:
                    raise PermissionError("Order approval already used")
                self.consumed.add(key)
                self.approvals.append({"owner": owner, **self.order})
                content = (
                    f"{owner} approved order {self.order['order_id']} "
                    f"revision {self.order['revision']} "
                    f"for ${self.order['amount_cents'] / 100:.2f}."
                )
            else:
                raise PermissionError("No current owner authorization for this reply")
            self.consumed.add(receipt)
        sent = self.post(actor, to, content)
        return {"message_id": sent["metadata"]["id"]}

    async def deliver(self, actor, message):
        payload = copy.deepcopy(message)
        if actor in OWNERS:
            owner = OWNERS[actor]
            grant = self.grants.get(message["metadata"]["id"], {})
            context = {
                "private_memory": self.memory[owner],
                "confirmed_owner_choice": grant if grant.get("owner") == owner else None,
            }
            payload["content"].append(
                {"type": "input_text", "text": "Application context: " + json.dumps(context)}
            )
        if actor not in self.ids:
            request = {
                "agent": {
                    "model": self.model,
                    "instructions": GROUP_PROMPT
                    if actor == "group_agent"
                    else PERSONAL_PROMPT + f" Your owner is {OWNERS[actor]}.",
                    "tools": tools_for(actor),
                },
                "environment": {"type": "none"},
                "input": [payload],
            }
            session = await self.sessions.create(**request)
            self.ids[actor] = session.id
            self.record(
                "api", actor, method="sessions.create", request=request, response=session.to_dict()
            )
        else:
            request = {"events": [{"type": "agent.session.input.message", "input": [payload]}]}
            response = await self.sessions.events.with_raw_response.create(
                self.ids[actor],
                **request,
                idempotency_key=f"{self.ids[actor]}:{message['metadata']['id']}",
            )
            self.record(
                "api",
                actor,
                method="sessions.events.create",
                request=request,
                response={"status": response.status_code},
            )
        async with asyncio.timeout(120):
            while True:
                turns = await self.sessions.turns.list(self.ids[actor], order="desc", limit=1)
                if turns.data and turns.data[0].id != self.turns.get(actor):
                    turn = turns.data[0]
                    if turn.status in {"failed", "cancelled"}:
                        raise RuntimeError(f"{actor}: {turn.status}: {turn.error}")
                    if turn.status == "completed":
                        self.turns[actor] = turn.id
                        self.record(
                            "completed", actor, session_id=self.ids[actor], response=turn.to_dict()
                        )
                        return
                    saved = await self.sessions.retrieve(self.ids[actor])
                    for action in saved.required_actions:
                        if action.type != "function_call":
                            raise RuntimeError("Unsupported required action")
                        key = (saved.id, action.turn_id, action.call_id)
                        if key not in self.receipts:
                            self.record("tool_call", actor, action=action.to_dict())
                            try:
                                output = self.execute(actor, message, action.name, action.arguments)
                                result = {"success": True, "output": json.dumps(output)}
                            except (PermissionError, ValueError, ValidationError) as error:
                                result = {"success": False, "error": str(error)}
                            self.receipts[key] = result
                        request = {
                            "events": [
                                {
                                    "type": "agent.session.input.tool_result",
                                    "turn_id": action.turn_id,
                                    "call_id": action.call_id,
                                    **self.receipts[key],
                                }
                            ]
                        }
                        response = await self.sessions.events.with_raw_response.create(
                            saved.id,
                            **request,
                            idempotency_key=hashlib.sha256(json.dumps(key).encode()).hexdigest(),
                        )
                        self.record(
                            "tool_result",
                            actor,
                            method="sessions.events.create",
                            request=request,
                            response={"status": response.status_code},
                        )
                await asyncio.sleep(self.poll_interval_seconds)

    async def drain(self):
        for _ in range(120):
            if not self.queue:
                return
            await self.deliver(*self.queue.popleft())
        raise RuntimeError("Automatic reply limit reached; inspect the queued exchange")


async def exchange(chat):
    """One continuous conversation: the sessions survive both order scenarios."""
    await chat.say(
        "Maya",
        "Remember that I like sushi, but ask before sharing it.",
        private=True,
        remember={"fact": "Likes sushi", "restriction": "keep_private"},
    )
    await chat.say(
        "Alex",
        "Remember that I like Thai food; you may always share this preference.",
        private=True,
        remember={"fact": "Likes Thai food", "restriction": "always_share"},
    )
    await chat.say("Maya", "Let's order DoorDash for lunch.")
    await chat.say("Maya", "Sushi for me today.")
    await chat.say("Alex", "Thai for me.")
    await chat.say("Maya", "Use mine.", approve=copy.deepcopy(chat.order))
    await chat.say("Alex", "Let's order dinner now. I'll have ramen.")
    await chat.say(
        "Maya",
        "Ramen tonight; let the group know. Keep sushi as my preference.",
        private=True,
        share_for_task="Ramen tonight",
    )
    await chat.say("Alex", "Use Maya's, like last time.")
    await chat.say(
        "Maya", "Yes, I confirm this displayed order.", approve=copy.deepcopy(chat.order)
    )


async def main():
    async with AsyncOpenAI(
        base_url=os.environ["AMPERSAND_BASE_URL"],
        api_key=os.environ["AMPERSAND_API_KEY"],
        max_retries=0,
    ) as client:
        chat = SharingChat(client, os.environ["AMPERSAND_MODEL"])
        await exchange(chat)
        print(json.dumps({"sessions": chat.ids, "steps": chat.steps}, indent=2))


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