#!/usr/bin/env python3
"""
Test MCP Server via MCP Protocol

This script connects to the MCP server using the actual MCP protocol:
- Option 1: Local MCP server via uvx (stdio transport)
- Option 2: Railway-deployed MCP server (HTTP/SSE transport)

Usage:
    python3 test_mcp_protocol.py [--railway]
"""

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

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


async def test_mcp_via_uvx():
    """Test MCP server via uvx (local stdio transport)"""
    print("\n" + "="*80)
    print("MCP SERVER TEST - Via uvx (stdio transport)")
    print("="*80)
    print("Connecting to: uvx dcisionai-mcp-server@latest")
    print("="*80)
    
    try:
        # Check if mcp package is available
        try:
            from mcp import ClientSession, StdioServerParameters
            from mcp.client.stdio import stdio_client
        except ImportError:
            print("❌ MCP SDK not installed")
            print("   Install with: pip install mcp --break-system-packages")
            print("   Or use: python3 -m pip install --user mcp")
            return False
        
        # Create server parameters for uvx
        server_params = StdioServerParameters(
            command="uvx",
            args=["dcisionai-mcp-server@latest"],
            env=None
        )
        
        print("\n🔄 Starting MCP server via uvx...")
        
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                # Initialize the session
                init_result = await session.initialize()
                print(f"✅ Connected to MCP server")
                print(f"   Server: {init_result.serverInfo.name if init_result.serverInfo else 'Unknown'}")
                print(f"   Version: {init_result.serverInfo.version if init_result.serverInfo else 'Unknown'}")
                
                # List resources
                print("\n📋 Listing resources...")
                resources = await session.list_resources()
                
                print(f"   Found {len(resources.resources)} resources:")
                for resource in resources.resources:
                    print(f"     • {resource.uri} - {resource.name}")
                
                # Read models resource
                print("\n📊 Reading models resource...")
                try:
                    result = await session.read_resource("dcisionai://models/list")
                    
                    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"   Type: {model.get('problem_type')}")
                            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 test_mcp_via_railway():
    """Test MCP server via Railway HTTP endpoint"""
    print("\n" + "="*80)
    print("MCP SERVER TEST - Via Railway (HTTP transport)")
    print("="*80)
    
    # Check if Railway MCP server URL is configured
    railway_url = os.getenv("DCISIONAI_MCP_SERVER_URL", "https://dcisionai-mcp-server-production.up.railway.app")
    
    print(f"Railway URL: {railway_url}")
    print("="*80)
    
    try:
        import aiohttp
        
        # Try to connect via HTTP
        async with aiohttp.ClientSession() as session:
            # Test health endpoint
            health_url = f"{railway_url}/health"
            print(f"\n🔄 Testing health endpoint: {health_url}")
            
            async with session.get(health_url) as response:
                if response.status == 200:
                    health_data = await response.json()
                    print(f"✅ Health check passed: {health_data}")
                else:
                    print(f"⚠️  Health check returned: {response.status}")
            
            # Try to list resources via MCP HTTP protocol
            # Note: This depends on Railway MCP server implementation
            print(f"\n📋 Attempting to list resources via HTTP...")
            print("   (Note: Railway MCP server may use SSE/HTTP transport)")
            
            # For now, we'll use direct backend API as fallback
            models_url = f"{railway_url.replace('/mcp', '')}/api/models" if '/mcp' in railway_url else f"{railway_url}/api/models"
            print(f"   Using backend API: {models_url}")
            
            async with session.get(models_url) as response:
                if response.status == 200:
                    data = await response.json()
                    models = data.get("models", [])
                    print(f"\n✅ Retrieved {len(models)} models via backend API:\n")
                    
                    for i, model in enumerate(models, 1):
                        print(f"{i}. {model.get('name')} ({model.get('id')})")
                        print(f"   Domain: {model.get('domain')}")
                        print()
                    
                    return True
                else:
                    print(f"❌ Failed to get models: {response.status}")
                    return False
                    
    except Exception as e:
        print(f"❌ Failed to connect to Railway MCP server: {e}")
        import traceback
        traceback.print_exc()
        return False


async def main():
    """Main function"""
    parser = argparse.ArgumentParser(description="Test MCP server via protocol")
    parser.add_argument("--railway", action="store_true", help="Use Railway HTTP transport instead of uvx stdio")
    parser.add_argument("--railway-url", help="Railway MCP server URL")
    args = parser.parse_args()
    
    if args.railway_url:
        import os
        os.environ["DCISIONAI_MCP_SERVER_URL"] = args.railway_url
    
    if args.railway:
        success = await test_mcp_via_railway()
    else:
        success = await test_mcp_via_uvx()
    
    print("\n" + "="*80)
    if success:
        print("✅ MCP Protocol Test: PASSED")
    else:
        print("❌ MCP Protocol Test: FAILED")
    print("="*80)
    
    return success


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

