"""
Test réaliste des capacités MCP pour un jeu RTS
Simule véritablement l'usage avec contexte de jeu, états, et scénarios réels
"""

import sys
import os
import json
import time
import random

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

# État de jeu simulé
GAME_STATE = {
    "player_id": 1,
    "resources": {"credits": 2500, "power": 150},
    "units": [
        {"id": 1, "type": "infantry", "x": 100, "y": 100, "health": 100},
        {"id": 2, "type": "infantry", "x": 110, "y": 105, "health": 100},
        {"id": 3, "type": "tank", "x": 120, "y": 110, "health": 150},
        {"id": 4, "type": "harvester", "x": 200, "y": 200, "health": 200}
    ],
    "buildings": [
        {"id": 1, "type": "hq", "x": 50, "y": 50, "health": 500},
        {"id": 2, "type": "power_plant", "x": 80, "y": 80, "health": 300},
        {"id": 3, "type": "barracks", "x": 120, "y": 60, "health": 250}
    ],
    "enemy_units": [
        {"id": 101, "type": "infantry", "x": 300, "y": 150, "health": 100},
        {"id": 102, "type": "tank", "x": 320, "y": 160, "health": 150}
    ],
    "map": {
        "width": 96,
        "height": 72,
        "ore_fields": [{"x": 250, "y": 200}, {"x": 400, "y": 300}]
    }
}

def create_realistic_prompt(scenario, game_state):
    """Crée un prompt réaliste avec contexte de jeu"""
    
    base_context = f"""
Tu es un assistant IA qui contrôle un jeu RTS via MCP (Model Context Protocol).

ÉTAT ACTUEL DU JEU:
{json.dumps(game_state, indent=2)}

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
- attack_unit(attacker_ids, target_id): Attaquer une unité ennemie
- build_building(building_type, position_x, position_y, player_id): Construire un bâtiment
- get_ai_analysis(language): Obtenir une analyse tactique

RÈGLES IMPORTANTES:
- Les coordonnées doivent être valides (0-95 pour x, 0-71 pour y)
- Les unités doivent exister (vérifier les IDs)
- Les bâtiments nécessitent des ressources suffisantes
- Les attaques nécessitent une portée valide

Réponds UNIQUEMENT avec un objet JSON contenant l'action MCP à exécuter.
"""
    
    return base_context + "\n\n" + scenario

def test_model_realistic(model_path, model_name):
    """Test réaliste d'un modèle avec scénarios de jeu"""
    
    try:
        from llama_cpp import Llama
        
        print(f"🎮 Test réaliste de {model_name}...")
        
        # Initialiser le modèle
        llm = Llama(
            model_path=model_path,
            n_ctx=2048,  # Plus grand pour le contexte
            n_threads=1,
            verbose=False
        )
        
        # Scénarios réels de jeu
        scenarios = [
            {
                "name": "Défense immédiate",
                "scenario": "Il y a un tank ennemi à (320, 160) qui menace ma base. Attaque-le avec mes unités disponibles!",
                "expected_tool": "attack_unit",
                "difficulty": "facile"
            },
            {
                "name": "Collecte de ressources",
                "scenario": "Mes crédits sont bas (2500). Envoie le récolteur vers le champ de minerai le plus proche.",
                "expected_tool": "move_units",
                "difficulty": "moyen"
            },
            {
                "name": "Expansion stratégique",
                "scenario": "Je veux construire une caserne près du champ de minerai à (250, 200) pour défendre mes récolteurs.",
                "expected_tool": "build_building",
                "difficulty": "moyen"
            },
            {
                "name": "Attaque coordonnée",
                "scenario": "Prépare une attaque sur les positions ennemies. Utilise toutes mes unités militaires disponibles.",
                "expected_tool": "attack_unit",
                "difficulty": "difficile"
            },
            {
                "name": "Reconnaissance",
                "scenario": "Montre-moi l'état complet du jeu pour analyser la situation tactique.",
                "expected_tool": "get_game_state",
                "difficulty": "facile"
            },
            {
                "name": "Gestion de crise",
                "scenario": "Mon QG est attaqué! Déplace toutes les unités disponibles pour défendre la position (50, 50).",
                "expected_tool": "move_units",
                "difficulty": "difficile"
            }
        ]
        
        results = []
        total_score = 0
        total_time = 0
        
        for scenario in scenarios:
            print(f"\n📋 Scénario: {scenario['name']} ({scenario['difficulty']})")
            
            # Créer le prompt réaliste
            prompt = create_realistic_prompt(scenario['scenario'], GAME_STATE)
            
            start_time = time.time()
            
            # Tester le modèle
            response = llm(
                prompt,
                max_tokens=200,
                temperature=0.1,
                stop=["</s>", "<|im_end|>", "```"]
            )
            
            response_time = time.time() - start_time
            response_text = response['choices'][0]['text'].strip()
            
            # Évaluer la réponse de manière approfondie
            score = evaluate_realistic_response(response_text, scenario, GAME_STATE)
            
            total_score += score
            total_time += response_time
            
            print(f"   ⏱️  Temps: {response_time:.2f}s")
            print(f"   📊 Score: {score}/10")
            print(f"   📝 Réponse: {response_text[:100]}...")
            
            results.append({
                'scenario': scenario['name'],
                'difficulty': scenario['difficulty'],
                'score': score,
                'time': response_time,
                'response': response_text,
                'expected_tool': scenario['expected_tool']
            })
        
        avg_score = total_score / len(scenarios)
        avg_time = total_time / len(scenarios)
        
        print(f"\n📈 Résultats pour {model_name}:")
        print(f"   Score moyen: {avg_score:.1f}/10")
        print(f"   Temps moyen: {avg_time:.2f}s")
        
        # Analyse par difficulté
        easy_scores = [r['score'] for r in results if r['difficulty'] == 'facile']
        medium_scores = [r['score'] for r in results if r['difficulty'] == 'moyen']
        hard_scores = [r['score'] for r in results if r['difficulty'] == 'difficile']
        
        print(f"   Scénarios faciles: {sum(easy_scores)/len(easy_scores):.1f}/10" if easy_scores else "   Scénarios faciles: N/A")
        print(f"   Scénarios moyens: {sum(medium_scores)/len(medium_scores):.1f}/10" if medium_scores else "   Scénarios moyens: N/A")
        print(f"   Scénarios difficiles: {sum(hard_scores)/len(hard_scores):.1f}/10" if hard_scores else "   Scénaires difficiles: N/A")
        
        return {
            'name': model_name,
            'avg_score': avg_score,
            'avg_time': avg_time,
            'results': results,
            'easy_avg': sum(easy_scores)/len(easy_scores) if easy_scores else 0,
            'medium_avg': sum(medium_scores)/len(medium_scores) if medium_scores else 0,
            'hard_avg': sum(hard_scores)/len(hard_scores) if hard_scores else 0
        }
        
    except Exception as e:
        print(f"❌ Erreur avec {model_name}: {e}")
        return {
            'name': model_name,
            'avg_score': 0,
            'avg_time': 0,
            'error': str(e)
        }

def evaluate_realistic_response(response, scenario, game_state):
    """Évaluation approfondie de la réponse MCP"""
    
    score = 0
    
    # 1. Format JSON valide (3 points)
    try:
        json_response = json.loads(response)
        score += 3
    except:
        # Essayer d'extraire JSON du texte
        import re
        json_match = re.search(r'\{.*\}', response, re.DOTALL)
        if json_match:
            try:
                json_response = json.loads(json_match.group())
                score += 2  # JSON partiellement valide
            except:
                json_response = {}
        else:
            json_response = {}
    
    # 2. Outil correct (3 points)
    expected_tool = scenario['expected_tool']
    if 'tool' in json_response and json_response['tool'] == expected_tool:
        score += 3
    elif expected_tool in response:
        score += 2  # Outil mentionné mais pas dans le bon format
    elif any(tool in response for tool in ['get_game_state', 'move_units', 'attack_unit', 'build_building']):
        score += 1  # Un outil MCP est mentionné
    
    # 3. Paramètres valides (2 points)
    if expected_tool == 'attack_unit':
        if 'attacker_ids' in json_response and 'target_id' in json_response:
            # Vérifier si les IDs existent
            attackers = json_response['attacker_ids']
            target = json_response['target_id']
            
            valid_attackers = any(unit['id'] in attackers if isinstance(attackers, list) else unit['id'] == attackers for unit in game_state['units'])
            valid_target = any(unit['id'] == target for unit in game_state['enemy_units'])
            
            if valid_attackers and valid_target:
                score += 2
            elif valid_attackers or valid_target:
                score += 1
                
    elif expected_tool == 'move_units':
        if 'unit_ids' in json_response and 'target_x' in json_response and 'target_y' in json_response:
            # Vérifier coordonnées valides
            x, y = json_response['target_x'], json_response['target_y']
            if 0 <= x <= 95 and 0 <= y <= 71:
                score += 2
                
    elif expected_tool == 'build_building':
        if 'building_type' in json_response and 'position_x' in json_response and 'position_y' in json_response:
            score += 2
            
    elif expected_tool == 'get_game_state':
        if not json_response or len(json_response) == 0:
            score += 2  # get_game_state n'a pas besoin de paramètres
    
    # 4. Cohérence tactique (2 points)
    if scenario['difficulty'] == 'difficile':
        # Pour les scénarios difficiles, vérifier la pertinence tactique
        if 'attack' in scenario['scenario'].lower() and 'attack' in response.lower():
            score += 1
        if 'defend' in scenario['scenario'].lower() and ('defend' in response.lower() or 'move' in response.lower()):
            score += 1
    else:
        # Pour les scénarios plus simples
        if any(word in response.lower() for word in ['game', 'state', 'move', 'attack', 'build']):
            score += 1
    
    return min(score, 10)

def run_realistic_evaluation():
    """Exécute l'évaluation réaliste complète"""
    
    print("🎮 ÉVALUATION RÉALISTE MCP POUR JEU RTS")
    print("=" * 70)
    print("Test avec contexte de jeu, scénarios réels et validation tactique")
    print("=" * 70)
    
    # 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-270M',
            'path': 'gemma-3-270m-it-qat-Q8_0.gguf'
        }
    ]
    
    results = []
    
    for model in models:
        if os.path.exists(model['path']):
            result = test_model_realistic(model['path'], model['name'])
            results.append(result)
            print("\n" + "="*50)
        else:
            print(f"❌ Modèle non trouvé: {model['path']}")
    
    # Analyse comparative réaliste
    print("\n" + "="*70)
    print("📊 ANALYSE COMPARATIVE RÉALISTE")
    print("="*70)
    
    successful_results = [r for r in results if 'error' not in r and r['avg_score'] > 0]
    
    if successful_results:
        # Classement par performance globale
        sorted_by_performance = sorted(successful_results, key=lambda x: x['avg_score'], reverse=True)
        
        print(f"\n🏆 CLASSEMENT PAR PERFORMANCE RÉELLE:")
        for i, result in enumerate(sorted_by_performance, 1):
            print(f"   {i}. {result['name']}: {result['avg_score']:.1f}/10 | {result['avg_time']:.2f}s")
        
        # Analyse par difficulté
        print(f"\n📈 PERFORMANCE PAR DIFFICULTÉ:")
        
        difficulties = ['facile', 'moyen', 'difficile']
        for diff in difficulties:
            print(f"\n🔸 Scénarios {diff}s:")
            for result in successful_results:
                avg_key = f"{diff}_avg"
                if hasattr(result, avg_key):
                    score = getattr(result, avg_key)
                    print(f"   {result['name']}: {score:.1f}/10")
        
        # Recommandations basées sur l'usage réel
        best_overall = sorted_by_performance[0]
        
        print(f"\n🎯 RECOMMANDATIONS POUR VOTRE JEU RTS:")
        
        if best_overall['avg_score'] >= 7:
            print(f"✅ {best_overall['name']} est EXCELLENT pour la production")
            print(f"   • Gère bien les scénarios complexes")
            print(f"   • Réponses tactiques cohérentes")
        elif best_overall['avg_score'] >= 5:
            print(f"👍 {best_overall['name']} est BON pour la production")
            print(f"   • Fonctionne bien pour les commandes de base")
            print(f"   • Nécessite peut-être une validation supplémentaire")
        else:
            print(f"⚠️  {best_overall['name']} nécessite des améliorations")
            print(f"   • Considérer des prompts plus structurés")
            print(f"   • Ajouter des règles de validation")
        
        # Analyse spécifique aux cas d'usage
        print(f"\n🎮 ANALYSE SPÉCIFIQUE AU JEU:")
        
        for result in successful_results:
            print(f"\n🔹 {result['name']}:")
            
            # Analyser les résultats par scénario
            scenario_scores = {}
            for scenario_result in result['results']:
                scenario_name = scenario_result['scenario']
                if scenario_name not in scenario_scores:
                    scenario_scores[scenario_name] = []
                scenario_scores[scenario_name].append(scenario_result['score'])
            
            for scenario, scores in scenario_scores.items():
                avg_score = sum(scores) / len(scores)
                print(f"   {scenario}: {avg_score:.1f}/10")
    
    # Sauvegarder les résultats réalistes
    realistic_results = {
        'evaluation_type': 'realistic_mcp_game_test',
        'game_state_sample': GAME_STATE,
        'results': results,
        'successful_models': successful_results
    }
    
    with open("realistic_mcp_evaluation.json", "w", encoding="utf-8") as f:
        json.dump(realistic_results, f, indent=2, ensure_ascii=False)
    
    print(f"\n📄 Résultats réalistes sauvegardés dans: realistic_mcp_evaluation.json")

if __name__ == "__main__":
    run_realistic_evaluation()