#!/usr/bin/env python3
"""Test all LLM MCP tools."""

import os
import sys
sys.path.insert(0, 'src')
from dotenv import load_dotenv

load_dotenv()

# Import the new tools (direct functions)
from llm_fusion_mcp.server import (
    llm_set_provider, llm_get_provider, llm_list_providers,
    llm_embed_text, llm_similarity, llm_utility_calculator, llm_health_check,
    get_client, PROVIDER_CONFIG
)

def test_provider_management():
    """Test provider management tools."""
    print("🔧 Testing Provider Management")
    print("=" * 50)
    
    # Test provider listing - call the actual function
    try:
        providers_info = {}
        for provider, config in PROVIDER_CONFIG.items():
            api_key_set = bool(os.getenv(config["api_key_env"]))
            providers_info[provider] = {
                "default_model": config["default_model"],
                "api_key_configured": api_key_set,
                "base_url": config["base_url"]
            }
        
        print(f"Available providers: {list(providers_info.keys())}")
        for provider, info in providers_info.items():
            status = "✓" if info["api_key_configured"] else "✗"
            print(f"  {provider}: {status} {info['default_model']}")
        
        print("✓ Provider listing working")
    except Exception as e:
        print(f"✗ Provider listing failed: {e}")

def test_embeddings_and_similarity():
    """Test embeddings and similarity tools."""
    print("\n📊 Testing Embeddings & Similarity")
    print("=" * 50)
    
    if not os.getenv("GOOGLE_API_KEY"):
        print("⚠️  Skipping embeddings test - no Google API key")
        return
    
    try:
        # Test embeddings
        texts = ["I love programming", "Coding is fun", "I hate bugs"]
        
        # Create embeddings using gemini
        embed_result = llm_embed_text(texts, "gemini")
        
        if embed_result.get("success"):
            print(f"✓ Created embeddings: {embed_result['count']} texts, {embed_result['dimensions']} dimensions")
            
            # Test similarity
            sim_result = llm_similarity(texts[0], texts[1], "gemini")
            if sim_result.get("success"):
                print(f"✓ Similarity between '{texts[0]}' and '{texts[1]}': {sim_result['similarity']:.3f}")
            else:
                print(f"✗ Similarity failed: {sim_result.get('error')}")
        else:
            print(f"✗ Embeddings failed: {embed_result.get('error')}")
            
    except Exception as e:
        print(f"✗ Embeddings test failed: {e}")

def test_basic_generation():
    """Test basic text generation."""
    print("\n💬 Testing Text Generation")
    print("=" * 50)
    
    if not os.getenv("GOOGLE_API_KEY"):
        print("⚠️  Skipping generation test - no Google API key")
        return
    
    try:
        # Test direct client usage
        client = get_client("gemini")
        
        response = client.chat.completions.create(
            model="gemini-2.5-flash",
            messages=[{"role": "user", "content": "Say hello in exactly 5 words"}]
        )
        
        text = response.choices[0].message.content
        word_count = len(text.split())
        
        print(f"✓ Generated text: '{text}' ({word_count} words)")
        
    except Exception as e:
        print(f"✗ Text generation failed: {e}")

def test_utility_tools():
    """Test utility and helper tools."""
    print("\n🛠️  Testing Utility Tools")
    print("=" * 50)
    
    # Test calculator
    try:
        calc_result = llm_utility_calculator("add", 15, 25)
        if calc_result.get("success"):
            print(f"✓ Calculator: 15 + 25 = {calc_result['result']}")
        else:
            print(f"✗ Calculator failed: {calc_result.get('error')}")
    except Exception as e:
        print(f"✗ Calculator test failed: {e}")
    
    # Test health check
    try:
        health_result = llm_health_check()
        if health_result.get("success"):
            print(f"✓ Health check: {health_result['overall_status']}")
            healthy_providers = sum(1 for p in health_result['providers'].values() 
                                  if p['status'] in ['healthy', 'configured'])
            total_providers = len(health_result['providers'])
            print(f"  Providers: {healthy_providers}/{total_providers} healthy")
        else:
            print("✗ Health check failed")
    except Exception as e:
        print(f"✗ Health check test failed: {e}")

def test_model_coverage():
    """Test model coverage across providers."""
    print("\n📋 Testing Model Coverage")
    print("=" * 50)
    
    for provider, config in PROVIDER_CONFIG.items():
        print(f"{provider.upper()}:")
        print(f"  Default: {config['default_model']}")
        print(f"  Models: {len(config['models'])} available")
        
        # Show some sample models
        models = config['models']
        if len(models) > 3:
            sample = models[:3] + ['...']
            print(f"  Sample: {', '.join(sample)}")
        else:
            print(f"  All: {', '.join(models)}")

if __name__ == "__main__":
    print("🚀 Comprehensive LLM MCP Server Test")
    print("=" * 70)
    
    test_provider_management()
    test_embeddings_and_similarity()
    test_basic_generation()
    test_utility_tools()
    test_model_coverage()
    
    print("\n" + "=" * 70)
    print("🎉 All tests completed!")
    
    # Summary
    configured_providers = sum(1 for config in PROVIDER_CONFIG.values() 
                             if os.getenv(config["api_key_env"]))
    total_providers = len(PROVIDER_CONFIG)
    
    print(f"📊 Summary: {configured_providers}/{total_providers} providers configured")
    print(f"🔧 Total tools: ~15 LLM tools available")
    print(f"🌐 Supported providers: {', '.join(PROVIDER_CONFIG.keys())}")