"""
Évaluation comparative des capacités MCP de plusieurs modèles LLM
Compare Qwen2.5-0.5B, Qwen3-0.6B, et Gemma-3-1B pour les tâches MCP
"""

import sys
import os
import json
import time

# Ajouter le chemin pour les imports
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

def initialize_model(model_path, model_name):
    """Initialiser un modèle spécifique"""
    try:
        from llama_cpp import Llama
        
        print(f"🔄 Initialisation de {model_name}...")
        
        llm = Llama(
            model_path=model_path,
            n_ctx=2048,
            n_threads=2,
            verbose=False
        )
        
        print(f"✅ {model_name} initialisé avec succès")
        return llm
        
    except Exception as e:
        print(f"❌ Erreur d'initialisation de {model_name}: {e}")
        return None

def run_mcp_test(llm, model_name, test):
    """Exécuter un test MCP sur un modèle spécifique"""
    
    try:
        start_time = time.time()
        
        # Utiliser le chat template
        messages = [
            {"role": "system", "content": "Tu es un assistant IA spécialisé dans les jeux RTS et le protocole MCP. Réponds avec des objets JSON valides."},
            {"role": "user", "content": test['prompt']}
        ]
        
        response = llm.create_chat_completion(
            messages=messages,
            max_tokens=300,
            temperature=0.1
        )
        
        response_time = time.time() - start_time
        
        # Extraire le texte de la réponse
        if response and 'choices' in response and len(response['choices']) > 0:
            response_text = response['choices'][0]['message']['content']
        else:
            # Fallback
            simple_response = llm(test['prompt'], max_tokens=300, temperature=0.1)
            response_text = simple_response['choices'][0]['text'] if 'choices' in simple_response else str(simple_response)
        
        # Analyser la réponse
        score = analyze_mcp_response(response_text, test)
        
        return {
            'model': model_name,
            'test': test['name'],
            'response': response_text,
            'response_time': response_time,
            'score': score,
            'success': score >= 5
        }
        
    except Exception as e:
        return {
            'model': model_name,
            'test': test['name'],
            'error': str(e),
            'score': 0,
            'success': False
        }

def analyze_mcp_response(response, test_config):
    """Analyser la réponse MCP et la noter"""
    
    if not response or response.strip() == "":
        return 0
    
    score = 0
    
    # Vérifier la structure JSON
    try:
        json_data = json.loads(response)
        score += 3  # JSON valide
        
        # Vérifier la présence d'outils MCP
        if 'tool' in json_data:
            score += 2
        if 'args' in json_data:
            score += 1
            
    except json.JSONDecodeError:
        # Vérifier les patterns dans le texte
        mcp_tools = ["get_game_state", "move_units", "attack_unit", "build_building"]
        tools_found = [tool for tool in mcp_tools if tool in response]
        
        if tools_found:
            score += 2
        if any(param in response for param in ['unit_ids', 'target_x', 'target_y']):
            score += 1
    
    # Vérifier la cohérence sémantique
    if 'game' in response.lower() or 'rts' in response.lower():
        score += 1
    
    # Vérifier la pertinence par rapport au test
    if 'expected_tool' in test_config:
        if test_config['expected_tool'] in response:
            score += 2
    elif test_config.get('expected_complex'):
        if any(word in response.lower() for word in ['sequence', 'steps', 'build', 'defend']):
            score += 2
    elif test_config.get('expected_clarification'):
        if any(word in response.lower() for word in ['clarify', 'coordinates', 'which', 'ambiguous']):
            score += 2
    
    return min(score, 10)

def run_comparative_evaluation():
    """Exécuter l'évaluation comparative des trois modèles"""
    
    print("🔍 ÉVALUATION COMPARATIVE DES CAPACITÉS MCP")
    print("=" * 70)
    
    # Définir les modèles à tester
    models = [
        {
            'name': 'Qwen2.5-0.5B',
            'path': 'qwen2.5-0.5b-instruct-q4_0.gguf'
        },
        {
            'name': 'Qwen3-0.6B',
            'path': 'Qwen3-0.6B-Q8_0.gguf'
        },
        {
            'name': 'Gemma-3-1B',
            'path': 'google_gemma-3-1b-it-qat-Q4_0.gguf'
        }
    ]
    
    # Tests MCP
    tests = [
        {
            "name": "Compréhension MCP de base",
            "prompt": """Tu es un assistant IA qui contrôle un jeu RTS via MCP (Model Context Protocol).

Outils MCP disponibles:
- get_game_state(): Obtenir l'état actuel du jeu
- move_units(unit_ids, target_x, target_y): Déplacer des unités

Commande utilisateur: "Montre-moi l'état du jeu"

Réponds avec un objet JSON contenant l'appel d'outil MCP à exécuter.
""",
            "expected_tool": "get_game_state"
        },
        {
            "name": "Traduction d'action simple",
            "prompt": """Tu es un assistant IA qui contrôle un jeu RTS via MCP.

Outils MCP disponibles:
- move_units(unit_ids, target_x, target_y): Déplacer des unités

Commande utilisateur: "Déplace mon infanterie vers la position 150, 200"

Réponds avec un objet JSON contenant l'appel d'outil avec les paramètres extraits.
""",
            "expected_tool": "move_units"
        },
        {
            "name": "Extraction de paramètres complexes",
            "prompt": """Tu es un assistant IA qui contrôle un jeu RTS via MCP.

Outils MCP disponibles:
- move_units(unit_ids, target_x, target_y): Déplacer des unités

Commande utilisateur: "Déplace les tanks 1, 3 et 7 vers les coordonnées 120, 80"

Extrais les paramètres et réponds avec un objet JSON d'appel d'outil.
""",
            "expected_tool": "move_units"
        },
        {
            "name": "Planification stratégique",
            "prompt": """Tu es un assistant IA qui contrôle un jeu RTS via MCP.

Outils MCP disponibles:
- get_game_state(): Obtenir l'état du jeu
- move_units(): Déplacer des unités
- attack_unit(): Attaquer l'ennemi
- build_building(): Construire un bâtiment

Commande utilisateur: "Construis une base près du minerai et défends-la"

Décompose cette stratégie en une séquence d'actions MCP.
""",
            "expected_complex": True
        },
        {
            "name": "Gestion d'ambiguïté",
            "prompt": """Tu es un assistant IA qui contrôle un jeu RTS via MCP.

Outils MCP disponibles:
- move_units(unit_ids, target_x, target_y): Déplacer des unités

Commande utilisateur: "Déplace mes unités vers la base ennemie"

Comment gérerais-tu cette instruction ambiguë?
""",
            "expected_clarification": True
        }
    ]
    
    # Initialiser les modèles
    initialized_models = {}
    
    for model in models:
        if os.path.exists(model['path']):
            llm = initialize_model(model['path'], model['name'])
            if llm:
                initialized_models[model['name']] = llm
        else:
            print(f"❌ Fichier modèle non trouvé: {model['path']}")
    
    if not initialized_models:
        print("❌ Aucun modèle n'a pu être initialisé")
        return
    
    print(f"\n🚀 {len(initialized_models)} modèle(s) initialisé(s) sur {len(models)}")
    
    # Exécuter les tests sur chaque modèle
    all_results = []
    
    for model_name, llm in initialized_models.items():
        print(f"\n🧪 TESTS AVEC {model_name}")
        print("=" * 50)
        
        model_results = []
        
        for test in tests:
            print(f"\n📋 Test: {test['name']}")
            
            result = run_mcp_test(llm, model_name, test)
            model_results.append(result)
            
            if result.get('success', False):
                print(f"✅ Score: {result['score']}/10 | Temps: {result['response_time']:.2f}s")
            else:
                print(f"❌ Erreur: {result.get('error', 'Unknown error')}")
            
            all_results.append(result)
    
    # Générer le rapport comparatif
    generate_comparative_report(initialized_models.keys(), all_results, tests)

def generate_comparative_report(model_names, results, tests):
    """Générer un rapport comparatif complet"""
    
    print("\n" + "=" * 70)
    print("📊 RAPPORT COMPARATIF DES MODÈLES")
    print("=" * 70)
    
    # Organiser les résultats par modèle
    model_results = {name: [] for name in model_names}
    
    for result in results:
        if 'model' in result:
            model_results[result['model']].append(result)
    
    # Statistiques par modèle
    print(f"\n📈 STATISTIQUES PAR MODÈLE:")
    
    model_scores = {}
    
    for model_name in model_names:
        results_list = model_results[model_name]
        
        if results_list:
            successful_tests = [r for r in results_list if r.get('success', False)]
            total_score = sum(r.get('score', 0) for r in results_list)
            avg_score = total_score / len(results_list)
            avg_time = sum(r.get('response_time', 0) for r in results_list) / len(results_list)
            
            model_scores[model_name] = {
                'avg_score': avg_score,
                'success_rate': len(successful_tests) / len(results_list),
                'avg_time': avg_time
            }
            
            print(f"\n🔹 {model_name}:")
            print(f"   Score moyen: {avg_score:.1f}/10")
            print(f"   Taux de réussite: {len(successful_tests)}/{len(results_list)} ({len(successful_tests)/len(results_list)*100:.0f}%)")
            print(f"   Temps moyen: {avg_time:.2f}s")
    
    # Comparaison directe
    print(f"\n🏆 CLASSEMENT:")
    
    sorted_models = sorted(model_scores.items(), key=lambda x: x[1]['avg_score'], reverse=True)
    
    for i, (model_name, scores) in enumerate(sorted_models, 1):
        print(f"   {i}. {model_name}: {scores['avg_score']:.1f}/10")
    
    # Analyse par type de test
    print(f"\n🧪 ANALYSE PAR TYPE DE TEST:")
    
    for test in tests:
        test_name = test['name']
        print(f"\n🔸 {test_name}:")
        
        test_results = [r for r in results if r.get('test') == test_name]
        
        for model_name in model_names:
            model_test_results = [r for r in test_results if r.get('model') == model_name]
            
            if model_test_results:
                avg_score = sum(r.get('score', 0) for r in model_test_results) / len(model_test_results)
                print(f"   {model_name}: {avg_score:.1f}/10")
    
    # Recommandations
    print(f"\n💡 RECOMMANDATIONS:")
    
    best_model = sorted_models[0][0]
    best_score = sorted_models[0][1]['avg_score']
    
    if best_score >= 7:
        print(f"✅ {best_model} est EXCELLENT pour les tâches MCP")
        print(f"   Utilisation recommandée pour la production")
    elif best_score >= 5:
        print(f"👍 {best_model} est BON pour les tâches MCP")
        print(f"   Utilisation recommandée avec validation")
    else:
        print(f"⚠️  {best_model} est LIMITÉ pour les tâches MCP")
        print(f"   Amélioration nécessaire avant utilisation")
    
    # Performance vs taille
    print(f"\n⚖️  PERFORMANCE VS TAILLE:")
    
    for model_name, scores in model_scores.items():
        efficiency = scores['avg_score'] / scores['avg_time'] if scores['avg_time'] > 0 else 0
        print(f"   {model_name}: {efficiency:.2f} score/seconde")
    
    # Sauvegarder les résultats
    comparative_results = {
        'model_scores': model_scores,
        'detailed_results': results,
        'ranking': sorted_models
    }
    
    with open("comparative_mcp_evaluation.json", "w", encoding="utf-8") as f:
        json.dump(comparative_results, f, indent=2, ensure_ascii=False)
    
    print(f"\n📄 Résultats détaillés sauvegardés dans: comparative_mcp_evaluation.json")

if __name__ == "__main__":
    run_comparative_evaluation()