"""
Test script for IBM watsonx.ai MCP Server
Tests all tools and resources to ensure proper functionality

Usage:
    python test_ibm_mcp.py

Prerequisites:
    - IBM credentials configured in .env file
    - ibm_watsonx_mcp_server.py in the same directory
    - MCP SDK installed (pip install mcp)
"""

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

async def test_ibm_mcp_server():
    """Test the IBM watsonx.ai MCP server functionality"""
    
    # Set up server parameters
    server_path = Path(__file__).parent / "ibm_watsonx_mcp_server.py"
    
    if not server_path.exists():
        print("Error: ibm_watsonx_mcp_server.py not found in current directory")
        sys.exit(1)
    
    server_params = StdioServerParameters(
        command="python",
        args=[str(server_path)]
    )
    
    try:
        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                
                print("Initializing connection to IBM watsonx.ai MCP server...")
                await session.initialize()
                print("Connected successfully!\n")
                
                # Test 1: List available resources
                print("=" * 50)
                print("Testing Resources")
                print("=" * 50)
                
                resources = await session.list_resources()
                print(f"Available resources: {len(resources.resources)}")
                
                for resource in resources.resources:
                    print(f"  {resource.name}")
                    print(f"     URI: {resource.uri}")
                    print(f"     Description: {resource.description}")
                    print()
                
                # Test 2: Read IBM models resource
                print("\n Reading IBM models resource...")
                try:
                    models_result = await session.read_resource("ibm://models/available")
                    models_data = json.loads(models_result.contents[0].text)
                    print(f"Available IBM models: {len(models_data)} models")
                    
                    for model_name, model_info in models_data.items():
                        print(f"  • {model_name}: {model_info.get('description', 'N/A')}")
                except Exception as e:
                    print(f"Error reading models: {e}")
                
                # Test 3: Read server info
                print("\n🔍 Reading server configuration...")
                try:
                    config_result = await session.read_resource("ibm://config/server-info")
                    config_data = json.loads(config_result.contents[0].text)
                    print(f"Server: {config_data['server_name']} v{config_data['version']}")
                    print(f"   Models available: {config_data['total_models']}")
                    print(f"   Status: {config_data['status']}")
                except Exception as e:
                    print(f"Error reading config: {e}")
                
                # Test 4: List available tools  
                print("\n" + "=" * 50)
                print("🛠️  Testing Tools")
                print("=" * 50)
                
                tools = await session.list_tools()
                print(f"Available tools: {len(tools.tools)}\n")
                
                for tool in tools.tools:
                    print(f"  🔧 {tool.name}")
                    print(f"     Description: {tool.description}")
                    print()
                
                # Test 5: Generate text with IBM Granite
                print("\n Testing Text Generation...")
                try:
                    generation_result = await session.call_tool("generate_text", {
                        "prompt": "Explain the key benefits of IBM watsonx.ai for enterprise AI development in 3 bullet points",
                        "model_name": "granite-13b-instruct",
                        "max_tokens": 200,
                        "temperature": 0.7
                    })
                    
                    response_data = json.loads(generation_result.content[0].text)
                    if response_data["status"] == "success":
                        print("Text generation successful!")
                        print(f"Model: {response_data['model_used']}")
                        print(f"Tokens: {response_data['completion_tokens']} generated")
                        print(f"\nGenerated text:\n{response_data['generated_text']}")
                    else:
                        print(f"Generation failed: {response_data.get('error')}")
                except Exception as e:
                    print(f"Error in text generation: {e}")
                
                # Test 6: Sentiment analysis
                print("\n💭 Testing Sentiment Analysis...")
                try:
                    sentiment_result = await session.call_tool("analyze_sentiment", {
                        "text": "IBM watsonx.ai is an excellent platform for enterprise AI development! Our team loves the new features and capabilities."
                    })
                    
                    sentiment_data = json.loads(sentiment_result.content[0].text)
                    if sentiment_data["status"] == "success":
                        print("Sentiment analysis successful!")
                        print(f"Text analyzed: {sentiment_data['text_analyzed']}")
                        print(f"\nAnalysis result:\n{sentiment_data['sentiment_analysis'][:300]}...")
                    else:
                        print(f"Sentiment analysis failed: {sentiment_data.get('error')}")
                except Exception as e:
                    print(f"Error in sentiment analysis: {e}")
                
                # Test 7: Text summarization
                print("\nTesting Text Summarization...")
                try:
                    long_text = """
                    IBM watsonx.ai is a comprehensive AI development platform that provides enterprise-grade 
                    tools for building, training, and deploying AI models. It offers access to foundation models 
                    like IBM Granite, along with capabilities for fine-tuning, prompt engineering, and model 
                    governance. The platform supports various AI use cases including natural language processing, 
                    code generation, and data analysis. With built-in security and compliance features, 
                    watsonx.ai enables organizations to develop trustworthy AI applications at scale. 
                    It integrates seamlessly with existing IBM Cloud services and provides comprehensive 
                    monitoring and lifecycle management for AI models in production environments.
                    """
                    
                    summary_result = await session.call_tool("summarize_text", {
                        "text": long_text,
                        "max_length": 100
                    })
                    
                    summary_data = json.loads(summary_result.content[0].text)
                    if summary_data["status"] == "success":
                        print(" Text summarization successful!")
                        print(f"Requested length: {summary_data['requested_length']} words")
                        print(f"\nSummary:\n{summary_data['summary']}")
                    else:
                        print(f" Summarization failed: {summary_data.get('error')}")
                except Exception as e:
                    print(f" Error in summarization: {e}")
                
                # Test 8: Model information
                print("\n🔍 Testing Model Information...")
                try:
                    model_info_result = await session.call_tool("get_model_info", {
                        "model_name": "granite-13b-instruct"
                    })
                    
                    model_data = json.loads(model_info_result.content[0].text)
                    if model_data["status"] == "success":
                        print("Model information retrieved!")
                        print(f"Model: {model_data['model_name']}")
                        print(f"Provider: {model_data['provider']}")
                        print(f"Description: {model_data['description']}")
                        if model_data.get('supported_tasks'):
                            print(f"Supported tasks: {', '.join(model_data['supported_tasks'])}")
                    else:
                        print(f" Model info failed: {model_data.get('error')}")
                except Exception as e:
                    print(f" Error getting model info: {e}")
                
                print("\n" + "=" * 50)
                print("🎉 IBM watsonx.ai MCP Server Testing Completed!")
                print("=" * 50)
                print("✅ All tests finished. Check results above for any errors.")
                
    except Exception as e:
        print(f"Failed to connect to MCP server: {e}")
        print("\nTroubleshooting tips:")
        print("1. Make sure your .env file has the correct IBM credentials")
        print("2. Verify ibm_watsonx_mcp_server.py is in the current directory")
        print("3. Check that the MCP SDK is installed: pip install mcp")
        print("4. Test IBM credentials separately")

if __name__ == "__main__":
    print("Starting IBM watsonx.ai MCP Server Tests...")
    print("=" * 50)
    asyncio.run(test_ibm_mcp_server())