from __future__ import annotations

import json
from dataclasses import replace

from rag.agent import RagAgent, _patient_generation_query
from rag.answer import GroundedAnswerService
from rag.clinical import (
    CaseContextAction,
    CandidateStatus,
    ConditionNormalizer,
    ConditionQuery,
    ConditionRelation,
    HepaticContext,
    PatientContext,
    RenalContext,
)
from rag.models import Evidence, SourceRef
from rag.sections import SECTION_PHRASES
from rag.understanding import (
    LlmQueryUnderstander,
    QueryFrame,
    _apply_condition_candidate_cue,
    _apply_contextual_candidate_safety,
    _apply_general_condition_scope,
    _apply_named_drug_cues,
    _apply_reverse_relation_cues,
    _merge_with_prior_frame,
)


SOURCE = SourceRef(physical_page=10, precision="chunk_page_range", printed_page=11)


def _evidence(drug_id: str, section: str = "chi_dinh") -> Evidence:
    return Evidence(
        evidence_id=f"{drug_id}__{section}__0",
        matched_doc_id=f"{drug_id}__{section}__0",
        kind="prose",
        text=f"{drug_id} có nội dung {section}.",
        score=1.0,
        source_refs=(SOURCE,),
        hydrated_from_parent=False,
        requires_visual_check=False,
        drug_id=drug_id,
        drug_name=drug_id.upper(),
        section_key=section,
        section_title=section,
    )


class _Understander:
    def __init__(self, frame: QueryFrame) -> None:
        self.frame = frame

    def understand(self, turn, history=(), budget=None, prior_frame=None):
        return self.frame


class _NoRetrieval:
    def retrieve_by_indication(self, indication):
        raise AssertionError(f"retrieval must not run for relation {indication}")


def _agent(frame: QueryFrame) -> RagAgent:
    return RagAgent(
        _Understander(frame),
        _NoRetrieval(),
        GroundedAnswerService(routing=None),
    )


def test_condition_normalizer_handles_professional_aliases_without_drug_mapping():
    normalizer = ConditionNormalizer()
    assert normalizer.normalize("THA dùng gì", "THA").normalized_condition == "tăng huyết áp"
    assert normalizer.normalize("cao huyết áp", "cao huyết áp").normalized_condition == "tăng huyết áp"
    assert normalizer.normalize("Gout", "gout").normalized_condition == "gút"
    assert normalizer.detect_known_alias("BN viêm phổi dùng thuốc gì?").normalized_condition == "viêm phổi"
    assert normalizer.normalize("bệnh lạ", "bệnh lạ").normalized_condition == "bệnh lạ"


def test_explicit_indication_relation_is_a_condition_candidate_lookup():
    noisy = QueryFrame(
        turn_type="condition_relation",
        condition=ConditionNormalizer().normalize("bệnh gút", "gút"),
        needs_clarify=True,
        clarify_reason="Hỏi lại sai hướng",
    )

    frame = _apply_condition_candidate_cue(
        noisy,
        "Thuốc nào có chỉ định liên quan bệnh gút?",
        ConditionNormalizer(),
    )

    assert frame.turn_type == "condition_to_drug"
    assert frame.condition_relation == ConditionRelation.INDICATION
    assert frame.needs_clarify is False


def test_general_condition_does_not_invent_patient_hepatic_context():
    frame = QueryFrame(
        turn_type="condition_to_drug",
        indication="viêm gan B mạn",
        condition=ConditionQuery(
            original_query="Viêm gan B mạn dùng thuốc gì?",
            normalized_condition="viêm gan B mạn",
        ),
        patient_context=PatientContext(
            primary_condition="viêm gan B mạn",
            hepatic=HepaticContext(description="viêm gan B mạn"),
        ),
    )

    cleaned = _apply_general_condition_scope(
        frame, "Viêm gan B mạn dùng thuốc gì?"
    )

    assert cleaned.patient_context.primary_condition == "viêm gan B mạn"
    assert cleaned.patient_context.requires_safety_review is False


def test_patient_allergy_condition_lookup_keeps_safety_context():
    patient = PatientContext(
        primary_condition="viêm phổi", allergies=("penicillin",)
    )
    frame = QueryFrame(
        turn_type="condition_to_drug",
        condition=ConditionQuery(
            original_query="BN dị ứng penicillin, viêm phổi dùng thuốc gì?",
            normalized_condition="viêm phổi",
        ),
        patient_context=patient,
    )

    kept = _apply_general_condition_scope(
        frame, "BN dị ứng penicillin, viêm phổi dùng thuốc gì?"
    )

    assert kept.patient_context == patient
    assert kept.patient_context.requires_safety_review is True


def test_candidate_safety_followup_stays_on_prior_condition_lookup():
    prior = QueryFrame(
        turn_type="condition_to_drug",
        indication="tăng huyết áp",
        condition=ConditionQuery(
            original_query="BN bị tăng huyết áp",
            normalized_condition="tăng huyết áp",
        ),
        patient_context=PatientContext(
            age_text="68 tuổi",
            renal=RenalContext(description="CKD", ckd_stage="G4"),
        ),
    )
    noisy = QueryFrame(
        turn_type="condition_relation",
        condition_relation=ConditionRelation.CONTRAINDICATION,
        depends_on_previous_turn=True,
        patient_context=prior.patient_context,
        needs_clarify=False,
    )

    corrected = _apply_contextual_candidate_safety(
        noisy,
        "Trong các thuốc trên cái nào cần lưu ý hơn với bệnh thận?",
        prior,
    )

    assert corrected.turn_type == "condition_to_drug"
    assert corrected.condition == prior.condition
    assert corrected.condition_relation == ConditionRelation.INDICATION


def test_broad_condition_is_clarified_but_specific_subtype_is_not():
    normalizer = ConditionNormalizer()
    broad = normalizer.normalize("Viêm gan dùng thuốc gì?", "viêm gan")
    specific = normalizer.normalize(
        "Viêm gan B mạn dùng thuốc gì?", "viêm gan B mạn", subtype="B mạn"
    )
    assert broad.ambiguous is True
    assert "A, B, C" in broad.clarify_question
    assert specific.ambiguous is False


def test_bare_broad_question_detector_does_not_overclarify_a_specific_site():
    normalizer = ConditionNormalizer()

    broad = normalizer.detect_broad_question("Nhiễm trùng dùng thuốc gì?")
    specific = normalizer.detect_broad_question(
        "Nhiễm trùng đường tiết niệu dùng thuốc gì?"
    )

    assert broad is not None
    assert broad.ambiguous is True
    assert broad.clarify_question
    assert specific is None


def test_relation_confusion_never_reaches_indication_retrieval():
    for relation, question in (
        (ConditionRelation.ADVERSE_EFFECT, "Thuốc nào gây tăng huyết áp?"),
        (ConditionRelation.CONTRAINDICATION, "Thuốc nào chống chỉ định ở bệnh nhân gout?"),
    ):
        reply = _agent(QueryFrame(
            turn_type="condition_relation",
            indication="tăng huyết áp",
            condition_relation=relation,
        )).handle(question)
        assert reply.decision == "abstain"
        assert reply.reason == "unsupported_reverse_relation"
        assert "không" in reply.answer.lower()


def test_explicit_reverse_relation_cue_overrides_a_noisy_llm_clarification():
    noisy = QueryFrame(
        turn_type="out_of_scope",
        needs_clarify=True,
        clarify_reason="Bạn muốn hỏi thuốc nào?",
    )

    adverse = _apply_reverse_relation_cues(noisy, "Thuốc nào gây tăng huyết áp?")
    contraindicated = _apply_reverse_relation_cues(
        noisy, "Thuốc nào chống chỉ định ở bệnh nhân gout?"
    )

    assert adverse.turn_type == "condition_relation"
    assert adverse.condition_relation == ConditionRelation.ADVERSE_EFFECT
    assert adverse.needs_clarify is False
    assert contraindicated.condition_relation == ConditionRelation.CONTRAINDICATION


def test_patient_candidate_wording_is_not_mistaken_for_reverse_contraindication():
    normalizer = ConditionNormalizer()
    noisy = QueryFrame(
        turn_type="drug_attribute",
        drugs=("digoxin",),
        needs_clarify=True,
        clarify_reason="Bạn muốn hỏi digoxin?",
    )

    candidate = _apply_condition_candidate_cue(
        noisy,
        "BN gout kèm suy thận nặng dùng thuốc nào cần thận trọng hoặc chống chỉ định?",
        normalizer,
    )
    after_relation_guard = _apply_reverse_relation_cues(
        candidate,
        "BN gout kèm suy thận nặng dùng thuốc nào cần thận trọng hoặc chống chỉ định?",
    )

    assert candidate.turn_type == "condition_to_drug"
    assert candidate.condition is not None
    assert candidate.condition.normalized_condition == "gút"
    assert after_relation_guard.turn_type == "condition_to_drug"


def test_named_drug_safety_and_purpose_cues_override_noisy_relation_frames():
    noisy = QueryFrame(
        turn_type="condition_relation",
        drugs=("probenecid",),
        condition_relation=ConditionRelation.CONTRAINDICATION,
        needs_clarify=True,
        clarify_reason="Cần làm rõ",
    )

    safety = _apply_named_drug_cues(
        noisy, "BN eGFR 25, probenecid có dùng được không?"
    )
    purpose = _apply_named_drug_cues(
        replace(noisy, drugs=("paracetamol_acetaminophen",)),
        "Paracetamol có tác dụng gì?",
    )
    golden_contraindication = _apply_named_drug_cues(
        replace(noisy, drugs=("paracetamol_acetaminophen",)),
        "Chống chỉ định của Paracetamol là gì?",
    )
    golden_contraindication_with_omitted_llm_drug = _apply_named_drug_cues(
        replace(noisy, drugs=()),
        "Chống chỉ định của Paracetamol là gì?",
        resolved_turn_drug="paracetamol_acetaminophen",
    )
    reverse_contraindication = _apply_named_drug_cues(
        replace(noisy, drugs=("warfarin",)),
        "Thuốc nào chống chỉ định ở bệnh nhân đang dùng warfarin?",
    )

    assert safety.turn_type == "drug_attribute"
    assert safety.attribute == "chong_chi_dinh"
    assert safety.needs_clarify is False
    assert purpose.turn_type == "drug_to_condition"
    assert purpose.attribute == "chi_dinh"
    assert golden_contraindication.turn_type == "drug_attribute"
    assert golden_contraindication.attribute == "chong_chi_dinh"
    assert golden_contraindication.condition_relation == ConditionRelation.CONTRAINDICATION
    assert golden_contraindication_with_omitted_llm_drug.turn_type == "drug_attribute"
    assert golden_contraindication_with_omitted_llm_drug.drugs == (
        "paracetamol_acetaminophen",
    )
    assert reverse_contraindication.turn_type == "condition_relation"


def test_owned_section_phrase_routes_every_monograph_section_without_llm_help():
    noisy = QueryFrame(
        turn_type="condition_relation",
        condition_relation=ConditionRelation.CONTRAINDICATION,
        needs_clarify=True,
        clarify_reason="Cần làm rõ",
    )

    for section_key, phrases in SECTION_PHRASES.items():
        phrase = phrases[0]
        corrected = _apply_named_drug_cues(
            noisy,
            f"{phrase} của Paracetamol là gì?",
            resolved_turn_drug="paracetamol_acetaminophen",
            resolved_section_key=section_key,
            resolved_section_phrase=phrase,
        )

        assert corrected.turn_type == "drug_attribute", section_key
        assert corrected.drugs == ("paracetamol_acetaminophen",), section_key
        assert corrected.attribute == section_key
        assert corrected.needs_clarify is False


def test_ambiguous_condition_clarifies_before_retrieval():
    reply = _agent(QueryFrame(
        turn_type="condition_to_drug",
        indication="viêm gan",
        condition=ConditionQuery(
            original_query="Viêm gan dùng thuốc gì?",
            normalized_condition="viêm gan",
            ambiguous=True,
            clarify_question="Bạn đang hỏi viêm gan A, B, C hay loại nào?",
        ),
    )).handle("Viêm gan dùng thuốc gì?")
    assert reply.decision == "clarify"
    assert reply.reason == "ambiguous_condition"


def test_patient_context_merges_only_for_explicit_same_case_continuation():
    prior = QueryFrame(
        turn_type="condition_to_drug",
        patient_context=PatientContext(
            age_text="68 tuổi",
            comorbidities=("CKD G4",),
            renal=RenalContext(description="CKD", ckd_stage="G4"),
        ),
    )
    current = QueryFrame(
        turn_type="condition_to_drug",
        indication="tăng huyết áp",
        condition=ConditionQuery("BN bị tăng huyết áp", "tăng huyết áp"),
        patient_context=PatientContext(primary_condition="tăng huyết áp"),
        context_action=CaseContextAction.CONTINUE,
    )
    merged = _merge_with_prior_frame(current, prior)
    assert merged.patient_context.age_text == "68 tuổi"
    assert merged.patient_context.renal.ckd_stage == "G4"

    new_case = _merge_with_prior_frame(
        QueryFrame(
            turn_type="condition_to_drug",
            indication="gút",
            context_action=CaseContextAction.NEW,
            patient_context=PatientContext(primary_condition="gút"),
        ),
        prior,
    )
    assert new_case.patient_context.age_text is None
    assert new_case.patient_context.renal.present is False


def test_patient_generation_query_keeps_task_but_not_user_only_numbers():
    frame = QueryFrame(
        turn_type="condition_to_drug",
        condition=ConditionQuery(
            "THA", "tăng huyết áp"
        ),
        patient_context=PatientContext(
            age_text="68 tuổi",
            primary_condition="tăng huyết áp",
            comorbidities=("CKD G4",),
            current_medications=("digoxin",),
            renal=RenalContext(description="eGFR 25", ckd_stage="G4", egfr="25"),
        ),
    )

    query = _patient_generation_query(frame)

    assert "tăng huyết áp" in query
    assert "68" not in query
    assert "G4" not in query
    assert "25" not in query
    assert "digoxin" not in query


def test_candidate_status_keeps_indication_separate_from_patient_safety():
    from rag.clinical import MedicationCandidateAssessment

    assessment = MedicationCandidateAssessment(
        drug_id="a",
        drug_name="A",
        indication_supported=True,
        indication_evidence=(_evidence("a"),),
        status=CandidateStatus.INSUFFICIENT_EVIDENCE,
    )
    assert assessment.indication_supported is True
    assert assessment.status == CandidateStatus.INSUFFICIENT_EVIDENCE
    assert len(assessment.evidence) == 1


class _NoCandidates:
    def resolve(self, query):
        class Result:
            status = "not_found"
            drug_id = None
            candidate_drug_ids = ()
        return Result()

    def suggest(self, query, k=3, min_score=0.5):
        return []


class _JsonLlm:
    def __init__(self, payload: dict) -> None:
        self.payload = payload

    def generate(self, system, user, schema):
        return json.dumps(self.payload, ensure_ascii=False)


def test_understanding_parses_condition_and_patient_context_without_inventing_fields():
    payload = {
        "turn_type": "condition_to_drug",
        "drugs": [],
        "unknown_drugs": [],
        "attribute": None,
        "population": "suy_than",
        "weight_kg": None,
        "age_text": "68 tuổi",
        "indication": "THA",
        "condition": {
            "original_text": "THA",
            "normalized_condition": "tăng huyết áp",
            "subtype": None,
            "qualifiers": [],
            "ambiguous": False,
            "clarify_question": None,
        },
        "condition_relation": "indication",
        "patient_context": {
            "age_text": "68 tuổi",
            "sex": None,
            "weight_kg": None,
            "primary_condition": "tăng huyết áp",
            "comorbidities": ["CKD G4", "gout"],
            "allergies": [],
            "previous_adverse_reactions": [],
            "current_medications": ["digoxin"],
            "pregnancy_status": None,
            "breastfeeding": None,
            "renal": {
                "description": "CKD",
                "ckd_stage": "G4",
                "egfr": None,
                "crcl": None,
                "creatinine": None,
            },
            "hepatic": {},
            "relevant_labs": ["K 5.7"],
            "treatment_history": [],
        },
        "context_action": "none",
        "route": None,
        "section_overview": False,
        "standalone_query": "BN 68 tuổi, THA + CKD G4 + gout, đang dùng digoxin",
        "depends_on_previous_turn": False,
        "needs_clarify": False,
        "clarify_reason": None,
        "quick_replies": [],
    }
    understander = LlmQueryUnderstander(
        _JsonLlm(payload), {}, _NoCandidates()
    )

    frame = understander.understand(
        "BN 68 tuổi, THA + CKD G4 + gout, K 5.7, đang dùng digoxin. Option hạ áp?"
    )

    assert frame.turn_type == "condition_to_drug"
    assert frame.condition.normalized_condition == "tăng huyết áp"
    assert frame.patient_context.comorbidities == ("CKD G4", "gout")
    assert frame.patient_context.current_medications == ("digoxin",)
    assert frame.patient_context.renal.ckd_stage == "G4"
    assert frame.patient_context.hepatic.present is False