from __future__ import annotations

import grantflow.swarm.critic_llm_policy as critic_llm_policy
from grantflow.swarm.nodes.critic import (
    _advisory_llm_findings_context,
    _apply_advisory_llm_score_cap,
    _classify_llm_finding_label,
    _downgrade_advisory_llm_findings,
    _is_advisory_llm_message,
)


def test_advisory_llm_findings_are_downgraded_and_score_penalty_is_capped():
    class _RuleCheck:
        def __init__(self, status: str):
            self.status = status

    class _RuleReport:
        def __init__(self):
            self.fatal_flaws = []
            self.checks = [_RuleCheck("pass")]

    state = {
        "citations": [
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.25},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.27},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.29},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.24},
            {"stage": "mel", "citation_type": "rag_result", "citation_confidence": 0.85},
        ]
    }
    llm_items = [
        {
            "code": "LLM_REVIEW_FLAG_1",
            "severity": "medium",
            "section": "logframe",
            "message": "Lack of defined baselines and targets for indicators, making it impossible to measure progress effectively.",
            "source": "llm",
        },
        {
            "code": "LLM_REVIEW_FLAG_2",
            "severity": "medium",
            "section": "logframe",
            "message": "Absence of detailed evidence excerpts to justify the selection of each indicator, reducing credibility.",
            "source": "llm",
        },
        {
            "code": "LLM_REVIEW_FLAG_3",
            "severity": "medium",
            "section": "toc",
            "message": "The Theory of Change objectives are not sufficiently specific and measurable.",
            "source": "llm",
        },
    ]

    advisory_ctx = _advisory_llm_findings_context(
        state=state, rule_report=_RuleReport(), llm_fatal_flaw_items=llm_items
    )
    assert advisory_ctx["applies"] is True
    downgraded, meta = _downgrade_advisory_llm_findings(llm_items, advisory_ctx=advisory_ctx)
    assert meta is not None and meta["applied"] is True
    assert all(item["severity"] == "low" for item in downgraded)

    combined, score_meta = _apply_advisory_llm_score_cap(
        combined_score=7.75,
        rule_score=9.25,
        llm_score=4.5,
        advisory_ctx=advisory_ctx,
    )
    assert combined == 8.5
    assert score_meta is not None and score_meta["applied"] is True


def test_usaid_early_draft_llm_findings_are_treated_as_advisory_when_grounded():
    class _RuleCheck:
        def __init__(self, status: str):
            self.status = status

    class _RuleReport:
        def __init__(self):
            self.fatal_flaws = []
            self.checks = [_RuleCheck("pass")]

    state = {
        "citations": [
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.51},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.49},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.53},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.47},
            {"stage": "mel", "citation_type": "rag_result", "citation_confidence": 0.8},
        ]
    }
    llm_items = [
        {
            "code": "LLM_REVIEW_FLAG_1",
            "severity": "medium",
            "section": "toc",
            "message": (
                "Weak causal links between Outputs and IRs: The proposal lacks detailed explanation on how training "
                "modules and workshops will lead to increased AI literacy and skills."
            ),
            "source": "llm",
        },
        {
            "code": "LLM_REVIEW_FLAG_2",
            "severity": "medium",
            "section": "toc",
            "message": (
                "Unrealistic assumptions: The proposal assumes civil servants are motivated to participate without "
                "evidence or strategies to sustain this motivation."
            ),
            "source": "llm",
        },
        {
            "code": "LLM_REVIEW_FLAG_3",
            "severity": "medium",
            "section": "general",
            "message": (
                "Missing cross-cutting themes: While gender equality is mentioned, the proposal lacks a detailed "
                "plan for implementation and measurement."
            ),
            "source": "llm",
        },
    ]

    advisory_ctx = _advisory_llm_findings_context(
        state=state, rule_report=_RuleReport(), llm_fatal_flaw_items=llm_items
    )
    assert advisory_ctx["applies"] is True
    downgraded, meta = _downgrade_advisory_llm_findings(llm_items, advisory_ctx=advisory_ctx)
    assert meta is not None and meta["downgraded_count"] == 3
    assert all(item["severity"] == "low" for item in downgraded)


def test_advisory_detector_matches_recent_usaid_llm_wording_variants():
    msgs = [
        (
            "Weak causal link between Outputs and IRs, particularly in explaining how AI training modules lead to "
            "a comprehensive AI training curriculum."
        ),
        (
            "Unrealistic assumption that AI training directly contributes to climate resilience without clear evidence "
            "or logical explanation."
        ),
        "Missing detailed strategies for ensuring gender equality in AI training participation.",
        "Lack of integration of climate resilience as a cross-cutting theme in AI training programs.",
        "Indicators lack baseline and target values, making it difficult to measure progress.",
    ]
    assert all(_is_advisory_llm_message(m) for m in msgs)


def test_llm_finding_classifier_assigns_stable_labels():
    cases = [
        (
            "Indicators lack baseline and target values, making it difficult to measure progress.",
            "logframe",
            "BASELINE_TARGET_MISSING",
        ),
        (
            "Weak causal links between Outputs and IRs: draft lacks a detailed causal explanation.",
            "toc",
            "CAUSAL_LINK_DETAIL",
        ),
        (
            "Missing cross-cutting themes: While gender equality is mentioned, the proposal lacks a detailed plan.",
            "general",
            "CROSS_CUTTING_INTEGRATION",
        ),
        (
            "Theory of Change objectives are not sufficiently specific and measurable.",
            "toc",
            "OBJECTIVE_SPECIFICITY",
        ),
    ]
    for message, section, expected in cases:
        assert _classify_llm_finding_label(message, section=section) == expected


def test_advisory_downgrade_uses_label_when_message_wording_varies():
    advisory_ctx = {"applies": True, "reason": "test", "architect_threshold_hit_rate": 0.75}
    llm_items = [
        {
            "code": "LLM_REVIEW_FLAG_1",
            "label": "CROSS_CUTTING_INTEGRATION",
            "severity": "medium",
            "section": "general",
            "message": "Coverage of gender and climate dimensions could be improved.",
            "source": "llm",
        },
        {
            "code": "LLM_REVIEW_FLAG_2",
            "label": "GENERIC_LLM_REVIEW_FLAG",
            "severity": "medium",
            "section": "general",
            "message": "Completely custom wording that should remain medium.",
            "source": "llm",
        },
    ]

    downgraded, meta = _downgrade_advisory_llm_findings(llm_items, advisory_ctx=advisory_ctx)
    assert meta is not None and meta["applied"] is True
    assert meta["downgraded_count"] == 1
    assert "CROSS_CUTTING_INTEGRATION" in (meta.get("labels_downgraded") or [])
    assert downgraded[0]["severity"] == "low"
    assert downgraded[1]["severity"] == "medium"


def test_llm_finding_policy_supports_donor_specific_overrides(monkeypatch):
    item = {"label": "CROSS_CUTTING_INTEGRATION", "message": "Placeholder wording"}
    assert critic_llm_policy.llm_finding_policy_class(item, donor_id="usaid") == "advisory"
    monkeypatch.setitem(
        critic_llm_policy.LLM_FINDING_LABEL_DONOR_POLICY_OVERRIDES,
        "usaid",
        {"CROSS_CUTTING_INTEGRATION": "default"},
    )
    assert critic_llm_policy.llm_finding_policy_class(item, donor_id="usaid") == "default"


def test_worldbank_default_override_keeps_baseline_target_missing_non_advisory():
    item = {"label": "BASELINE_TARGET_MISSING", "message": "Indicators lack baseline and target values."}
    assert critic_llm_policy.llm_finding_policy_class(item, donor_id="worldbank") == "default"
    assert critic_llm_policy.llm_finding_policy_class(item, donor_id="usaid") == "advisory"


def test_usaid_generic_toc_and_logframe_labels_are_advisory():
    toc_item = {"label": "GENERIC_TOC_REVIEW_FLAG", "message": "General TOC refinement note."}
    logframe_item = {"label": "GENERIC_LOGFRAME_REVIEW_FLAG", "message": "General logframe refinement note."}
    assert critic_llm_policy.llm_finding_policy_class(toc_item, donor_id="usaid") == "advisory"
    assert critic_llm_policy.llm_finding_policy_class(logframe_item, donor_id="usaid") == "advisory"
    assert critic_llm_policy.llm_finding_policy_class(toc_item, donor_id="worldbank") == "default"


def test_advisory_llm_findings_allow_good_enough_architect_grounding_without_fallback():
    class _RuleCheck:
        def __init__(self, status: str):
            self.status = status

    class _RuleReport:
        def __init__(self):
            self.fatal_flaws = []
            self.checks = [_RuleCheck("pass")]

    state = {
        "citations": [
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.55},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.52},
            {"stage": "architect", "citation_type": "rag_claim_support", "citation_confidence": 0.48},
            {"stage": "architect", "citation_type": "rag_low_confidence", "citation_confidence": 0.28},
            {"stage": "architect", "citation_type": "rag_low_confidence", "citation_confidence": 0.27},
            {"stage": "mel", "citation_type": "rag_result", "citation_confidence": 0.8},
        ]
    }
    llm_items = [
        {
            "code": "LLM_REVIEW_FLAG_1",
            "severity": "medium",
            "section": "toc",
            "message": "Weak causal links between Outputs and IRs: draft lacks a detailed causal explanation.",
            "source": "llm",
        }
    ]

    advisory_ctx = _advisory_llm_findings_context(
        state=state, rule_report=_RuleReport(), llm_fatal_flaw_items=llm_items
    )
    assert advisory_ctx["applies"] is True
    assert advisory_ctx["architect_threshold_hit_rate"] == 0.6
    assert advisory_ctx["architect_rag_low_ratio"] == 0.4
