#!/usr/bin/env python3
"""
Test to reproduce issue #3763: AgentTool creates multiple duplicate MCP connections
"""

import asyncio
from typing import Optional

from google.adk.agents.readonly_context import ReadonlyContext
from google.adk.tools.base_tool import BaseTool
from google.adk.tools.base_toolset import BaseToolset


class MockMCPTool(BaseTool):
    """Simulates an MCP tool."""
    
    async def run_async(self, *, args, tool_context):
        return f"Result from {self.name}"


class MockMCPToolset(BaseToolset):
    """Simulates McpToolset behavior - creates tools on each get_tools() call."""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.get_tools_call_count = 0
        self.connection_count = 0
    
    async def get_tools(
        self, readonly_context: Optional[ReadonlyContext] = None
    ) -> list[BaseTool]:
        self.get_tools_call_count += 1
        self.connection_count += 1
        print(f"  📡 Creating MCP connection #{self.connection_count}")
        print(f"  🔧 Fetching tools from MCP server (call #{self.get_tools_call_count})")
        
        return [
            MockMCPTool(name='github_search', description='Search GitHub'),
            MockMCPTool(name='github_create_issue', description='Create GitHub issue'),
            MockMCPTool(name='github_list_repos', description='List repositories'),
        ]


async def simulate_agent_preprocessing(toolset: MockMCPToolset):
    """Simulates what happens during agent preprocessing in base_llm_flow.py."""
    print("\n🚀 Starting agent preprocessing...")
    print("=" * 70)
    
    print("\n📝 Call 1: Processing toolset for LLM request...")
    tools1 = await toolset.get_tools_with_prefix()
    print(f"   Got {len(tools1)} tools")
    
    print("\n📝 Call 2: Building tools_dict...")
    tools2 = await toolset.get_tools_with_prefix()
    print(f"   Got {len(tools2)} tools")
    
    print("\n📝 Call 3: Additional processing...")
    tools3 = await toolset.get_tools_with_prefix()
    print(f"   Got {len(tools3)} tools")
    
    print("\n" + "=" * 70)
    print(f"📊 RESULTS:")
    print(f"   Total get_tools() calls: {toolset.get_tools_call_count}")
    print(f"   Total MCP connections: {toolset.connection_count}")
    
    if toolset.get_tools_call_count > 1:
        print(f"   ⚠️  ISSUE: get_tools() called {toolset.get_tools_call_count} times!")
    else:
        print(f"   ✅ FIXED: get_tools() called only once (cached)")


async def main():
    toolset = MockMCPToolset(tool_name_prefix='mcp')
    await simulate_agent_preprocessing(toolset)


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