# /// script
# requires-python = ">=3.12"
# dependencies = ["openai==3.24.0"]
# ///
"""Create a hosted session, answer a read-only tool call, and continue the session."""

import asyncio
import hashlib
import json
import os
from dataclasses import asdict, dataclass

from openai import AsyncOpenAI
from openai.types.beta.agent_session import RequiredActionSessionRequiredActionResourceFunctionCall
from openai.types.beta.agent_session_input_message_param import AgentSessionInputMessageParam
from openai.types.beta.agent_session_input_param import (
    SessionInputParamAgentSessionInputMessage as MessageInput,
)
from openai.types.beta.agent_session_input_param import (
    SessionInputParamAgentSessionInputToolResult as ToolResultInput,
)
from openai.types.beta.agent_session_message import AgentSessionMessage
from openai.types.beta.agent_session_message_content import MessageContentResourceOutputText
from openai.types.beta.input_content_param import InputContentParamInputText


@dataclass(frozen=True)
class TurnResult:
    turn_id: str
    output: str


@dataclass(frozen=True)
class ExampleResult:
    session_id: str
    first_output: str
    followup_output: str


async def finish_turn(
    client: AsyncOpenAI, session_id: str, previous_turn_id: str | None = None
) -> TurnResult:
    """Poll saved state for one session with one application writer."""
    sessions = client.beta.agents.sessions
    async with asyncio.timeout(120):
        while True:
            turns = await sessions.turns.list(session_id, limit=1, order="desc")
            if not turns.data or turns.data[0].id == previous_turn_id:
                await asyncio.sleep(0.25)
                continue
            turn = turns.data[0]
            if turn.status in {"failed", "cancelled"}:
                raise RuntimeError(f"Turn {turn.id} ended with {turn.status}: {turn.error}")
            if turn.status == "completed":
                texts: list[str] = []
                async for item in sessions.items.list(session_id, order="asc"):
                    if (
                        isinstance(item, AgentSessionMessage)
                        and item.turn_id == turn.id
                        and item.role == "assistant"
                    ):
                        texts.extend(
                            part.text
                            for part in item.content
                            if isinstance(part, MessageContentResourceOutputText)
                        )
                return TurnResult(turn.id, "\n".join(texts))
            session = await sessions.retrieve(session_id)
            for action in session.required_actions:
                if not isinstance(action, RequiredActionSessionRequiredActionResourceFunctionCall):
                    raise ValueError("Unexpected required action")
                if action.name != "lookup_launch_status" or action.arguments != {}:
                    raise ValueError("Unexpected tool name or arguments")
                # This fixed lookup is safe to repeat. External effects need durable receipts.
                identity = json.dumps([session_id, action.turn_id, action.call_id])
                event: ToolResultInput = {
                    "type": "agent.session.input.tool_result",
                    "turn_id": action.turn_id,
                    "call_id": action.call_id,
                    "success": True,
                    "output": json.dumps({"status": "ready", "day": "Friday"}),
                }
                await sessions.events.create(
                    session_id,
                    events=[event],
                    idempotency_key=hashlib.sha256(identity.encode()).hexdigest(),
                )
            await asyncio.sleep(0.25)


async def run_example(base_url: str, api_key: str, model: str) -> ExampleResult:
    async with AsyncOpenAI(base_url=base_url, api_key=api_key, max_retries=0) as client:
        session = await client.beta.agents.sessions.create(
            agent={
                "model": model,
                "instructions": "Use lookup_launch_status for the first question. Answer in text.",
                "tools": [
                    {
                        "type": "function",
                        "name": "lookup_launch_status",
                        "description": "Read the launch status.",
                        "parameters": {
                            "type": "object",
                            "properties": {},
                            "required": [],
                            "additionalProperties": False,
                        },
                    }
                ],
            },
            environment={"type": "none"},
            input="Check the launch status and tell me the planned day.",
        )
        # Save this ID in your application before you do more work.
        print(f"Session: {session.id}", flush=True)
        first = await finish_turn(client, session.id)
        text_content: InputContentParamInputText = {
            "type": "input_text",
            "text": "Summarize that in one sentence.",
        }
        message: AgentSessionInputMessageParam = {"role": "user", "content": [text_content]}
        followup_event: MessageInput = {
            "type": "agent.session.input.message",
            "input": [message],
        }
        await client.beta.agents.sessions.events.create(
            session.id,
            events=[followup_event],
            idempotency_key="quickstart-followup-1",
        )
        followup = await finish_turn(client, session.id, first.turn_id)
        return ExampleResult(session.id, first.output, followup.output)


if __name__ == "__main__":
    result = asyncio.run(
        run_example(
            os.environ["AMPERSAND_BASE_URL"],
            os.environ["AMPERSAND_API_KEY"],
            os.environ["AMPERSAND_MODEL"],
        )
    )
    print(json.dumps(asdict(result), indent=2))