t2a_v2/tests/test_rag.py

"""Tests pour le RAG CIM-10 (modèles, chunking, intégration)."""

from __future__ import annotations

from pathlib import Path
from unittest.mock import patch, MagicMock

import pytest

from src.config import RAGSource, Diagnostic, ActeCCAM, DossierMedical, CIM10_PDF, GUIDE_METHODO_PDF, CCAM_PDF
from src.medical.ollama_cache import OllamaCache


class TestRAGSource:
    def test_create_minimal(self):
        src = RAGSource(document="cim10")
        assert src.document == "cim10"
        assert src.page is None
        assert src.code is None
        assert src.extrait is None

    def test_create_full(self):
        src = RAGSource(
            document="guide_methodo",
            page=42,
            code="K85",
            extrait="Pancréatite aiguë biliaire...",
        )
        assert src.document == "guide_methodo"
        assert src.page == 42
        assert src.code == "K85"
        assert src.extrait == "Pancréatite aiguë biliaire..."

    def test_serialization(self):
        src = RAGSource(document="ccam", page=1, code="HMFC004")
        data = src.model_dump(exclude_none=True)
        assert data == {"document": "ccam", "page": 1, "code": "HMFC004"}


class TestDiagnosticExtended:
    def test_backward_compatible(self):
        """Les nouveaux champs sont optionnels — rétrocompatible."""
        d = Diagnostic(texte="Pancréatite aiguë", cim10_suggestion="K85.9")
        assert d.texte == "Pancréatite aiguë"
        assert d.cim10_suggestion == "K85.9"
        assert d.cim10_confidence is None
        assert d.justification is None
        assert d.raisonnement is None
        assert d.sources_rag == []

    def test_with_rag_fields(self):
        d = Diagnostic(
            texte="Lithiase cholédoque",
            cim10_suggestion="K80.5",
            cim10_confidence="high",
            justification="Code K80.5 correspond à la lithiase du cholédoque",
            raisonnement="1. ANALYSE CLINIQUE : La lithiase cholédoque est...",
            sources_rag=[
                RAGSource(document="cim10", page=480, code="K80"),
            ],
        )
        assert d.cim10_confidence == "high"
        assert d.justification is not None
        assert d.raisonnement is not None
        assert d.raisonnement.startswith("1. ANALYSE CLINIQUE")
        assert len(d.sources_rag) == 1
        assert d.sources_rag[0].code == "K80"

    def test_serialization_exclude_none(self):
        """Vérifier que le JSON n'inclut pas les champs None."""
        d = Diagnostic(texte="Test", cim10_suggestion="K85.9")
        data = d.model_dump(exclude_none=True)
        assert "cim10_confidence" not in data
        assert "justification" not in data
        assert "raisonnement" not in data
        assert "sources_rag" in data  # list vide incluse

    def test_dossier_with_extended_diagnostic(self):
        """Un DossierMedical avec des diagnostics enrichis par le RAG."""
        dossier = DossierMedical(
            diagnostic_principal=Diagnostic(
                texte="Pancréatite aiguë biliaire",
                cim10_suggestion="K85.1",
                cim10_confidence="high",
                justification="Confirmé par CIM-10 FR 2026",
                raisonnement="Le DP K85.1 est le code le plus spécifique...",
                sources_rag=[
                    RAGSource(document="cim10", page=496, code="K85"),
                    RAGSource(document="guide_methodo", page=30),
                ],
            ),
        )
        assert dossier.diagnostic_principal.cim10_confidence == "high"
        assert dossier.diagnostic_principal.raisonnement is not None
        assert len(dossier.diagnostic_principal.sources_rag) == 2


class TestExtractMedicalInfoRAGFlag:
    def test_use_rag_false_no_change(self):
        """use_rag=False ne modifie pas le comportement existant."""
        from src.medical.cim10_extractor import extract_medical_info

        parsed = {
            "type": "crh",
            "patient": {"sexe": "M"},
            "sejour": {},
            "diagnostics": [],
        }
        text = "Pancréatite aiguë biliaire.\nTTT de sortie :\nParacétamol\n\nDevenir : retour."

        dossier = extract_medical_info(parsed, text, use_rag=False)
        assert dossier.diagnostic_principal is not None
        assert dossier.diagnostic_principal.cim10_suggestion == "K85.1"
        # Pas de sources RAG
        assert dossier.diagnostic_principal.sources_rag == []
        assert dossier.diagnostic_principal.justification is None

    def test_use_rag_true_calls_enrich(self):
        """use_rag=True appelle _enrich_dp_only et _enrich_das_and_actes (mockés)."""
        from src.medical.cim10_extractor import extract_medical_info

        parsed = {
            "type": "crh",
            "patient": {"sexe": "M"},
            "sejour": {},
            "diagnostics": [],
        }
        text = "Pancréatite aiguë biliaire.\nTTT de sortie :\nParacétamol\n\nDevenir : retour."

        with patch("src.medical.cim10_extractor._enrich_dp_only") as mock_dp, \
             patch("src.medical.cim10_extractor._enrich_das_and_actes") as mock_das, \
             patch("src.medical.cim10_extractor._extract_das_llm"), \
             patch("src.medical.cim10_extractor._validate_justifications"):
            dossier = extract_medical_info(parsed, text, use_rag=True)
            mock_dp.assert_called_once_with(dossier)
            mock_das.assert_called_once_with(dossier)

    def test_use_rag_default_false(self):
        """Par défaut, use_rag=False."""
        from src.medical.cim10_extractor import extract_medical_info

        parsed = {
            "type": "crh",
            "patient": {"sexe": "M"},
            "sejour": {},
            "diagnostics": [],
        }
        text = "Test simple."

        with patch("src.medical.cim10_extractor._enrich_dp_only") as mock_dp, \
             patch("src.medical.cim10_extractor._enrich_das_and_actes") as mock_das:
            extract_medical_info(parsed, text)
            mock_dp.assert_not_called()
            mock_das.assert_not_called()


class TestChunkingCIM10:
    @pytest.mark.skipif(
        not CIM10_PDF.exists(),
        reason=f"PDF CIM-10 non trouvé : {CIM10_PDF}",
    )
    def test_chunks_contain_known_codes(self):
        from src.medical.rag_index import _chunk_cim10

        chunks = _chunk_cim10(CIM10_PDF)
        assert len(chunks) > 100, f"Trop peu de chunks : {len(chunks)}"

        codes = {c.code for c in chunks if c.code}
        # Codes parents 3-char
        assert "K85" in codes, "K85 (pancréatite) non trouvé"
        assert "K80" in codes, "K80 (lithiase biliaire) non trouvé"
        assert "E66" in codes, "E66 (obésité) non trouvé"

    @pytest.mark.skipif(
        not CIM10_PDF.exists(),
        reason=f"PDF CIM-10 non trouvé : {CIM10_PDF}",
    )
    def test_double_chunking_subcodes(self):
        """Le double chunking produit des chunks sous-codes (X99.9) en plus des parents."""
        from src.medical.rag_index import _chunk_cim10

        chunks = _chunk_cim10(CIM10_PDF)
        codes = {c.code for c in chunks if c.code}

        # Il doit y avoir des sous-codes (avec un point)
        subcodes = {c for c in codes if "." in c}
        assert len(subcodes) > 100, f"Trop peu de sous-codes : {len(subcodes)}"

        # Le nombre total de chunks doit être significativement plus grand
        # qu'un chunking simple par code 3-char
        parent_codes = {c for c in codes if "." not in c}
        assert len(chunks) > len(parent_codes) * 2, \
            f"Double chunking inefficace : {len(chunks)} chunks pour {len(parent_codes)} codes parents"

    @pytest.mark.skipif(
        not CIM10_PDF.exists(),
        reason=f"PDF CIM-10 non trouvé : {CIM10_PDF}",
    )
    def test_chunk_content(self):
        from src.medical.rag_index import _chunk_cim10

        chunks = _chunk_cim10(CIM10_PDF)
        k85_chunks = [c for c in chunks if c.code and c.code.startswith("K85")]
        assert len(k85_chunks) >= 2, "Il devrait y avoir au moins un chunk parent K85 + des sous-codes"
        texts_lower = " ".join(c.text.lower() for c in k85_chunks)
        assert "pancréatite" in texts_lower or "pancreatite" in texts_lower


class TestChunkingGuideMethodo:
    @pytest.mark.skipif(
        not GUIDE_METHODO_PDF.exists(),
        reason=f"PDF Guide Métho non trouvé : {GUIDE_METHODO_PDF}",
    )
    def test_chunks_extracted(self):
        from src.medical.rag_index import _chunk_guide_methodo

        chunks = _chunk_guide_methodo(GUIDE_METHODO_PDF)
        assert len(chunks) >= 10, f"Trop peu de chunks : {len(chunks)}"
        assert all(c.document == "guide_methodo" for c in chunks)


class TestChunkingCCAM:
    @pytest.mark.skipif(
        not CCAM_PDF.exists(),
        reason=f"PDF CCAM non trouvé : {CCAM_PDF}",
    )
    def test_chunks_extracted(self):
        from src.medical.rag_index import _chunk_ccam

        chunks = _chunk_ccam(CCAM_PDF)
        assert len(chunks) >= 1, f"Aucun chunk CCAM extrait"
        assert all(c.document == "ccam" for c in chunks)


class TestParseOllamaResponse:
    """Tests pour _parse_ollama_response en mode JSON structuré."""

    def test_parse_structured_json(self):
        """Le mode JSON retourne un objet avec champs de raisonnement."""
        from src.medical.rag_search import _parse_ollama_response
        import json

        raw = json.dumps({
            "analyse_clinique": "La pancréatite aiguë biliaire est une inflammation...",
            "codes_candidats": "K85.0, K85.1, K85.9",
            "discrimination": "K85.1 est spécifique à l'origine biliaire",
            "regle_pmsi": "Conforme pour un DP",
            "code": "K85.1",
            "confidence": "high",
            "justification": "Pancréatite aiguë d'origine biliaire",
        })

        result = _parse_ollama_response(raw)
        assert result is not None
        assert result["code"] == "K85.1"
        assert result["confidence"] == "high"
        assert result["justification"] == "Pancréatite aiguë d'origine biliaire"
        assert "raisonnement" in result
        assert "ANALYSE CLINIQUE" in result["raisonnement"]
        assert "CODES CANDIDATS" in result["raisonnement"]
        # Les champs de raisonnement sont retirés du dict
        assert "analyse_clinique" not in result
        assert "codes_candidats" not in result

    def test_parse_minimal_json(self):
        """JSON sans champs de raisonnement — pas de clé raisonnement."""
        from src.medical.rag_search import _parse_ollama_response
        import json

        raw = json.dumps({
            "code": "E66.0",
            "confidence": "medium",
            "justification": "Obésité due à un excès calorique",
        })

        result = _parse_ollama_response(raw)
        assert result is not None
        assert result["code"] == "E66.0"
        assert result["confidence"] == "medium"
        assert "raisonnement" not in result

    def test_parse_empty_response(self):
        from src.medical.rag_search import _parse_ollama_response

        result = _parse_ollama_response("")
        assert result is None

    def test_parse_no_json(self):
        from src.medical.rag_search import _parse_ollama_response

        result = _parse_ollama_response("Réponse sans aucun JSON valide.")
        assert result is None

    def test_parse_invalid_json(self):
        from src.medical.rag_search import _parse_ollama_response

        raw = """{code: K85.1, invalid json}"""
        result = _parse_ollama_response(raw)
        assert result is None

    def test_parse_partial_reasoning_fields(self):
        """Seuls certains champs de raisonnement sont présents."""
        from src.medical.rag_search import _parse_ollama_response
        import json

        raw = json.dumps({
            "analyse_clinique": "Diagnostic clair",
            "code": "K85.1",
            "confidence": "high",
            "justification": "Biliaire confirmé",
        })

        result = _parse_ollama_response(raw)
        assert result is not None
        assert result["code"] == "K85.1"
        assert "raisonnement" in result
        assert "ANALYSE CLINIQUE" in result["raisonnement"]


class TestBuildPrompt:
    """Tests pour le nouveau _build_prompt avec raisonnement structuré."""

    def test_prompt_contains_diagnostic(self):
        from src.medical.rag_search import _build_prompt

        sources = [{"document": "cim10", "code": "K85", "page": 1, "extrait": "K85 Pancréatite"}]
        contexte = {"sexe": "F", "age": 43}
        prompt = _build_prompt("Pancréatite aiguë biliaire", sources, contexte, est_dp=True)

        assert "Pancréatite aiguë biliaire" in prompt
        assert "DP (diagnostic principal)" in prompt
        assert "analyse_clinique" in prompt
        assert "objet JSON" in prompt

    def test_prompt_das_type(self):
        from src.medical.rag_search import _build_prompt

        sources = [{"document": "cim10", "code": "E66", "page": 1, "extrait": "E66 Obésité"}]
        contexte = {"sexe": "F", "age": 43}
        prompt = _build_prompt("Obésité", sources, contexte, est_dp=False)

        assert "DAS (diagnostic associé significatif)" in prompt

    def test_prompt_enriched_context(self):
        from src.medical.rag_search import _build_prompt

        sources = [{"document": "cim10", "code": "K85", "page": 1, "extrait": "K85"}]
        contexte = {
            "sexe": "F",
            "age": 43,
            "imc": 34.4,
            "duree_sejour": 6,
            "antecedents": ["HTA", "diabète type 2"],
            "biologie_cle": [("Lipasémie", "850", True), ("CRP", "45", True)],
            "imagerie": [("TDM abdominal", "pancréatite stade C Balthazar")],
            "complications": ["éruption cutanée"],
            "dp_texte": "Pancréatite aiguë biliaire",
        }
        prompt = _build_prompt("Éruption cutanée", sources, contexte, est_dp=False)

        assert "IMC 34.4" in prompt
        assert "6 jours" in prompt
        assert "HTA" in prompt
        assert "Lipasémie" in prompt
        assert "TDM abdominal" in prompt
        assert "éruption cutanée" in prompt
        assert "Pancréatite aiguë biliaire" in prompt


class TestSearchSimilar:
    """Tests pour search_similar avec score minimum et priorisation CIM-10."""

    def test_filters_low_scores(self):
        """Les résultats avec score < 0.3 sont éliminés."""
        from src.medical.rag_search import search_similar
        import numpy as np

        mock_metadata = [
            {"document": "cim10", "code": "K85", "page": 1, "extrait": "K85"},
            {"document": "cim10", "code": "K86", "page": 2, "extrait": "K86"},
        ]

        mock_index = MagicMock()
        mock_index.ntotal = 2
        # Premier résultat score=0.9 (bon), second score=0.1 (sous le seuil)
        mock_index.search.return_value = (
            np.array([[0.9, 0.1]], dtype=np.float32),
            np.array([[0, 1]], dtype=np.int64),
        )

        with patch("src.medical.rag_index.get_index", return_value=(mock_index, mock_metadata)), \
             patch("src.medical.rag_search._get_embed_model") as mock_model:
            mock_model.return_value.encode.return_value = np.array([[0.1] * 768], dtype=np.float32)
            results = search_similar("pancréatite")

        assert len(results) == 1
        assert results[0]["code"] == "K85"

    def test_prioritizes_cim10(self):
        """Les sources CIM-10 sont priorisées (au moins 6 sur 10)."""
        from src.medical.rag_search import search_similar
        import numpy as np

        # 8 sources CIM-10 + 8 sources guide_methodo, toutes avec bon score
        mock_metadata = []
        for i in range(8):
            mock_metadata.append({"document": "cim10", "code": f"K8{i}", "page": i, "extrait": f"K8{i}"})
        for i in range(8):
            mock_metadata.append({"document": "guide_methodo", "page": i + 10, "extrait": f"Guide {i}"})

        mock_index = MagicMock()
        mock_index.ntotal = 16
        scores = np.array([[0.9 - i * 0.03 for i in range(16)]], dtype=np.float32)
        indices = np.array([list(range(16))], dtype=np.int64)
        mock_index.search.return_value = (scores, indices)

        with patch("src.medical.rag_index.get_index", return_value=(mock_index, mock_metadata)), \
             patch("src.medical.rag_search._get_embed_model") as mock_model:
            mock_model.return_value.encode.return_value = np.array([[0.1] * 768], dtype=np.float32)
            results = search_similar("pancréatite", top_k=10)

        cim10_count = sum(1 for r in results if r["document"] == "cim10")
        assert cim10_count >= 6, f"Seulement {cim10_count} sources CIM-10 sur {len(results)}"


class TestRAGSearchMocked:
    def test_search_similar_no_index(self):
        """search_similar retourne une liste vide si l'index n'existe pas."""
        from src.medical.rag_search import search_similar

        with patch("src.medical.rag_index.get_index", return_value=None):
            results = search_similar("pancréatite aiguë")
            assert results == []

    def test_enrich_diagnostic_no_sources(self):
        """enrich_diagnostic ne plante pas si aucune source trouvée."""
        from src.medical.rag_search import enrich_diagnostic

        diag = Diagnostic(texte="test quelconque", cim10_suggestion="Z99.9")

        with patch("src.medical.rag_search.search_similar", return_value=[]):
            enrich_diagnostic(diag, {"sexe": "M", "age": 50})

        assert diag.sources_rag == []
        assert diag.justification is None
        assert diag.raisonnement is None

    def test_enrich_diagnostic_with_sources_no_ollama(self):
        """Enrichissement avec sources FAISS mais sans Ollama."""
        from src.medical.rag_search import enrich_diagnostic

        diag = Diagnostic(texte="Pancréatite aiguë", cim10_suggestion="K85.9")
        mock_sources = [
            {
                "document": "cim10",
                "page": 496,
                "code": "K85",
                "extrait": "K85 Pancréatite aiguë...",
                "score": 0.92,
            },
        ]

        with patch("src.medical.rag_search.search_similar", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=None):
            enrich_diagnostic(diag, {"sexe": "M", "age": 50})

        assert len(diag.sources_rag) == 1
        assert diag.sources_rag[0].document == "cim10"
        assert diag.sources_rag[0].code == "K85"
        assert diag.justification is None
        assert diag.raisonnement is None

    def test_enrich_diagnostic_with_ollama(self):
        """Enrichissement complet avec sources + Ollama + raisonnement."""
        from src.medical.rag_search import enrich_diagnostic

        diag = Diagnostic(texte="Pancréatite aiguë biliaire")
        mock_sources = [
            {
                "document": "cim10",
                "page": 496,
                "code": "K85",
                "extrait": "K85 Pancréatite aiguë...",
                "score": 0.95,
            },
        ]
        mock_llm = {
            "code": "K85.1",
            "confidence": "high",
            "justification": "Pancréatite aiguë d'origine biliaire = K85.1",
            "raisonnement": "1. ANALYSE CLINIQUE : La pancréatite...",
        }

        with patch("src.medical.rag_search.search_similar", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=mock_llm):
            enrich_diagnostic(diag, {"sexe": "F", "age": 43})

        assert diag.cim10_suggestion == "K85.1"
        assert diag.cim10_confidence == "high"
        assert diag.justification == "Pancréatite aiguë d'origine biliaire = K85.1"
        assert diag.raisonnement == "1. ANALYSE CLINIQUE : La pancréatite..."
        assert len(diag.sources_rag) == 1

    def test_enrich_diagnostic_invalid_code_ignored(self):
        """Un code Ollama invalide ne remplace pas le code existant."""
        from src.medical.rag_search import enrich_diagnostic

        diag = Diagnostic(texte="Pancréatite aiguë", cim10_suggestion="K85.9")
        mock_sources = [
            {"document": "cim10", "page": 496, "code": "K85", "extrait": "K85", "score": 0.9},
        ]
        mock_llm = {
            "code": "QQ9.99",  # code invalide (pas de parent valide)
            "confidence": "high",
            "justification": "Hallucination",
        }

        with patch("src.medical.rag_search.search_similar", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=mock_llm):
            enrich_diagnostic(diag, {"sexe": "M", "age": 50})

        # Le code original est conservé (pas remplacé par le code invalide)
        assert diag.cim10_suggestion == "K85.9"

    def test_enrich_diagnostic_fallback_parent_code(self):
        """Un code invalide D71.9 est corrigé en D71 (code parent standalone)."""
        from src.medical.rag_search import enrich_diagnostic

        diag = Diagnostic(texte="Anomalie des leucocytes")
        mock_sources = [
            {"document": "cim10", "page": 100, "code": "D71", "extrait": "D71", "score": 0.9},
        ]
        mock_llm = {
            "code": "D71.9",  # invalide : D71 est standalone
            "confidence": "high",
            "justification": "Anomalie leucocytaire",
        }

        with patch("src.medical.rag_search.search_similar", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=mock_llm):
            enrich_diagnostic(diag, {"sexe": "M", "age": 60})

        assert diag.cim10_suggestion == "D71"

    def test_enrich_diagnostic_normalizes_code(self):
        """Un code Ollama sans point est normalisé (K851 → K85.1)."""
        from src.medical.rag_search import enrich_diagnostic

        diag = Diagnostic(texte="Pancréatite aiguë biliaire")
        mock_sources = [
            {"document": "cim10", "page": 496, "code": "K85", "extrait": "K85", "score": 0.9},
        ]
        mock_llm = {
            "code": "K851",  # sans point
            "confidence": "high",
            "justification": "Pancréatite biliaire",
        }

        with patch("src.medical.rag_search.search_similar", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=mock_llm):
            enrich_diagnostic(diag, {"sexe": "F", "age": 43})

        assert diag.cim10_suggestion == "K85.1"

    def test_enrich_diagnostic_est_dp_flag(self):
        """Le flag est_dp est bien passé à _build_prompt."""
        from src.medical.rag_search import enrich_diagnostic

        diag = Diagnostic(texte="Obésité")
        mock_sources = [
            {"document": "cim10", "page": 1, "code": "E66", "extrait": "E66 Obésité", "score": 0.9},
        ]

        with patch("src.medical.rag_search.search_similar", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=None) as mock_ollama, \
             patch("src.medical.rag_search._build_prompt", return_value="prompt") as mock_prompt:
            enrich_diagnostic(diag, {"sexe": "F", "age": 43}, est_dp=False)
            mock_prompt.assert_called_once_with("Obésité", mock_sources, {"sexe": "F", "age": 43}, est_dp=False)


class TestEnrichDossier:
    """Tests pour enrich_dossier avec le contexte enrichi."""

    def test_enriched_context(self):
        """enrich_dossier passe le contexte enrichi (bio, imagerie, etc.)."""
        from src.medical.rag_search import enrich_dossier
        from src.config import Sejour, BiologieCle, Imagerie

        dossier = DossierMedical(
            sejour=Sejour(sexe="F", age=43, duree_sejour=6, imc=34.4),
            diagnostic_principal=Diagnostic(texte="Pancréatite aiguë biliaire"),
            antecedents=["HTA", "diabète type 2"],
            biologie_cle=[
                BiologieCle(test="Lipasémie", valeur="850", anomalie=True),
            ],
            imagerie=[
                Imagerie(type="TDM abdominal", conclusion="pancréatite stade C"),
            ],
            complications=["éruption cutanée"],
        )

        captured_contexts = []

        def mock_enrich(diag, contexte, est_dp=True, cache=None):
            captured_contexts.append(contexte.copy())

        with patch("src.medical.rag_search.enrich_diagnostic", side_effect=mock_enrich), \
             patch("src.medical.rag_search.OllamaCache") as mock_cache_cls:
            mock_cache_cls.return_value = MagicMock()
            enrich_dossier(dossier)

        assert len(captured_contexts) == 1  # DP seulement (pas de DAS)
        ctx = captured_contexts[0]
        assert ctx["sexe"] == "F"
        assert ctx["age"] == 43
        assert ctx["duree_sejour"] == 6
        assert ctx["imc"] == 34.4
        assert ctx["antecedents"] == ["HTA", "diabète type 2"]
        assert ctx["biologie_cle"] == [("Lipasémie", "850", True)]
        assert ctx["imagerie"] == [("TDM abdominal", "pancréatite stade C")]
        assert ctx["complications"] == ["éruption cutanée"]

    def test_das_gets_dp_context(self):
        """Les DAS reçoivent le texte du DP dans leur contexte."""
        from src.medical.rag_search import enrich_dossier

        dossier = DossierMedical(
            diagnostic_principal=Diagnostic(texte="Pancréatite aiguë biliaire"),
            diagnostics_associes=[
                Diagnostic(texte="Obésité"),
            ],
        )

        captured = []

        def mock_enrich(diag, contexte, est_dp=True, cache=None):
            captured.append({"texte": diag.texte, "est_dp": est_dp, "dp_texte": contexte.get("dp_texte")})

        with patch("src.medical.rag_search.enrich_diagnostic", side_effect=mock_enrich), \
             patch("src.medical.rag_search.OllamaCache") as mock_cache_cls:
            mock_cache_cls.return_value = MagicMock()
            enrich_dossier(dossier)

        assert len(captured) == 2
        # DP n'a pas dp_texte dans son contexte
        assert captured[0]["est_dp"] is True
        assert captured[0]["dp_texte"] is None
        # DAS a dp_texte
        assert captured[1]["est_dp"] is False
        assert captured[1]["dp_texte"] == "Pancréatite aiguë biliaire"


class TestNormalizeCode:
    def test_insert_dot(self):
        from src.medical.cim10_dict import normalize_code
        assert normalize_code("K810") == "K81.0"

    def test_already_dotted(self):
        from src.medical.cim10_dict import normalize_code
        assert normalize_code("k85.1") == "K85.1"

    def test_three_chars(self):
        from src.medical.cim10_dict import normalize_code
        assert normalize_code("K85") == "K85"

    def test_strip_spaces(self):
        from src.medical.cim10_dict import normalize_code
        assert normalize_code(" E660 ") == "E66.0"


class TestValidateCodeCIM10:
    def test_known_code(self):
        from src.medical.cim10_dict import validate_code
        is_valid, label = validate_code("K81.9")
        assert is_valid is True
        assert label  # non vide

    def test_unknown_code(self):
        from src.medical.cim10_dict import validate_code
        is_valid, label = validate_code("Z99.99")
        assert is_valid is False
        assert label == ""

    def test_normalize_before_validate(self):
        """K810 doit être normalisé en K81.0 et trouvé."""
        from src.medical.cim10_dict import validate_code
        is_valid, label = validate_code("K810")
        assert is_valid is True

    def test_three_char_code(self):
        """Code parent sans point (K85) doit être validé."""
        from src.medical.cim10_dict import validate_code
        is_valid, label = validate_code("K85")
        assert is_valid is True


class TestFallbackParentCode:
    def test_d71_9_to_d71(self):
        """D71.9 (invalide) → D71 (standalone valide)."""
        from src.medical.cim10_dict import fallback_parent_code
        assert fallback_parent_code("D71.9") == "D71"

    def test_r69_8_to_r69(self):
        """R69.8 (invalide) → R69 (standalone valide)."""
        from src.medical.cim10_dict import fallback_parent_code
        assert fallback_parent_code("R69.8") == "R69"

    def test_valid_code_no_fallback(self):
        """Un code déjà valide ne devrait pas matcher (parent aussi valide)."""
        from src.medical.cim10_dict import fallback_parent_code
        # K85.1 est valide, donc on ne devrait pas appeler fallback
        # Mais si on l'appelle, K85 est aussi valide → retourne K85
        result = fallback_parent_code("K85.1")
        assert result == "K85"  # le parent est valide

    def test_truly_invalid_no_fallback(self):
        """Un code sans parent valide retourne None."""
        from src.medical.cim10_dict import fallback_parent_code
        assert fallback_parent_code("QQ9.99") is None

    def test_three_char_code_no_fallback(self):
        """Un code 3 caractères sans point ne peut pas remonter."""
        from src.medical.cim10_dict import fallback_parent_code
        assert fallback_parent_code("QQ9") is None


class TestValidateCIM10PostProcessing:
    def test_hallucination_rejected(self):
        """Les codes hallucination (Aucun, N/A...) sont rejetés."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostic_principal=Diagnostic(texte="Test", cim10_suggestion="Aucun"),
        )
        _validate_cim10(dossier)
        assert dossier.diagnostic_principal.cim10_suggestion is None
        assert any("rejeté" in a for a in dossier.alertes_codage)

    def test_normalizes_format(self):
        """K810 est normalisé en K81.0."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostic_principal=Diagnostic(texte="Cholécystite", cim10_suggestion="K810"),
        )
        _validate_cim10(dossier)
        assert dossier.diagnostic_principal.cim10_suggestion == "K81.0"

    def test_invalid_code_gets_low_confidence(self):
        """Un code inexistant reçoit confidence=low et une alerte."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostics_associes=[
                Diagnostic(texte="Chose bizarre", cim10_suggestion="Z99.99"),
            ],
        )
        _validate_cim10(dossier)
        assert dossier.diagnostics_associes[0].cim10_confidence == "low"
        assert any("absent du dictionnaire" in a for a in dossier.alertes_codage)

    def test_valid_code_unchanged(self):
        """Un code valide n'est pas modifié et pas d'alerte."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostic_principal=Diagnostic(texte="Pancréatite", cim10_suggestion="K85.1"),
        )
        _validate_cim10(dossier)
        assert dossier.diagnostic_principal.cim10_suggestion == "K85.1"
        assert not any("CIM-10" in a for a in dossier.alertes_codage)

    def test_non_codable_rejected(self):
        """'non_codable' est rejeté comme hallucination."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostics_associes=[
                Diagnostic(texte="Truc", cim10_suggestion="non_codable"),
            ],
        )
        _validate_cim10(dossier)
        assert dossier.diagnostics_associes[0].cim10_suggestion is None

    def test_hallucination_fallback_found(self):
        """Hallucination rejetée mais fallback dictionnaire trouve un code."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostic_principal=Diagnostic(texte="Cholécystite aiguë", cim10_suggestion="Aucun"),
        )
        _validate_cim10(dossier)
        assert dossier.diagnostic_principal.cim10_suggestion == "K81.0"
        assert dossier.diagnostic_principal.cim10_confidence == "medium"
        assert any("fallback" in a for a in dossier.alertes_codage)

    def test_invalid_code_fallback_found(self):
        """Code invalide remplacé par fallback dictionnaire."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostics_associes=[
                Diagnostic(texte="Hypertension artérielle", cim10_suggestion="I99.99"),
            ],
        )
        _validate_cim10(dossier)
        assert dossier.diagnostics_associes[0].cim10_suggestion == "I10"
        assert dossier.diagnostics_associes[0].cim10_confidence == "medium"
        assert any("fallback" in a for a in dossier.alertes_codage)

    def test_invalid_code_no_fallback(self):
        """Code invalide sans fallback possible → low confidence."""
        from src.medical.cim10_extractor import _validate_cim10

        dossier = DossierMedical(
            diagnostics_associes=[
                Diagnostic(texte="Chose bizarre inconnue", cim10_suggestion="Z99.99"),
            ],
        )
        _validate_cim10(dossier)
        assert dossier.diagnostics_associes[0].cim10_suggestion == "Z99.99"
        assert dossier.diagnostics_associes[0].cim10_confidence == "low"
        assert any("absent du dictionnaire" in a for a in dossier.alertes_codage)


class TestFormatContexte:
    """Tests pour _format_contexte."""

    def test_minimal_context(self):
        from src.medical.rag_search import _format_contexte

        result = _format_contexte({})
        assert result == "Non précisé"

    def test_full_context(self):
        from src.medical.rag_search import _format_contexte

        ctx = {
            "sexe": "F",
            "age": 43,
            "imc": 34.4,
            "duree_sejour": 6,
            "antecedents": ["HTA", "diabète type 2"],
            "biologie_cle": [("Lipasémie", "850", True), ("CRP", "45", True)],
            "imagerie": [("TDM abdominal", "pancréatite stade C Balthazar")],
            "complications": ["éruption cutanée"],
            "dp_texte": "Pancréatite aiguë biliaire",
        }
        result = _format_contexte(ctx)

        assert "F, 43 ans, IMC 34.4" in result
        assert "6 jours" in result
        assert "HTA" in result
        assert "Lipasémie 850" in result
        assert "TDM abdominal" in result
        assert "éruption cutanée" in result
        assert "Pancréatite aiguë biliaire" in result


class TestActeCCAMExtended:
    def test_backward_compatible(self):
        """Les nouveaux champs RAG sont optionnels — rétrocompatible."""
        a = ActeCCAM(texte="Cholécystectomie", code_ccam_suggestion="HMFC004")
        assert a.texte == "Cholécystectomie"
        assert a.code_ccam_suggestion == "HMFC004"
        assert a.ccam_confidence is None
        assert a.justification is None
        assert a.raisonnement is None
        assert a.sources_rag == []

    def test_with_rag_fields(self):
        a = ActeCCAM(
            texte="Cholécystectomie par coelioscopie",
            code_ccam_suggestion="HMFC004",
            ccam_confidence="high",
            justification="HMFC004 correspond à la cholécystectomie par coelioscopie",
            raisonnement="ANALYSE ACTE : Cholécystectomie par voie coelioscopique...",
            sources_rag=[
                RAGSource(document="ccam", page=10, code="HMFC004"),
            ],
        )
        assert a.ccam_confidence == "high"
        assert a.justification is not None
        assert len(a.sources_rag) == 1
        assert a.sources_rag[0].code == "HMFC004"

    def test_serialization_exclude_none(self):
        a = ActeCCAM(texte="Test", code_ccam_suggestion="HMFC004")
        data = a.model_dump(exclude_none=True)
        assert "ccam_confidence" not in data
        assert "justification" not in data
        assert "raisonnement" not in data
        assert "sources_rag" in data


class TestSearchSimilarCCAM:
    def test_prioritizes_ccam(self):
        """Les sources CCAM sont priorisées (au moins 5 sur 8)."""
        from src.medical.rag_search import search_similar_ccam
        import numpy as np

        mock_metadata = []
        for i in range(6):
            mock_metadata.append({"document": "ccam", "code": f"HMFC00{i}", "page": i, "extrait": f"CCAM {i}"})
        for i in range(6):
            mock_metadata.append({"document": "guide_methodo", "page": i + 10, "extrait": f"Guide {i}"})

        mock_index = MagicMock()
        mock_index.ntotal = 12
        scores = np.array([[0.9 - i * 0.03 for i in range(12)]], dtype=np.float32)
        indices = np.array([list(range(12))], dtype=np.int64)
        mock_index.search.return_value = (scores, indices)

        with patch("src.medical.rag_index.get_index", return_value=(mock_index, mock_metadata)), \
             patch("src.medical.rag_search._get_embed_model") as mock_model:
            mock_model.return_value.encode.return_value = np.array([[0.1] * 768], dtype=np.float32)
            results = search_similar_ccam("cholécystectomie", top_k=8)

        ccam_count = sum(1 for r in results if r["document"] == "ccam")
        assert ccam_count >= 5, f"Seulement {ccam_count} sources CCAM sur {len(results)}"

    def test_no_index(self):
        """search_similar_ccam retourne une liste vide si l'index n'existe pas."""
        from src.medical.rag_search import search_similar_ccam

        with patch("src.medical.rag_index.get_index", return_value=None):
            results = search_similar_ccam("cholécystectomie")
            assert results == []


class TestEnrichActe:
    def test_enrich_with_ollama(self):
        """Enrichissement complet avec sources + Ollama."""
        from src.medical.rag_search import enrich_acte

        acte = ActeCCAM(texte="Cholécystectomie par coelioscopie")
        mock_sources = [
            {
                "document": "ccam",
                "page": 10,
                "code": "HMFC004",
                "extrait": "HMFC004 Cholécystectomie par coelioscopie...",
                "score": 0.92,
            },
        ]
        mock_llm = {
            "code": "HMFC004",
            "confidence": "high",
            "justification": "Cholécystectomie par coelioscopie = HMFC004",
            "raisonnement": "ANALYSE ACTE : Cholécystectomie par voie coelioscopique...",
        }

        with patch("src.medical.rag_search.search_similar_ccam", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=mock_llm), \
             patch("src.medical.rag_search.ccam_validate", return_value=(True, "Cholécystectomie")):
            enrich_acte(acte, {"sexe": "F", "age": 43})

        assert acte.code_ccam_suggestion == "HMFC004"
        assert acte.ccam_confidence == "high"
        assert acte.justification == "Cholécystectomie par coelioscopie = HMFC004"
        assert acte.raisonnement is not None
        assert len(acte.sources_rag) == 1

    def test_enrich_no_sources(self):
        """enrich_acte ne plante pas si aucune source trouvée."""
        from src.medical.rag_search import enrich_acte

        acte = ActeCCAM(texte="Acte inconnu", code_ccam_suggestion="ABCD123")

        with patch("src.medical.rag_search.search_similar_ccam", return_value=[]):
            enrich_acte(acte, {"sexe": "M", "age": 50})

        assert acte.sources_rag == []
        assert acte.justification is None

    def test_enrich_no_ollama(self):
        """Enrichissement avec sources FAISS mais sans Ollama."""
        from src.medical.rag_search import enrich_acte

        acte = ActeCCAM(texte="Cholécystectomie", code_ccam_suggestion="HMFC004")
        mock_sources = [
            {"document": "ccam", "page": 10, "code": "HMFC004", "extrait": "HMFC004", "score": 0.9},
        ]

        with patch("src.medical.rag_search.search_similar_ccam", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=None):
            enrich_acte(acte, {"sexe": "M", "age": 50})

        assert len(acte.sources_rag) == 1
        assert acte.justification is None
        assert acte.raisonnement is None

    def test_enrich_invalid_code(self):
        """Un code CCAM invalide d'Ollama ne remplace pas le code existant."""
        from src.medical.rag_search import enrich_acte

        acte = ActeCCAM(texte="Cholécystectomie", code_ccam_suggestion="HMFC004")
        mock_sources = [
            {"document": "ccam", "page": 10, "code": "HMFC004", "extrait": "HMFC004", "score": 0.9},
        ]
        mock_llm = {
            "code": "ZZZZ999",
            "confidence": "high",
            "justification": "Hallucination",
        }

        with patch("src.medical.rag_search.search_similar_ccam", return_value=mock_sources), \
             patch("src.medical.rag_search._call_ollama", return_value=mock_llm), \
             patch("src.medical.rag_search.ccam_validate", return_value=(False, "")):
            enrich_acte(acte, {"sexe": "M", "age": 50})

        # Le code original est conservé
        assert acte.code_ccam_suggestion == "HMFC004"
        # Mais la confidence est quand même affectée
        assert acte.ccam_confidence == "high"


class TestEnrichDossierCCAM:
    def test_enriches_actes(self):
        """enrich_dossier enrichit aussi les actes CCAM."""
        from src.medical.rag_search import enrich_dossier

        dossier = DossierMedical(
            diagnostic_principal=Diagnostic(texte="Lithiase vésiculaire"),
            actes_ccam=[
                ActeCCAM(texte="Cholécystectomie par coelioscopie"),
                ActeCCAM(texte="Anesthésie générale"),
            ],
        )

        enriched = []

        def mock_enrich_diag(diag, contexte, est_dp=True, cache=None):
            pass

        def mock_enrich_acte(acte, contexte, cache=None):
            enriched.append(acte.texte)

        with patch("src.medical.rag_search.enrich_diagnostic", side_effect=mock_enrich_diag), \
             patch("src.medical.rag_search.enrich_acte", side_effect=mock_enrich_acte), \
             patch("src.medical.rag_search.OllamaCache") as mock_cache_cls:
            mock_cache_cls.return_value = MagicMock()
            enrich_dossier(dossier)

        assert len(enriched) == 2
        assert "Cholécystectomie par coelioscopie" in enriched
        assert "Anesthésie générale" in enriched


class TestBuildPromptCCAM:
    def test_prompt_contains_acte(self):
        from src.medical.rag_search import _build_prompt_ccam

        sources = [{"document": "ccam", "code": "HMFC004", "page": 10, "extrait": "HMFC004 Cholécystectomie"}]
        contexte = {"sexe": "F", "age": 43}
        prompt = _build_prompt_ccam("Cholécystectomie par coelioscopie", sources, contexte)

        assert "Cholécystectomie par coelioscopie" in prompt
        assert "CCAM" in prompt
        assert "analyse_acte" in prompt
        assert "objet JSON" in prompt

    def test_prompt_contains_source_info(self):
        from src.medical.rag_search import _build_prompt_ccam

        sources = [{"document": "ccam", "code": "HMFC004", "page": 10, "extrait": "HMFC004 Cholécystectomie par coelioscopie"}]
        contexte = {}
        prompt = _build_prompt_ccam("Cholécystectomie", sources, contexte)

        assert "CCAM PMSI V4 2025" in prompt
        assert "HMFC004" in prompt


class TestParseOllamaResponseCCAM:
    def test_parse_ccam_structured_json(self):
        """Le parsing extrait analyse_acte dans le raisonnement."""
        from src.medical.rag_search import _parse_ollama_response
        import json

        raw = json.dumps({
            "analyse_acte": "Cholécystectomie par voie coelioscopique",
            "codes_candidats": "HMFC004, HMFC003",
            "discrimination": "HMFC004 est le code spécifique à la coelioscopie",
            "code": "HMFC004",
            "confidence": "high",
            "justification": "Cholécystectomie coelioscopique = HMFC004",
        })

        result = _parse_ollama_response(raw)
        assert result is not None
        assert result["code"] == "HMFC004"
        assert "raisonnement" in result
        assert "ANALYSE ACTE" in result["raisonnement"]
        assert "CODES CANDIDATS" in result["raisonnement"]
        assert "analyse_acte" not in result