from __future__ import annotations

from unittest.mock import AsyncMock

from google.adk.agents import Agent
from google.adk.apps import App
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_response import LlmResponse
from google.genai import types  
from google.adk.runners import Runner
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
from google.adk.memory.in_memory_memory_service import InMemoryMemoryService


def create_mock_llm_with_transfer(agent_name: str) -> BaseLlm:
  """Create a mock LLM that returns a transfer_to_agent call."""
  mock_llm = AsyncMock(spec=BaseLlm)
  mock_llm.model = 'mock-model'
  
  # This simulates an LLM that has inferred sibling names from history
  async def mock_generate(*args, **kwargs):
    yield LlmResponse(
        content=types.Content(
            role='model',
            parts=[
                types.Part.from_function_call(
                    name='transfer_to_agent',
                    args={'agent_name': agent_name}
                )
            ],
        ),
        usage_metadata=None,
    )
  
  mock_llm.generate_content_async.side_effect = mock_generate
  
  return mock_llm


async def test_disallow_transfer_to_peers_bypass():
  """Test that demonstrates the bug.
  
  Expected: sub_agent_1 cannot transfer to sub_agent_2 (sibling)
  Actual: sub_agent_1 can transfer to sub_agent_2 by calling transfer_to_agent
  """
  
  # Create mock LLMs
  # root_agent's LLM transfers to sub_agent_1
  root_mock_llm = create_mock_llm_with_transfer('sub_agent_1')
  
  # sub_agent_1's LLM transfers to sub_agent_2 (its sibling) - THIS SHOULD FAIL
  sub_agent_1_mock_llm = create_mock_llm_with_transfer('sub_agent_2')
  
  # sub_agent_2 just responds normally
  sub_agent_2_mock_llm = AsyncMock(spec=BaseLlm)
  sub_agent_2_mock_llm.model = 'mock-model'

  async def mock_sub_agent_2_generate(*args, **kwargs):
    yield LlmResponse(
        content=types.Content(
            role='model',
            parts=[types.Part(text='Hello from sub_agent_2')],
        ),
        usage_metadata=None,
    )
  sub_agent_2_mock_llm.generate_content_async.side_effect = mock_sub_agent_2_generate
  
  # root_agent
  #   ├── sub_agent_1 (disallow_transfer_to_peers=True)
  #   └── sub_agent_2 (disallow_transfer_to_peers=True)
  
  sub_agent_1 = Agent(
      name='sub_agent_1',
      model=sub_agent_1_mock_llm,
      description='Sub agent 1',
      disallow_transfer_to_peers=True,  # Should prevent transfer to sub_agent_2
  )
  
  sub_agent_2 = Agent(
      name='sub_agent_2',
      model=sub_agent_2_mock_llm,
      description='Sub agent 2',
      disallow_transfer_to_peers=True,  # Should prevent transfer to sub_agent_1
  )
  
  root_agent = Agent(
      name='root_agent',
      model=root_mock_llm,
      sub_agents=[sub_agent_1, sub_agent_2],
  )
  
  app = App(name='test_app', root_agent=root_agent)
  
  session_service = InMemorySessionService()
  
  runner = Runner(
      app=app,
      session_service=session_service,
      artifact_service=InMemoryArtifactService(),
      memory_service=InMemoryMemoryService(),
  )
  
  # Create session
  session = await session_service.create_session(
      app_name='test_app',
      user_id='test_user'
  )
  
  events = []
  async for event in runner.run_async(
      user_id=session.user_id,
      session_id=session.id,
      new_message=types.Content(
          role='user',
          parts=[types.Part.from_text(text='test')]
      )
  ):
    events.append(event)
  
  for i, event in enumerate(events):
    print(f"Event {i + 1}:")
    print(f"  Author: {event.author}")
    print(f"  Content: {event.content}")
    if event.actions and event.actions.transfer_to_agent:
      print(f"  Transfer to: {event.actions.transfer_to_agent}")
    print()
  
  # Check if transfer to sibling happened
  transfer_to_sibling_events = [
      e for e in events 
      if e.actions and e.actions.transfer_to_agent == 'sub_agent_2'
      and e.author == 'sub_agent_1'
  ]
  
  assert not transfer_to_sibling_events, "BUG: Transfer to sibling agent should be disallowed but occurred."


if __name__ == '__main__':
  import asyncio
  
  asyncio.run(test_disallow_transfer_to_peers_bypass())