import asyncio, random
from dataclasses import replace
from agents import Agent, Runner, function_tool, handoff
from agents.handoffs import HandoffInputData
from agents import MessageOutputItem, ToolCallOutputItem, ToolCallItem, ReasoningItem
from openai.types.responses import ResponseOutputMessage, ResponseOutputText

# ── filter that nukes reasoning + fn‑call, keeps content as text ─────────────
def strip_reasoning_bundle(data: HandoffInputData) -> HandoffInputData:
    def transform(seq):
        cleaned = []
        for item in seq:
            if isinstance(item, ReasoningItem):
                continue
            cleaned.append(item)
        return cleaned
    return replace(
        data,
        pre_handoff_items = transform(data.pre_handoff_items),
        new_items         = transform(data.new_items),
    )

# ── tools ────────────────────────────────────────────────────────────────────
@function_tool
def make_haiku_about_haikus() -> str:
    return "\n".join([
        "Seventeen small breaths,",
        "A world folded into three,",
        "Haiku makes haiku."
    ])

@function_tool
def choose_random_word(words: list[str]) -> str:
    return random.choice(words)

# ── agent 3 ──────────────────────────────────────────────────────────────────
final_haiku_agent = Agent(
    name="Final‑Haikuist",
    model="o3",
    instructions=("The prior agent will supply ONE word. "
                  "Write a 5‑7‑5 haiku containing that word exactly once."),
)

# ── agent 2 (general model) ─────────────────────────────────────────────────
word_picker_agent = Agent(
    name="Word‑Picker",
    model="gpt-4.1",
    instructions=("Call `choose_random_word` on the list you receive, then hand off "
                  "to Final‑Haikuist; add no extra text."),
    tools=[choose_random_word],
    handoffs=[handoff(final_haiku_agent)],
)

# ── agent 1 ──────────────────────────────────────────────────────────────────
intro_haiku_agent = Agent(
    name="Intro‑Haikuist",
    model="o3",
    instructions=("Call `make_haiku_about_haikus`, then hand off to Word‑Picker."),
    tools=[make_haiku_about_haikus],
    # filter strips reasoning, fn‑call, result bundle just before the hand‑off
    handoffs=[handoff(word_picker_agent, input_filter=strip_reasoning_bundle)],
)

# ── driver ───────────────────────────────────────────────────────────────────
WORDS = ["moon", "blossom", "mountain", "breeze", "river"]

async def main() -> None:
    run = await Runner.run(intro_haiku_agent,
                           f"The candidate words are: {WORDS}")

    for t in run.turns:
        print(f"\n— {t.agent.name} —\n{t.text}")

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