from google.adk.agents.llm_agent import Agent

root_agent = Agent(
    model='gemini-3-flash-preview',
    name='root_agent',
    description="A helpful assistant for user questions.",
    instruction="Answer any user question to the best of your knowledge."
)

import asyncio
from dotenv import load_dotenv
import agent
from google.adk.runners import InMemoryRunner
from google.adk.agents.run_config import RunConfig
from google.adk.agents.run_config import StreamingMode
from google.adk.runners import StopSignal
from google.genai.types import Part, Content

load_dotenv()

APP_NAME = "adk_test_app"
USER_ID = "adk_test_user"
STOP_SIGNAL = StopSignal()

async def consume_stream(runner, **kwargs):
    async for event in runner.run_async(**kwargs):
        content = ""
        if hasattr(event, 'text'):
            content = event.text
        elif hasattr(event, 'content') and event.content.parts:
            content = event.content.parts[0].text

        if event.interrupted:
            print('Runner interrupted')
        if content:
            print(content)

async def manual_test_interrupt():
    client = agent.root_agent
    runner = InMemoryRunner(agent=client, app_name=APP_NAME)
    
    session = await runner.session_service.create_session(
        app_name=APP_NAME, user_id=USER_ID
    )

    message = Content(
        role="user",
        parts=[Part(text="Generate a long story...")]
    )
    
    print("Starting run_async")
    
    task = asyncio.create_task(consume_stream(
        runner=runner, 
        session_id=session.id, 
        user_id=USER_ID, 
        new_message=message, 
        stop_signal=STOP_SIGNAL,
        run_config=RunConfig(streaming_mode=StreamingMode.SSE)
    ))

    # Wait for a few seconds to let it start generating text
    await asyncio.sleep(5) 
    
    # Interrupts the run
    STOP_SIGNAL.stop()
    
    try:
        # Checks to make sure the task is actually ending
        # This should print immediately after "Runner interrupted"
        await asyncio.wait_for(task, timeout=10)
        print("run_async complete")
    except asyncio.TimeoutError:
        print("Fail: took more than 10 seconds for the runner to stop")

asyncio.run(manual_test_interrupt())