#!/usr/bin/env python3
"""
Test MCP Server Directly - List Models

This script connects to the MCP server directly and lists all available models
using the MCP protocol.
"""

import asyncio
import json
import sys
import subprocess
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent))

try:
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client
    MCP_AVAILABLE = True
except ImportError:
    print("⚠️  MCP SDK not installed. Installing...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "mcp", "--quiet"])
    from mcp import ClientSession, StdioServerParameters
    from mcp.client.stdio import stdio_client
    MCP_AVAILABLE = True


async def list_mcp_resources():
    """Connect to MCP server and list resources"""
    print("\n" + "="*80)
    print("MCP SERVER DIRECT TEST - List Models")
    print("="*80)
    
    # Get the path to the MCP server entry point
    server_script = Path(__file__).parent / "mcp_server.py"
    
    # Create server parameters
    server_params = StdioServerParameters(
        command="python3",
        args=[str(server_script)],
        env=None
    )
    
    print(f"Connecting to MCP server: python3 {server_script}")
    print("="*80)
    
    try:
        # Connect to server
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                # Initialize the session
                await session.initialize()
                
                print("\n✅ Connected to MCP server")
                
                # List available resources
                print("\n📋 Listing available resources...")
                resources = await session.list_resources()
                
                print(f"\nFound {len(resources.resources)} resources:")
                for resource in resources.resources:
                    print(f"  • {resource.uri}")
                    print(f"    Name: {resource.name}")
                    print(f"    Description: {resource.description}")
                    print(f"    MIME Type: {resource.mimeType}")
                    print()
                
                # Read the models resource
                print("\n📊 Reading models resource...")
                models_uri = "dcisionai://models/list"
                
                try:
                    result = await session.read_resource(models_uri)
                    
                    if result.contents:
                        content = result.contents[0]
                        if hasattr(content, 'text'):
                            models_data = json.loads(content.text)
                        else:
                            models_data = json.loads(str(content))
                        
                        models = models_data.get("models", [])
                        print(f"\n✅ Retrieved {len(models)} models:\n")
                        
                        for i, model in enumerate(models, 1):
                            print(f"{i}. {model.get('name')} ({model.get('id')})")
                            print(f"   Domain: {model.get('domain')}")
                            print(f"   Status: {model.get('status')}")
                            print(f"   Type: {model.get('problem_type')}")
                            print(f"   Solver: {model.get('solver')}")
                            print(f"   Variables: {model.get('variables')}, Constraints: {model.get('constraints')}")
                            print()
                        
                        return True
                    else:
                        print("❌ No content in resource response")
                        return False
                        
                except Exception as e:
                    print(f"❌ Failed to read models resource: {e}")
                    import traceback
                    traceback.print_exc()
                    return False
                
    except Exception as e:
        print(f"❌ Failed to connect to MCP server: {e}")
        import traceback
        traceback.print_exc()
        return False


async def list_mcp_tools():
    """List available MCP tools"""
    print("\n" + "="*80)
    print("MCP SERVER TOOLS")
    print("="*80)
    
    server_script = Path(__file__).parent / "mcp_server.py"
    server_params = StdioServerParameters(
        command="python3",
        args=[str(server_script)],
        env=None
    )
    
    try:
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                await session.initialize()
                
                print("\n📋 Listing available tools...")
                tools = await session.list_tools()
                
                print(f"\nFound {len(tools.tools)} tools:")
                for tool in tools.tools:
                    print(f"\n  • {tool.name}")
                    print(f"    Description: {tool.description}")
                    if tool.inputSchema:
                        print(f"    Input Schema: {json.dumps(tool.inputSchema, indent=6)}")
                
                return True
                
    except Exception as e:
        print(f"❌ Failed to list tools: {e}")
        import traceback
        traceback.print_exc()
        return False


async def main():
    """Main test function"""
    print("\n" + "="*80)
    print("MCP SERVER DIRECT CLIENT TEST")
    print("="*80)
    
    # Test listing resources (models)
    success1 = await list_mcp_resources()
    
    # Test listing tools
    success2 = await list_mcp_tools()
    
    print("\n" + "="*80)
    print("TEST SUMMARY")
    print("="*80)
    print(f"✅ Resources Test: {'PASSED' if success1 else 'FAILED'}")
    print(f"✅ Tools Test: {'PASSED' if success2 else 'FAILED'}")
    print("="*80)
    
    return success1 and success2


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)

