import asyncio
import random
from typing import Any

from agents import Agent, AgentHooks, ItemHelpers, set_default_openai_key, RunContextWrapper, Runner, Tool, function_tool

set_default_openai_key('sk-your-openai-key')

class JokeAgentHooks(AgentHooks):
    async def on_tool_start(self, wrapper: RunContextWrapper[Any], agent: Agent, tool: Tool):
        print(f"-- Hook: On Tool Start --")

    async def on_tool_end(
        self, wrapper: RunContextWrapper[Any], agent: Agent, tool: Tool, result: Any
    ):
        print(f"-- Hook: On Tool End --")

@function_tool
def how_many_jokes() -> int:
    print(f"-- Tool Execution --")
    # await asyncio.sleep(1)
    return random.randint(1, 10)


async def main():
    agent = Agent(
        name="Joker",
        model="gpt-4.1",
        instructions='Tell the user you are searching jokes for them and call the `how_many_jokes` tool simultaneously. Then tell that many jokes',
        tools=[how_many_jokes],
        hooks=JokeAgentHooks()
    )

    result = Runner.run_streamed(
        agent,
        input="Hello",
    )
    print("=== Run starting ===")
    async for event in result.stream_events():
        if event.type == "raw_response_event":
            if event.data.type == "response.output_item.added":
                print(f"-- Raw Output: Item done --")
        elif event.type == "agent_updated_stream_event":
            print(f"-- Agent updated --")
            continue
        elif event.type == "run_item_stream_event":
            if event.item.type == "tool_call_item":
                print("-- Run Item: Tool Called --")
            elif event.item.type == "tool_call_output_item":
                print(f"-- Run Item: Tool Output --")
            elif event.item.type == "message_output_item":
                print(f"-- Run Item: Message Output --")
            else:
                pass  # Ignore other event types

    print("=== Run complete ===")


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