#!/usr/bin/env python3
"""
Test script for the MCP protocol implementation.
Tests MCP client functionality separately before integration.
"""

import asyncio
import logging
import os
import sys
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

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

from src.llm_fusion_mcp.config import load_config, MCPServerConfig
from src.llm_fusion_mcp.error_handling import ErrorHandler
from src.llm_fusion_mcp.mcp_client import MCPClientManager, MCPSTDIOClient, MCPHTTPClient

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

async def test_mcp_client_manager():
    """Test the MCP client manager functionality."""
    print("=" * 60)
    print("TESTING MCP CLIENT MANAGER")
    print("=" * 60)
    
    try:
        error_handler = ErrorHandler()
        manager = MCPClientManager(error_handler)
        
        print("✅ MCP Client Manager initialized")
        
        # Test with a simple mock STDIO server config
        test_config = MCPServerConfig(
            type="stdio",
            namespace="test",
            command=["echo", "test-mcp-server"],  # Simple command that won't actually work as MCP
            auto_start=True,
            timeout=5
        )
        
        print("🔧 Testing connection (will likely fail - this is expected for test command)...")
        
        # This will fail because echo is not an MCP server, but tests our connection logic
        result = await manager.connect_server(test_config)
        print(f"   - Connection result: {result}")
        
        # Test tool and resource listing
        tools = manager.get_available_tools()
        resources = manager.get_available_resources()
        status = manager.get_connection_status()
        
        print(f"   - Available tools: {len(tools)}")
        print(f"   - Available resources: {len(resources)}")
        print(f"   - Connection status: {status}")
        
        print("✅ MCP Client Manager basic functionality test passed")
        return True
        
    except Exception as e:
        print(f"❌ MCP Client Manager test failed: {e}")
        return False

async def test_mcp_client_creation():
    """Test creating MCP clients without connecting."""
    print("\n" + "=" * 60)
    print("TESTING MCP CLIENT CREATION")
    print("=" * 60)
    
    try:
        error_handler = ErrorHandler()
        
        # Test STDIO client creation
        stdio_config = MCPServerConfig(
            type="stdio", 
            namespace="test_stdio",
            command=["echo", "test"],
            timeout=5
        )
        
        stdio_client = MCPSTDIOClient(stdio_config, error_handler)
        print("✅ STDIO MCP Client created")
        
        # Test HTTP client creation
        http_config = MCPServerConfig(
            type="http",
            namespace="test_http", 
            url="https://httpbin.org/get",  # Test URL
            timeout=5
        )
        
        http_client = MCPHTTPClient(http_config, error_handler)
        print("✅ HTTP MCP Client created")
        
        print("✅ MCP Client creation test passed")
        return True
        
    except Exception as e:
        print(f"❌ MCP Client creation test failed: {e}")
        return False

async def test_config_loading():
    """Test loading MCP server configurations."""
    print("\n" + "=" * 60)
    print("TESTING MCP CONFIGURATION LOADING")
    print("=" * 60)
    
    try:
        config = load_config()
        
        print(f"✅ Configuration loaded")
        print(f"   - MCP servers configured: {len(config.mcp_servers)}")
        
        for server_name, server_config in config.mcp_servers.items():
            print(f"   - {server_name}: {server_config.type} ({server_config.namespace})")
            
            # Test config validation
            if server_config.type == "stdio":
                if not server_config.command:
                    print(f"     ❌ STDIO server missing command")
                else:
                    print(f"     ✅ Command: {' '.join(server_config.command)}")
            elif server_config.type == "http":
                if not server_config.url:
                    print(f"     ❌ HTTP server missing URL")
                else:
                    print(f"     ✅ URL: {server_config.url}")
        
        print("✅ MCP configuration loading test passed")
        return True
        
    except Exception as e:
        print(f"❌ MCP configuration loading test failed: {e}")
        return False

async def test_realistic_mcp_connection():
    """Test connection to a realistic MCP server (if available)."""
    print("\n" + "=" * 60)
    print("TESTING REALISTIC MCP CONNECTION")
    print("=" * 60)
    
    print("🔧 Checking for available MCP servers...")
    
    # Check if uvx is available (for filesystem server)
    import shutil
    if shutil.which("uvx"):
        print("   ✅ uvx found - can test filesystem MCP server")
        
        try:
            error_handler = ErrorHandler()
            manager = MCPClientManager(error_handler)
            
            # Test filesystem server
            fs_config = MCPServerConfig(
                type="stdio",
                namespace="fs",
                command=["uvx", "mcp-server-filesystem"],
                args=["/tmp"],  # Safe directory to test with
                auto_start=True,
                timeout=10
            )
            
            print("🔧 Attempting to connect to filesystem MCP server...")
            result = await manager.connect_server(fs_config)
            
            if result:
                print("🎉 Successfully connected to real MCP server!")
                
                # Test tool discovery
                tools = manager.get_available_tools()
                print(f"   - Discovered tools: {list(tools.keys())}")
                
                # Test a simple tool execution (if any tools were discovered)
                if tools:
                    tool_name = list(tools.keys())[0]
                    print(f"   - Testing tool execution: {tool_name}")
                    
                    result = await manager.execute_tool(tool_name, {})
                    print(f"   - Tool execution result: {result.get('success', False)}")
                
                # Clean up
                await manager.disconnect_server("fs")
                
            else:
                print("⚠️  Could not connect to filesystem MCP server (this might be expected)")
                
        except Exception as e:
            print(f"⚠️  Realistic MCP connection test encountered error: {e}")
            print("   (This is often expected - MCP servers need specific setup)")
    
    else:
        print("   ⚠️  uvx not found - skipping filesystem MCP server test")
    
    # Check for npx (for other servers)
    if shutil.which("npx"):
        print("   ✅ npx found - could test NPM-based MCP servers")
    else:
        print("   ⚠️  npx not found - can't test NPM-based MCP servers")
    
    print("✅ Realistic MCP connection test completed")
    return True

async def main():
    """Run all MCP protocol tests."""
    print("🚀 Starting MCP Protocol Tests")
    print("=" * 60)
    
    tests = [
        ("Configuration Loading", test_config_loading),
        ("MCP Client Creation", test_mcp_client_creation),  
        ("MCP Client Manager", test_mcp_client_manager),
        ("Realistic MCP Connection", test_realistic_mcp_connection)
    ]
    
    passed = 0
    total = len(tests)
    
    for test_name, test_func in tests:
        print(f"\n🧪 Running test: {test_name}")
        try:
            if await test_func():
                passed += 1
                print(f"✅ {test_name} PASSED")
            else:
                print(f"❌ {test_name} FAILED")
        except Exception as e:
            print(f"❌ {test_name} FAILED with exception: {e}")
    
    print("\n" + "=" * 60)
    print("MCP PROTOCOL TEST RESULTS")
    print("=" * 60)
    print(f"📊 Tests passed: {passed}/{total}")
    
    if passed == total:
        print("🎉 All MCP protocol tests passed!")
        print("\n💡 Next steps:")
        print("   - MCP protocol implementation is ready")
        print("   - Can integrate with main orchestrator")
        print("   - Ready to test with real MCP servers")
    else:
        print("⚠️  Some tests failed, but this is often expected:")
        print("   - MCP servers need specific installation")
        print("   - Configuration may need adjustment") 
        print("   - Core protocol implementation should still work")
    
    return passed == total

if __name__ == "__main__":
    try:
        success = asyncio.run(main())
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\n\n⚠️  Tests interrupted by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ Unexpected error during testing: {e}")
        sys.exit(1)