from agno.agent import Agent
from agno.team import Team
from agno.team.mode import TeamMode
from agno.tools.telegram import TelegramTools
from agno.db.postgres import AsyncPostgresDb

# Shared DB for team + member
db = AsyncPostgresDb(
    db_url="postgresql+psycopg_async://...",
    db_schema="ai",
    session_table="agno_sessions",
    approvals_table="agno_approvals",
)

# Member agent with HITL
telegram_tools = TelegramTools(
    token="...", chat_id="...",
    enable_send_message=True,
    requires_confirmation_tools=["send_message"],
)

member = Agent(
    id="telegram-agent",
    name="Telegram Agent",
    model=model,
    tools=[telegram_tools],
    db=db,  # Same DB as team
)

team = Team(
    id="my-team",
    name="My Team",
    model=model,
    members=[member],
    mode=TeamMode.coordinate,
    db=db,
    session_id="test-session",
)

# Step 1: Run — member pauses at send_message (requires_confirmation)
run_id = None
async for ev in team.arun("send a telegram message", stream=True, stream_events=True):
    if isinstance(ev, TeamRunPaused):
        run_id = ev.run_id
        requirements = ev.requirements
        # requirements[0].member_run_id exists
        # requirements[0]._member_run_response exists (in-memory)
        break

# Step 2: Simulate process restart — recreate team from scratch
# (In real scenario: backend restarts, all in-memory objects lost)
del team, member

member2 = Agent(id="telegram-agent", name="Telegram Agent", model=model, tools=[telegram_tools], db=db)
team2 = Team(id="my-team", name="My Team", model=model, members=[member2], mode=TeamMode.coordinate, db=db, session_id="test-session")

# Step 3: Load paused run from DB
session = await db.get_session("test-session")
paused_run = next(r for r in session.runs if r.run_id == run_id)
reqs = paused_run.requirements
for req in reqs:
    req.confirm()

# Step 4: Continue — THIS FAILS
async for ev in team2.acontinue_run(
    run_id=run_id,
    session_id="test-session",
    requirements=reqs,
    stream=True,
    stream_events=True,
):
    pass  # ERROR: "No runs found for run ID {member_run_id}"

# Member agent
Agent(
    id="telegram-agent",
    name="Telegram Agent",
    model=AzureOpenAI(id="gpt-4o"),
    tools=[TelegramTools(requires_confirmation_tools=["send_message"])],
    db=shared_async_postgres_db,  # Same instance as team
    add_history_to_context=False,
)

# Team
Team(
    id="my-team",
    mode=TeamMode.coordinate,
    members=[member_agent],
    db=shared_async_postgres_db,
    session_id="test-session",
    stream_member_events=True,
)

async def asave_session(agent, session):
    if agent.db is not None and session.session_data is not None:
        # Remove the team_id/workflow_id check — members should persist too
        await agent.db.upsert_session(session)

if member_run_output is None and member_run_id:
    # Search team session for member run (it's stored there via upsert_run)
    for r in (run_response.member_responses or []) + (session.runs or []):
        if getattr(r, "run_id", None) == member_run_id:
            member_run_output = r
            break