"""
Example usage of the MCP-RLM Proxy Server

This demonstrates how to use the proxy server with Python REPL-based filtering and search.
"""

import asyncio
import os
from pathlib import Path
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


async def example_with_projection():
    """Example: Using proxy_filter with Python REPL to transform data."""
    print("=== Example: proxy_filter with Python REPL ===\n")
    
    # Get the project root directory (parent of examples/)
    project_root = Path(__file__).parent.parent
    os.chdir(project_root)
    
    config_file = project_root / "mcp.json"
    server_params = StdioServerParameters(
        command="mcp-rlm-proxy",
        args=["--config", str(config_file)],
    )
    
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            
            # Wait for underlying servers to connect
            await asyncio.sleep(3)
            
            # List available tools
            tools = await session.list_tools()
            print(f"Available tools: {[tool.name for tool in tools.tools[:10]]}")
            if len(tools.tools) > 10:
                print(f"... and {len(tools.tools) - 10} more\n")
            else:
                print()
            
            # Check for proxy tools
            proxy_tools = [t for t in tools.tools if t.name.startswith("proxy_")]
            print(f"Proxy tools available: {[t.name for t in proxy_tools]}\n")
            
            print("Typical RLM-style flow:")
            print("  1. Call an underlying tool (e.g., filesystem_read_file)")
            print("  2. If response is large, it's auto-truncated and cached with a cache_id")
            print("  3. Use proxy_filter with Python code to transform the cached data\n")
            print("Example proxy_filter request:")
            print('  Tool: "proxy_filter"')
            print("  Arguments:")
            print("  {")
            print('    "cache_id": "agent_1:ABC123DEF456",')
            print('    "code": "[{k: item[k] for k in [\'name\', \'email\', \'status\']} for item in data]",')
            print('    "return_format": "json"')
            print("  }")
            print("\nThis executes Python code on the cached data and returns only the specified fields.\n")
            
            print("More examples:")
            print("  - Filter with conditions:")
            print('    "code": "[item for item in data if item.get(\'status\') == \'active\']"')
            print("  - Aggregate data:")
            print('    "code": "{\'total\': len(data), \'avg\': sum(x.get(\'score\', 0) for x in data) / len(data) if data else 0}"')
            print("  - Exclude sensitive fields:")
            print('    "code": "{k: v for k, v in data.items() if k not in {\'password\', \'ssn\'}}"')


async def example_with_grep():
    """Example: Using proxy_search to filter tool outputs."""
    print("=== Example: proxy_search ===\n")
    
    # Get the project root directory (parent of examples/)
    project_root = Path(__file__).parent.parent
    os.chdir(project_root)
    
    config_file = project_root / "mcp.json"
    server_params = StdioServerParameters(
        command="mcp-rlm-proxy",
        args=["--config", str(config_file)],
    )
    
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            
            # Wait for underlying servers to connect
            await asyncio.sleep(3)
            
            print("Typical RLM-style flow:")
            print("  1. Call the underlying tool (filesystem_read_file, api_call, etc.)")
            print("  2. If response is large, it's auto-truncated and cached with a cache_id")
            print("  3. Use proxy_search with the cache_id to search inside the cached content\n")
            
            print("Example proxy_search requests:")
            print("\n1. Regex search with context:")
            print('  Tool: "proxy_search"')
            print("  Arguments:")
            print("  {")
            print('    "cache_id": "agent_1:ABC123DEF456",')
            print('    "pattern": "ERROR|WARN",')
            print('    "mode": "regex",')
            print('    "case_insensitive": true,')
            print('    "context_lines": 2,')
            print('    "max_results": 20')
            print("  }")
            
            print("\n2. BM25 relevance search:")
            print('  Tool: "proxy_search"')
            print("  Arguments:")
            print("  {")
            print('    "cache_id": "agent_1:ABC123DEF456",')
            print('    "pattern": "database connection timeout",')
            print('    "mode": "bm25",')
            print('    "top_k": 5')
            print("  }")
            
            print("\n3. Fuzzy search (handles typos):")
            print('  Tool: "proxy_search"')
            print("  Arguments:")
            print("  {")
            print('    "cache_id": "agent_1:ABC123DEF456",')
            print('    "pattern": "authentication",')
            print('    "mode": "fuzzy",')
            print('    "threshold": 0.7')
            print("  }")
            
            print("\n4. Fresh call (call tool and search in one step):")
            print('  Tool: "proxy_search"')
            print("  Arguments:")
            print("  {")
            print('    "tool": "filesystem_read_file",')
            print('    "arguments": {"path": "logs/app.log"},')
            print('    "pattern": "ERROR",')
            print('    "mode": "regex"')
            print("  }")
            print("\nThis calls the tool fresh, caches it, and searches in one step.\n")


async def example_combined():
    """Example: Using proxy_explore, proxy_filter, and proxy_search together."""
    print("=== Example: Combined RLM Workflow ===\n")
    
    # Get the project root directory (parent of examples/)
    project_root = Path(__file__).parent.parent
    os.chdir(project_root)
    
    config_file = project_root / "mcp.json"
    server_params = StdioServerParameters(
        command="mcp-rlm-proxy",
        args=["--config", str(config_file)],
    )
    
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            
            # Wait for underlying servers to connect
            await asyncio.sleep(3)
            
            print("Example: Complete RLM workflow")
            print("This demonstrates the recursive exploration pattern:\n")
            print("Step 1: Call underlying tool (e.g., filesystem_read_file)")
            print("  -> Response may be truncated if large, returns cache_id")
            print("  -> Example cache_id: 'agent_1:ABC123DEF456'\n")
            
            print("Step 2: Explore structure (optional but recommended)")
            print("  Tool: proxy_explore")
            print("  Arguments:")
            print("  {")
            print('    "cache_id": "agent_1:ABC123DEF456",')
            print('    "max_depth": 3')
            print("  }")
            print("  -> Returns: types, field names, sizes, sample\n")
            
            print("Step 3: Filter/transform using Python REPL")
            print("  Tool: proxy_filter")
            print("  Arguments:")
            print("  {")
            print('    "cache_id": "agent_1:ABC123DEF456",')
            print('    "code": "[{k: item[k] for k in [\'name\', \'email\', \'status\']} for item in data]",')
            print('    "return_format": "json"')
            print("  }")
            print("  -> Returns: Only selected fields\n")
            
            print("Step 4: Search within cached data")
            print("  Tool: proxy_search")
            print("  Arguments:")
            print("  {")
            print('    "cache_id": "agent_1:ABC123DEF456",')
            print('    "pattern": "gmail\\.com",')
            print('    "mode": "regex"')
            print("  }")
            print("  -> Returns: Matching lines\n")
            
            print("Alternative: Fresh call mode")
            print("  You can also call tools fresh and process in one step:")
            print("  Tool: proxy_filter")
            print("  Arguments:")
            print("  {")
            print('    "tool": "filesystem_read_file",')
            print('    "arguments": {"path": "data.json"},')
            print('    "code": "[item[\'name\'] for item in data]"')
            print("  }")
            print("  -> Calls tool, caches it, and filters in one step\n")


async def main():
    """Run all examples."""
    print("MCP-RLM Proxy Server Usage Examples\n")
    print("=" * 50 + "\n")
    
    await example_with_projection()
    await example_with_grep()
    await example_combined()
    
    print("\n" + "=" * 50)
    print("Note: Configure underlying servers in mcp.json to test these examples.")
    print("See mcp.json.example for configuration format.")


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

