from __future__ import annotations

from grantflow.swarm.nodes.critic import _citation_grounding_context, _combine_critic_scores


def test_citation_grounding_context_tracks_fallback_and_weak_grounding():
    state = {
        "architect_retrieval": {"enabled": True, "hits_count": 0},
        "citations": [
            {"citation_type": "fallback_namespace", "citation_confidence": 0.1},
            {"citation_type": "fallback_namespace", "citation_confidence": 0.1},
            {"citation_type": "rag_low_confidence", "citation_confidence": 0.2},
            {"citation_type": "rag_low_confidence", "citation_confidence": 0.25},
            {"citation_type": "rag_claim_support", "citation_confidence": 0.9},
        ],
    }

    ctx = _citation_grounding_context(state)
    assert ctx["citation_count"] == 5
    assert ctx["fallback_namespace_citation_count"] == 2
    assert ctx["strategy_reference_citation_count"] == 0
    assert ctx["retrieval_grounded_citation_count"] == 3
    assert ctx["non_retrieval_citation_count"] == 2
    assert ctx["rag_low_confidence_citation_count"] == 2
    assert ctx["low_confidence_citation_count"] == 4
    assert ctx["traceability_complete_citation_count"] == 0
    assert ctx["traceability_partial_citation_count"] == 0
    assert ctx["traceability_missing_citation_count"] == 5
    assert ctx["traceability_gap_citation_count"] == 5
    assert ctx["traceability_gap_ratio"] == 1.0
    assert ctx["architect_retrieval_hits_count"] == 0
    assert ctx["weak_grounding"] is True
    assert "architect_retrieval_no_hits" in ctx["weak_grounding_reasons"]
    assert "citation_traceability_gaps_dominate" in ctx["weak_grounding_reasons"]


def test_citation_grounding_context_flags_non_retrieval_dominance_when_retrieval_expected():
    state = {
        "architect_retrieval": {"enabled": True, "hits_count": 2},
        "citations": [
            {"citation_type": "strategy_reference", "citation_confidence": 0.75},
            {"citation_type": "strategy_reference", "citation_confidence": 0.75},
            {"citation_type": "strategy_reference", "citation_confidence": 0.75},
            {"citation_type": "strategy_reference", "citation_confidence": 0.75},
            {"citation_type": "rag_claim_support", "citation_confidence": 0.9},
        ],
    }

    ctx = _citation_grounding_context(state)
    assert ctx["strategy_reference_citation_count"] == 4
    assert ctx["non_retrieval_citation_count"] == 4
    assert ctx["retrieval_grounded_citation_count"] == 1
    assert ctx["weak_grounding"] is True
    assert "non_retrieval_citations_dominate_when_retrieval_enabled" in ctx["weak_grounding_reasons"]


def test_citation_grounding_context_flags_low_retrieval_metadata_completeness():
    state = {
        "architect_retrieval": {"enabled": True, "hits_count": 5},
        "citations": [
            {
                "citation_type": "rag_claim_support",
                "doc_id": "doc-1",
                "retrieval_rank": 1,
                "retrieval_confidence": 0.9,
                "citation_confidence": 0.9,
            },
            {"citation_type": "rag_claim_support", "doc_id": "doc-2", "citation_confidence": 0.9},
            {"citation_type": "rag_claim_support", "doc_id": "doc-3", "citation_confidence": 0.9},
            {"citation_type": "rag_claim_support", "doc_id": "doc-4", "citation_confidence": 0.9},
            {"citation_type": "rag_claim_support", "doc_id": "doc-5", "citation_confidence": 0.9},
        ],
    }
    ctx = _citation_grounding_context(state)
    assert ctx["retrieval_grounded_citation_count"] == 5
    assert ctx["retrieval_metadata_complete_citation_count"] == 1
    assert ctx["retrieval_metadata_complete_rate"] == 0.2
    assert ctx["weak_grounding"] is True
    assert "retrieval_metadata_completeness_below_min" in ctx["weak_grounding_reasons"]


def test_combine_critic_scores_caps_llm_penalty_in_weak_grounding_context():
    state = {
        "architect_retrieval": {"enabled": True, "hits_count": 0},
        "citations": [
            {"citation_type": "fallback_namespace", "citation_confidence": 0.1},
            {"citation_type": "fallback_namespace", "citation_confidence": 0.1},
            {"citation_type": "rag_low_confidence", "citation_confidence": 0.2},
            {"citation_type": "rag_low_confidence", "citation_confidence": 0.2},
            {"citation_type": "rag_low_confidence", "citation_confidence": 0.2},
        ],
    }

    score, meta = _combine_critic_scores(rule_score=9.25, llm_score=3.0, state=state)
    assert score == 7.75  # 9.25 - capped 1.5 penalty
    assert meta is not None
    assert meta["applied"] is True
    assert meta["raw_llm_score"] == 3.0
    assert meta["calibrated_llm_score"] == 7.75


def test_combine_critic_scores_keeps_strict_min_when_grounding_is_not_weak():
    state = {
        "architect_retrieval": {"enabled": True, "hits_count": 4},
        "citations": [
            {
                "citation_type": "rag_claim_support",
                "citation_confidence": 0.92,
                "doc_id": "doc-1",
                "retrieval_rank": 1,
            },
            {
                "citation_type": "rag_claim_support",
                "citation_confidence": 0.88,
                "doc_id": "doc-2",
                "retrieval_rank": 2,
            },
        ],
    }

    score, meta = _combine_critic_scores(rule_score=9.25, llm_score=5.0, state=state)
    assert score == 5.0
    assert meta is not None
    assert meta["applied"] is False
