#!/usr/bin/env python3
"""
Test d'intégration du nouveau système d'embeddings dans l'Orchestrator.
"""

import sys
import os
import time
import numpy as np
from PIL import Image

# Ajouter le chemin du projet
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'geniusia2'))

from core.orchestrator import Orchestrator
from core.learning_manager import LearningManager
from core.utils.vision_utils import VisionUtils
from core.llm_manager import LLMManager
from core.logger import Logger
from core.embeddings_manager import EmbeddingsManager
from core.config import get_config, ensure_directories

def test_embedding_system_initialization():
    """Test 1: Vérifier que le système d'embeddings s'initialise correctement."""
    print("\n" + "="*60)
    print("TEST 1: Initialisation du système d'embeddings")
    print("="*60)
    
    try:
        # Créer les composants
        ensure_directories()
        logger = Logger()
        embeddings_manager = EmbeddingsManager()
        learning_manager = LearningManager(embeddings_manager, logger, get_config())
        vision_utils = VisionUtils()
        llm_manager = LLMManager(logger=logger, fallback_to_vision=True)
        
        # Créer l'orchestrator
        orchestrator = Orchestrator(
            learning_manager=learning_manager,
            vision_utils=vision_utils,
            llm_manager=llm_manager,
            logger=logger,
            gui=None
        )
        
        # Vérifier que le nouveau système est initialisé
        assert hasattr(orchestrator, 'new_embedding_manager'), "❌ new_embedding_manager non initialisé"
        assert hasattr(orchestrator, 'faiss_index'), "❌ faiss_index non initialisé"
        assert hasattr(orchestrator, 'fine_tuner'), "❌ fine_tuner non initialisé"
        
        print("✅ Système d'embeddings initialisé")
        print(f"   DEBUG: faiss_index = {orchestrator.faiss_index}")
        print(f"   DEBUG: type = {type(orchestrator.faiss_index)}")
        
        # Vérifier les composants
        if orchestrator.new_embedding_manager:
            print(f"   ✓ EmbeddingManager: {orchestrator.new_embedding_manager.get_model_name()}")
            print(f"   ✓ Dimension: {orchestrator.new_embedding_manager.get_dimension()}")
        else:
            print("   ⚠️  EmbeddingManager non disponible")
        
        if orchestrator.faiss_index is not None:
            stats = orchestrator.faiss_index.get_stats()
            print(f"   ✓ FAISS Index: {stats['num_embeddings']} embeddings")
        else:
            print("   ⚠️  FAISS Index non disponible")
        
        if orchestrator.fine_tuner:
            stats = orchestrator.fine_tuner.get_stats()
            print(f"   ✓ Fine-tuner: {stats['total_examples']} exemples")
        else:
            print("   ⚠️  Fine-tuner désactivé")
        
        return orchestrator
    
    except Exception as e:
        print(f"❌ Erreur: {e}")
        import traceback
        traceback.print_exc()
        return None


def test_workflow_indexing(orchestrator):
    """Test 2: Vérifier l'indexation d'un workflow dans FAISS."""
    print("\n" + "="*60)
    print("TEST 2: Indexation d'un workflow dans FAISS")
    print("="*60)
    
    if not orchestrator or orchestrator.faiss_index is None:
        print("⚠️  Test ignoré (FAISS non disponible)")
        return
    
    try:
        # Créer un workflow de test avec des screenshots
        test_workflow = {
            "workflow_id": "test_workflow_001",
            "name": "Test Workflow",
            "screenshots": [
                Image.new('RGB', (100, 100), color='red'),
                Image.new('RGB', (100, 100), color='green'),
                Image.new('RGB', (100, 100), color='blue')
            ]
        }
        
        # Indexer le workflow
        initial_count = orchestrator.faiss_index.get_stats()['num_embeddings']
        orchestrator._index_workflow_in_faiss(test_workflow)
        final_count = orchestrator.faiss_index.get_stats()['num_embeddings']
        
        # Vérifier que les embeddings ont été ajoutés
        added = final_count - initial_count
        assert added == 3, f"❌ Attendu 3 embeddings, obtenu {added}"
        
        print(f"✅ Workflow indexé: {added} embeddings ajoutés")
        print(f"   ✓ Total dans l'index: {final_count}")
        
    except Exception as e:
        print(f"❌ Erreur: {e}")
        import traceback
        traceback.print_exc()


def test_positive_example_collection(orchestrator):
    """Test 3: Vérifier la collection d'exemples positifs."""
    print("\n" + "="*60)
    print("TEST 3: Collection d'exemples positifs")
    print("="*60)
    
    if not orchestrator or not orchestrator.fine_tuner:
        print("⚠️  Test ignoré (Fine-tuner non disponible)")
        return
    
    try:
        # Créer une suggestion de test
        test_suggestion = {
            "task_id": "test_task_001",
            "screenshot": Image.new('RGB', (100, 100), color='yellow'),
            "confidence": 0.85
        }
        
        # Ajouter comme exemple positif
        initial_count = orchestrator.fine_tuner.get_stats()['total_examples']
        orchestrator._add_positive_example_for_finetuning(test_suggestion)
        final_count = orchestrator.fine_tuner.get_stats()['total_examples']
        
        # Vérifier que l'exemple a été ajouté
        added = final_count - initial_count
        assert added == 1, f"❌ Attendu 1 exemple, obtenu {added}"
        
        print(f"✅ Exemple positif ajouté")
        print(f"   ✓ Total d'exemples: {final_count}")
        
    except Exception as e:
        print(f"❌ Erreur: {e}")
        import traceback
        traceback.print_exc()


def test_negative_example_collection(orchestrator):
    """Test 4: Vérifier la collection d'exemples négatifs."""
    print("\n" + "="*60)
    print("TEST 4: Collection d'exemples négatifs")
    print("="*60)
    
    if not orchestrator or not orchestrator.fine_tuner:
        print("⚠️  Test ignoré (Fine-tuner non disponible)")
        return
    
    try:
        # Créer une suggestion de test
        test_suggestion = {
            "task_id": "test_task_002",
            "screenshot": Image.new('RGB', (100, 100), color='purple'),
            "confidence": 0.45
        }
        
        # Ajouter comme exemple négatif
        initial_count = orchestrator.fine_tuner.get_stats()['total_examples']
        orchestrator._add_negative_example_for_finetuning(test_suggestion)
        final_count = orchestrator.fine_tuner.get_stats()['total_examples']
        
        # Vérifier que l'exemple a été ajouté
        added = final_count - initial_count
        assert added == 1, f"❌ Attendu 1 exemple, obtenu {added}"
        
        print(f"✅ Exemple négatif ajouté")
        print(f"   ✓ Total d'exemples: {final_count}")
        
    except Exception as e:
        print(f"❌ Erreur: {e}")
        import traceback
        traceback.print_exc()


def test_cache_functionality(orchestrator):
    """Test 5: Vérifier le fonctionnement du cache."""
    print("\n" + "="*60)
    print("TEST 5: Fonctionnement du cache")
    print("="*60)
    
    if not orchestrator or not orchestrator.new_embedding_manager:
        print("⚠️  Test ignoré (EmbeddingManager non disponible)")
        return
    
    try:
        # Créer une image de test
        test_image = Image.new('RGB', (100, 100), color='cyan')
        
        # Premier embedding (cache miss)
        start = time.time()
        embedding1 = orchestrator.new_embedding_manager.embed(test_image)
        time1 = (time.time() - start) * 1000
        
        # Deuxième embedding (cache hit)
        start = time.time()
        embedding2 = orchestrator.new_embedding_manager.embed(test_image)
        time2 = (time.time() - start) * 1000
        
        # Vérifier que les embeddings sont identiques
        assert np.allclose(embedding1, embedding2), "❌ Embeddings différents"
        
        # Vérifier que le cache est plus rapide
        speedup = time1 / time2 if time2 > 0 else float('inf')
        
        print(f"✅ Cache fonctionnel")
        print(f"   ✓ Premier appel: {time1:.2f}ms")
        print(f"   ✓ Cache hit: {time2:.2f}ms")
        print(f"   ✓ Speedup: {speedup:.1f}x")
        
        # Afficher les stats du cache
        stats = orchestrator.new_embedding_manager.get_stats()
        print(f"   ✓ Hit rate: {stats['cache_hit_rate']:.1%}")
        
    except Exception as e:
        print(f"❌ Erreur: {e}")
        import traceback
        traceback.print_exc()


def test_save_and_load(orchestrator):
    """Test 6: Vérifier la sauvegarde et le chargement."""
    print("\n" + "="*60)
    print("TEST 6: Sauvegarde et chargement")
    print("="*60)
    
    if not orchestrator:
        print("⚠️  Test ignoré (Orchestrator non disponible)")
        return
    
    try:
        # Sauvegarder l'état
        orchestrator._save_embedding_system_on_shutdown()
        
        print("✅ Sauvegarde effectuée")
        
        # Vérifier que les fichiers existent
        if orchestrator.faiss_index:
            index_path = orchestrator._faiss_index_path
            if os.path.exists(f"{index_path}.index"):
                print(f"   ✓ Index FAISS sauvegardé: {index_path}.index")
            else:
                print(f"   ⚠️  Index FAISS non trouvé: {index_path}.index")
        
        if orchestrator.fine_tuner:
            checkpoint_name = orchestrator._fine_tuner_checkpoint_name
            checkpoint_path = os.path.join("data/fine_tuning", f"{checkpoint_name}.pkl")
            if os.path.exists(checkpoint_path):
                print(f"   ✓ Checkpoint fine-tuner sauvegardé: {checkpoint_path}")
            else:
                print(f"   ⚠️  Checkpoint non trouvé: {checkpoint_path}")
        
    except Exception as e:
        print(f"❌ Erreur: {e}")
        import traceback
        traceback.print_exc()


def main():
    """Exécuter tous les tests."""
    print("\n" + "="*60)
    print("TEST D'INTÉGRATION DU SYSTÈME D'EMBEDDINGS")
    print("="*60)
    
    # Test 1: Initialisation
    orchestrator = test_embedding_system_initialization()
    
    if not orchestrator:
        print("\n❌ Échec de l'initialisation, arrêt des tests")
        return
    
    # Test 2: Indexation workflow
    test_workflow_indexing(orchestrator)
    
    # Test 3: Exemples positifs
    test_positive_example_collection(orchestrator)
    
    # Test 4: Exemples négatifs
    test_negative_example_collection(orchestrator)
    
    # Test 5: Cache
    test_cache_functionality(orchestrator)
    
    # Test 6: Sauvegarde
    test_save_and_load(orchestrator)
    
    # Résumé
    print("\n" + "="*60)
    print("RÉSUMÉ")
    print("="*60)
    print("✅ Tous les tests terminés")
    print("\nLe système d'embeddings est maintenant intégré dans l'Orchestrator.")
    print("Les workflows seront automatiquement indexés dans FAISS.")
    print("Les suggestions acceptées/rejetées alimenteront le fine-tuning.")


if __name__ == "__main__":
    main()