"""
Tests pour la vectorisation et l'indexation des référentiels.

Ce module teste:
- Le chargement du modèle d'embeddings
- La vectorisation des chunks
- La construction de l'index HNSW avec FAISS
- La vectorisation des index alphabétiques
- La génération de hash d'index pour versionnement
"""

import json
import tempfile
from pathlib import Path

import faiss
import numpy as np
import pytest

from pipeline_mco_pmsi.rag.referentiels_manager import Chunk, ReferentielsManager, VectorIndex


class TestEmbeddingsModel:
    """Tests pour le chargement du modèle d'embeddings."""

    def test_load_embeddings_model_success(self):
        """Test que le modèle d'embeddings se charge correctement."""
        manager = ReferentielsManager()
        model = manager._load_embeddings_model()
        
        assert model is not None
        # Vérifier que le modèle peut encoder du texte
        embedding = model.encode("Test médical", convert_to_numpy=True)
        assert embedding is not None
        assert len(embedding.shape) == 1  # Vecteur 1D
        assert embedding.shape[0] > 0  # Dimension > 0

    def test_embeddings_model_produces_consistent_vectors(self):
        """Test que le modèle produit des vecteurs cohérents pour le même texte."""
        manager = ReferentielsManager()
        model = manager._load_embeddings_model()
        
        text = "Gastrite aiguë avec hémorragie"
        embedding1 = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
        embedding2 = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
        
        # Les embeddings doivent être identiques (ou très proches)
        np.testing.assert_array_almost_equal(embedding1, embedding2, decimal=5)

    def test_embeddings_model_normalized(self):
        """Test que les embeddings sont normalisés (L2 norm = 1)."""
        manager = ReferentielsManager()
        model = manager._load_embeddings_model()
        
        text = "Diagnostic principal"
        embedding = model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
        
        # Calculer la norme L2
        norm = np.linalg.norm(embedding)
        
        # La norme doit être proche de 1 (normalisé)
        assert abs(norm - 1.0) < 0.01


class TestVectorization:
    """Tests pour la vectorisation des chunks."""

    def test_vectorize_single_chunk(self):
        """Test la vectorisation d'un seul chunk."""
        manager = ReferentielsManager()
        
        chunk = Chunk(
            chunk_id="test_001",
            referentiel_type="cim10",
            referentiel_version="2026",
            content="K29.7 Gastrite, sans précision",
            metadata={"chunk_type": "code_block"},
            chunk_index=0,
        )
        
        # Vectoriser le chunk
        model = manager._load_embeddings_model()
        vector = model.encode(chunk.content, convert_to_numpy=True, normalize_embeddings=True)
        
        assert vector is not None
        assert vector.shape[0] > 0
        assert isinstance(vector, np.ndarray)

    def test_vectorize_multiple_chunks(self):
        """Test la vectorisation de plusieurs chunks."""
        manager = ReferentielsManager()
        
        chunks = [
            Chunk(
                chunk_id=f"test_{i:03d}",
                referentiel_type="cim10",
                referentiel_version="2026",
                content=f"Code test {i}",
                metadata={"chunk_type": "code_block"},
                chunk_index=i,
            )
            for i in range(10)
        ]
        
        model = manager._load_embeddings_model()
        vectors = [
            model.encode(chunk.content, convert_to_numpy=True, normalize_embeddings=True)
            for chunk in chunks
        ]
        
        assert len(vectors) == 10
        # Tous les vecteurs doivent avoir la même dimension
        dimensions = [v.shape[0] for v in vectors]
        assert len(set(dimensions)) == 1  # Toutes les dimensions sont identiques


class TestBuildIndex:
    """Tests pour la construction de l'index HNSW."""

    def test_build_index_success(self, tmp_path):
        """Test la construction réussie d'un index HNSW."""
        manager = ReferentielsManager(data_dir=tmp_path)
        
        # Créer des chunks de test
        chunks = [
            Chunk(
                chunk_id=f"cim10_2026_{i}",
                referentiel_type="cim10",
                referentiel_version="2026",
                content=f"K29.{i} Gastrite type {i}",
                metadata={"chunk_type": "code_block"},
                chunk_index=i,
            )
            for i in range(20)
        ]
        
        # Construire l'index
        vector_index = manager.build_index(chunks)
        
        # Vérifications
        assert isinstance(vector_index, VectorIndex)
        assert vector_index.index_hash is not None
        assert len(vector_index.index_hash) == 64  # SHA-256 hex
        assert vector_index.dimension > 0
        assert vector_index.num_vectors == 20
        assert vector_index.index_type == "HNSW"
        assert vector_index.created_at is not None

    def test_build_index_saves_to_disk(self, tmp_path):
        """Test que l'index est sauvegardé sur disque."""
        manager = ReferentielsManager(data_dir=tmp_path)
        
        chunks = [
            Chunk(
                chunk_id=f"cim10_2026_{i}",
                referentiel_type="cim10",
                referentiel_version="2026",
                content=f"Code {i}",
                metadata={},
                chunk_index=i,
            )
            for i in range(10)
        ]
        
        manager.build_index(chunks)
        
        # Vérifier que les fichiers sont créés
        index_file = tmp_path / "cim10_2026_index.faiss"
        chunks_file = tmp_path / "cim10_2026_chunks.json"
        
        assert index_file.exists()
        assert chunks_file.exists()
        
        # Vérifier que l'index peut être rechargé
        loaded_index = faiss.read_index(str(index_file))
        assert loaded_index.ntotal == 10

    def test_build_index_chunks_json_valid(self, tmp_path):
        """Test que le fichier JSON des chunks est valide."""
        manager = ReferentielsManager(data_dir=tmp_path)
        
        chunks = [
            Chunk(
                chunk_id=f"ccam_2025_{i}",
                referentiel_type="ccam",
                referentiel_version="2025",
                content=f"YYYY00{i} Acte {i}",
                metadata={"section": "Section A"},
                chunk_index=i,
            )
            for i in range(5)
        ]
        
        manager.build_index(chunks)
        
        # Charger et vérifier le JSON
        chunks_file = tmp_path / "ccam_2025_chunks.json"
        with open(chunks_file, "r", encoding="utf-8") as f:
            loaded_chunks = json.load(f)
        
        assert len(loaded_chunks) == 5
        assert loaded_chunks[0]["chunk_id"] == "ccam_2025_0"
        assert loaded_chunks[0]["referentiel_type"] == "ccam"
        assert loaded_chunks[0]["content"] == "YYYY000 Acte 0"

    def test_build_index_empty_chunks_raises_error(self):
        """Test qu'une liste vide de chunks lève une erreur."""
        manager = ReferentielsManager()
        
        with pytest.raises(ValueError, match="ne peut pas être vide"):
            manager.build_index([])

    def test_build_index_hash_consistency(self, tmp_path):
        """Test que le hash de l'index est cohérent."""
        manager = ReferentielsManager(data_dir=tmp_path)
        
        chunks = [
            Chunk(
                chunk_id=f"test_{i}",
                referentiel_type="cim10",
                referentiel_version="2026",
                content=f"Content {i}",
                metadata={},
                chunk_index=i,
            )
            for i in range(10)
        ]
        
        # Construire l'index deux fois
        index1 = manager.build_index(chunks)
        
        # Le hash devrait être le même (même paramètres, même nombre de vecteurs)
        # Note: En pratique, le hash dépend des paramètres, pas du contenu exact
        assert len(index1.index_hash) == 64


class TestAlphabeticalIndexVectorization:
    """Tests pour la vectorisation des index alphabétiques."""

    def test_vectorize_alphabetical_index_cim10(self):
        """Test la vectorisation d'un index alphabétique CIM-10."""
        manager = ReferentielsManager()
        
        # Simuler un index alphabétique CIM-10
        alpha_text = """A
Abcès - voir A00.1
Abdomen aigu - voir R10.0

B
Bronchite - voir J20.9
Bronchopneumonie - voir J18.0

C
Cataracte - voir H26.9
Cholécystite - voir K81.9
"""
        
        chunks = manager.vectorize_alphabetical_indexes(
            alpha_text, "cim10", "2026"
        )
        
        assert len(chunks) > 0
        # Vérifier que les chunks ont le bon type
        for chunk in chunks:
            assert chunk.referentiel_type == "cim10"
            assert chunk.referentiel_version == "2026"
            assert chunk.metadata["chunk_type"] == "alphabetical_index"
            assert "letter" in chunk.metadata

    def test_vectorize_alphabetical_index_ccam(self):
        """Test la vectorisation d'un index alphabétique CCAM."""
        manager = ReferentielsManager()
        
        # Simuler un index alphabétique CCAM
        alpha_text = """A
Ablation - YYYY001
Amputation - YYYY002

B
Biopsie - ZZZZ001
Bronchoscopie - ZZZZ002
"""
        
        chunks = manager.vectorize_alphabetical_indexes(
            alpha_text, "ccam", "2025"
        )
        
        assert len(chunks) > 0
        for chunk in chunks:
            assert chunk.referentiel_type == "ccam"
            assert chunk.referentiel_version == "2025"
            assert chunk.metadata["chunk_type"] == "alphabetical_index"

    def test_alphabetical_index_extracts_codes(self):
        """Test que les codes sont extraits dans les métadonnées."""
        manager = ReferentielsManager()
        
        alpha_text = """Gastrite - voir K29.7
Gastro-entérite - voir A09
Glaucome - voir H40.9
"""
        
        chunks = manager.vectorize_alphabetical_indexes(
            alpha_text, "cim10", "2026"
        )
        
        # Au moins un chunk devrait contenir des codes dans les métadonnées
        codes_found = False
        for chunk in chunks:
            if chunk.metadata.get("codes"):
                codes_found = True
                # Vérifier que les codes sont au format CIM-10
                codes = chunk.metadata["codes"].split(",")
                for code in codes:
                    assert len(code) >= 3  # Au moins A00
        
        assert codes_found

    def test_alphabetical_index_chunk_size(self):
        """Test que les chunks d'index alphabétique respectent la taille cible."""
        manager = ReferentielsManager()
        
        # Créer un long index alphabétique
        entries = []
        for i in range(100):
            entries.append(f"Terme médical {i} - voir K{i:02d}.{i % 10}")
        
        alpha_text = "\n".join(entries)
        
        chunks = manager.vectorize_alphabetical_indexes(
            alpha_text, "cim10", "2026"
        )
        
        # Vérifier que les chunks ne sont pas trop grands
        for chunk in chunks:
            assert len(chunk.content) <= 2048  # max_chunk_size


class TestIntegrationVectorizationIndexation:
    """Tests d'intégration pour le workflow complet."""

    def test_full_workflow_import_chunk_vectorize_index(self, tmp_path):
        """Test le workflow complet: import → chunk → vectorize → index."""
        # Créer un fichier PDF de test
        pdf_path = tmp_path / "test_ref.pdf"
        
        # Pour ce test, on va créer un fichier texte au lieu d'un PDF
        # car créer un PDF valide est complexe
        text_content = """CHAPITRE I - Maladies infectieuses

A00-A09 Maladies intestinales infectieuses

A00 Choléra
A00.0 Choléra dû à Vibrio cholerae 01, biovar cholerae
A00.1 Choléra dû à Vibrio cholerae 01, biovar El Tor
A00.9 Choléra, sans précision

Inclus: infection à Vibrio cholerae
Exclut: intoxication alimentaire (A05.-)
"""
        
        # Sauvegarder le texte directement (simuler l'extraction PDF)
        manager = ReferentielsManager(data_dir=tmp_path)
        text_file = tmp_path / "cim10_2026_text.txt"
        with open(text_file, "w", encoding="utf-8") as f:
            f.write(text_content)
        
        # Créer une version de référentiel manuellement
        from pipeline_mco_pmsi.models.metadata import ReferentielVersion
        from datetime import datetime
        import hashlib
        
        ref_version = ReferentielVersion(
            type="cim10",
            version="2026",
            import_date=datetime.now(),
            file_hash=hashlib.sha256(text_content.encode()).hexdigest(),
            chunk_count=0,
            index_hash="0" * 64,
        )
        
        # Chunker
        chunks = manager.chunk_referentiel(ref_version)
        assert len(chunks) > 0
        
        # Construire l'index
        vector_index = manager.build_index(chunks)
        assert vector_index.num_vectors == len(chunks)
        
        # Vérifier que les fichiers sont créés
        assert (tmp_path / "cim10_2026_index.faiss").exists()
        assert (tmp_path / "cim10_2026_chunks.json").exists()


class TestIndexSearch:
    """Tests pour la recherche dans l'index vectoriel."""

    def test_index_search_basic(self, tmp_path):
        """Test une recherche basique dans l'index."""
        manager = ReferentielsManager(data_dir=tmp_path)
        
        # Créer des chunks avec du contenu médical
        chunks = [
            Chunk(
                chunk_id=f"cim10_2026_{i}",
                referentiel_type="cim10",
                referentiel_version="2026",
                content=content,
                metadata={},
                chunk_index=i,
            )
            for i, content in enumerate([
                "K29.7 Gastrite, sans précision",
                "K29.0 Gastrite hémorragique aiguë",
                "J18.0 Bronchopneumonie, sans précision",
                "I10 Hypertension essentielle",
                "E11 Diabète sucré non insulino-dépendant",
            ])
        ]
        
        # Construire l'index
        manager.build_index(chunks)
        
        # Charger l'index
        index_file = tmp_path / "cim10_2026_index.faiss"
        index = faiss.read_index(str(index_file))
        
        # Effectuer une recherche
        model = manager._load_embeddings_model()
        query = "gastrite"
        query_vector = model.encode(query, convert_to_numpy=True, normalize_embeddings=True)
        query_vector = query_vector.reshape(1, -1).astype(np.float32)
        
        # Rechercher les 2 plus proches voisins
        distances, indices = index.search(query_vector, 2)
        
        # Les résultats devraient inclure les chunks sur la gastrite
        assert len(indices[0]) == 2
        # Les indices devraient être 0 ou 1 (les deux chunks sur la gastrite)
        assert indices[0][0] in [0, 1]
        assert indices[0][1] in [0, 1]

    def test_index_search_similarity_scores(self, tmp_path):
        """Test que les scores de similarité sont cohérents."""
        manager = ReferentielsManager(data_dir=tmp_path)
        
        chunks = [
            Chunk(
                chunk_id=f"test_{i}",
                referentiel_type="cim10",
                referentiel_version="2026",
                content=content,
                metadata={},
                chunk_index=i,
            )
            for i, content in enumerate([
                "Gastrite aiguë avec hémorragie",
                "Gastrite chronique",
                "Pneumonie bactérienne",
            ])
        ]
        
        manager.build_index(chunks)
        
        # Charger l'index
        index_file = tmp_path / "cim10_2026_index.faiss"
        index = faiss.read_index(str(index_file))
        
        # Rechercher avec une requête proche du premier chunk
        model = manager._load_embeddings_model()
        query = "gastrite hémorragique"
        query_vector = model.encode(query, convert_to_numpy=True, normalize_embeddings=True)
        query_vector = query_vector.reshape(1, -1).astype(np.float32)
        
        distances, indices = index.search(query_vector, 3)
        
        # Le premier résultat devrait être le plus proche (distance la plus petite)
        # Avec des embeddings normalisés, la distance est 1 - cosine_similarity
        # Donc une distance plus petite = plus similaire
        assert distances[0][0] <= distances[0][1]
        assert distances[0][1] <= distances[0][2]


if __name__ == "__main__":
    pytest.main([__file__, "-v"])