import pytest
from types import SimpleNamespace

# Agents SDK imports
from agents.run import Runner
from agents import (
    Agent,
    AgentHooks,
    RunContextWrapper,
    handoff,
    SQLiteSession,
)

# --- Hooks for Foo ----------------------------------------------------------

class FooHooks(AgentHooks):
    """Inject an extra prompt into the session during hand-off."""

    async def on_handoff(
        self,
        ctx: RunContextWrapper,
        agent: Agent,  # Destination agent (Bar)
        source: Agent,  # Source agent (Foo)
    ) -> None:
        # Append a follow-up question to the conversation history so that Bar
        # sees it when it takes over.
        print(await ctx.context.session.get_items())
        await ctx.context.session.add_items(
            [{"role": "assistant", "content": "name a blue fruit"}]
        )
        print(await ctx.context.session.get_items())


# --- Test -------------------------------------------------------------------

@pytest.mark.asyncio
async def test_handoff_session_runs():
    """Run a simple Foo → Bar conversation with a session in the context."""

    # Shared session so that hooks can write to the conversation history.
    session = SQLiteSession(session_id="test_session")
    runtime_context = SimpleNamespace(session=session)

    # Destination agent (Bar)
    bar = Agent(
        name="Bar",
        instructions="Respond politely to follow-up questions.",
    )

    # Starting agent (Foo)
    foo = Agent(
        name="Foo",
        instructions="Answer with the name of a red fruit, then hand off to Bar.",
        hooks=FooHooks(),
    )

    # Foo hands the conversation to Bar once it's answered.
    foo.handoffs = [handoff(bar)]

    # Kick off the run; no assertions needed – success is that it executes.
    await Runner.run(starting_agent=foo, input="name a red fruit", context=runtime_context)