from __future__ import annotations

from grantflow.swarm.nodes.architect_generation import (
    _architect_grounded_confidence_bonus,
    _extract_claim_strings,
    build_architect_claim_citations,
    extract_architect_claim_records,
    summarize_architect_claim_citations,
)
from grantflow.swarm.nodes.architect_policy import architect_claim_confidence_threshold


def test_architect_claim_citation_policy_marks_low_confidence_hits():
    toc_payload = {"project_goal": "Improve water sanitation outcomes", "objectives": []}
    citations = build_architect_claim_citations(
        toc_payload=toc_payload,
        namespace="usaid_ads201",
        donor_id="usaid",
        evidence_hits=[
            {
                "rank": 1,
                "retrieval_rank": 1,
                "retrieval_confidence": 0.62,
                "doc_id": "usaid_ads201_p3_c0",
                "chunk_id": "usaid_ads201_p3_c0",
                "label": "Unrelated guidance",
                "source": "a.pdf",
                "excerpt": "Procurement templates and financial reporting annexes only",
            }
        ],
    )
    assert citations
    assert citations[0]["citation_type"] in {"rag_low_confidence", "fallback_namespace", "rag_claim_support"}
    assert "citation_confidence" in citations[0]
    assert "confidence_threshold" in citations[0]
    assert "doc_id" in citations[0]
    assert "retrieval_rank" in citations[0]
    assert "retrieval_confidence" in citations[0]
    assert "evidence_signal" in citations[0]
    assert "review_hint" in citations[0]
    assert 0.0 <= float(citations[0]["citation_confidence"]) <= 1.0
    assert 0.0 < float(citations[0]["confidence_threshold"]) < 1.0
    if citations[0]["citation_type"] == "rag_low_confidence":
        assert float(citations[0]["citation_confidence"]) < float(citations[0]["confidence_threshold"])


def test_architect_claim_citations_without_hits_emit_per_claim_fallback_records():
    toc_payload = {
        "project_goal": "Improve water sanitation outcomes",
        "development_objectives": [
            {"description": "Improve governance and service quality"},
            {"description": "Improve coverage and reliability"},
        ],
    }
    citations = build_architect_claim_citations(
        toc_payload=toc_payload,
        namespace="usaid_ads201",
        donor_id="usaid",
        evidence_hits=[],
    )
    assert len(citations) >= 3
    assert all(c["citation_type"] == "fallback_namespace" for c in citations)
    assert all(c.get("used_for") == "toc_claim" for c in citations)
    assert all(str(c.get("statement_path") or "").strip() for c in citations)
    assert all(c.get("result_level") in {"impact", "outcome", "output", "general"} for c in citations)
    assert all(c.get("traceability_status") == "missing" for c in citations)
    assert all(c.get("traceability_complete") is False for c in citations)
    assert all(c.get("evidence_signal") == "fallback namespace only" for c in citations)
    assert all(str(c.get("review_hint") or "").strip() for c in citations)
    assert all("no retrieved evidence" in str(c.get("label") or "").lower() for c in citations)
    assert all(0.0 < float(c.get("confidence_threshold") or 0.0) < 1.0 for c in citations)


def test_architect_claim_citations_without_hits_use_strategy_reference_when_retrieval_disabled():
    toc_payload = {
        "project_goal": "Improve water sanitation outcomes",
        "development_objectives": [
            {"description": "Improve governance and service quality"},
            {"description": "Improve coverage and reliability"},
        ],
    }
    citations = build_architect_claim_citations(
        toc_payload=toc_payload,
        namespace="usaid_ads201",
        donor_id="usaid",
        evidence_hits=[],
        retrieval_expected=False,
    )
    assert len(citations) >= 3
    assert all(c["citation_type"] == "strategy_reference" for c in citations)
    assert all(c.get("used_for") == "toc_claim" for c in citations)
    assert all(str(c.get("statement_path") or "").strip() for c in citations)
    assert all(c.get("result_level") in {"impact", "outcome", "output", "general"} for c in citations)
    assert all(c.get("traceability_status") == "complete" for c in citations)
    assert all(c.get("traceability_complete") is True for c in citations)
    assert all(c.get("evidence_signal") == "strategy reference" for c in citations)
    assert all(str(c.get("doc_id") or "").startswith("strategy::") for c in citations)
    assert all(str(c.get("source") or "").startswith("strategy::") for c in citations)
    assert all(float(c.get("citation_confidence") or 0.0) >= 0.7 for c in citations)


def test_architect_claim_support_requires_traceable_hit_even_when_overlap_is_high():
    statement = "Improve water sanitation outcomes in Kenya through community systems"
    toc_payload = {"project_goal": statement}
    citations = build_architect_claim_citations(
        toc_payload=toc_payload,
        namespace="usaid_ads201",
        donor_id="usaid",
        evidence_hits=[
            {
                "rank": 1,
                "retrieval_rank": 1,
                "label": "Untyped fragment",
                "excerpt": statement,
                "source": None,
                "doc_id": None,
                "chunk_id": None,
            }
        ],
    )
    assert citations
    citation = citations[0]
    assert citation["citation_type"] == "rag_low_confidence"
    assert citation["traceability_complete"] is False
    assert citation["traceability_status"] == "missing"


def test_architect_grounded_confidence_bonus_rewards_complete_rank1_hits():
    bonus = _architect_grounded_confidence_bonus(
        {
            "rank": 1,
            "retrieval_rank": 1,
            "retrieval_confidence": 0.87,
            "doc_id": "doc-1",
            "chunk_id": "chunk-1",
        },
        traceability_status="complete",
    )
    assert 0.0 < bonus <= 0.06


def test_extract_claim_strings_skips_identifier_fields():
    toc_payload = {
        "project_goal": "Improve services",
        "development_objectives": [
            {
                "do_id": "DO 1",
                "description": "Institutional capacity improves",
                "intermediate_results": [{"ir_id": "IR 1", "description": "Teams apply SOPs"}],
            }
        ],
    }
    claims = _extract_claim_strings(toc_payload, "toc")
    paths = [p for p, _ in claims]
    assert "toc.development_objectives[0].do_id" not in paths
    assert "toc.development_objectives[0].intermediate_results[0].ir_id" not in paths
    assert "toc.development_objectives[0].description" in paths


def test_extract_claim_strings_skips_nested_indicator_fields():
    toc_payload = {
        "development_objectives": [
            {
                "intermediate_results": [
                    {
                        "description": "Teams improve delivery practices",
                        "outputs": [
                            {
                                "description": "Training package deployed",
                                "indicators": [
                                    {
                                        "name": "Number of staff trained",
                                        "target": "500",
                                        "justification": "Tracks coverage",
                                        "citation": "usaid",
                                    }
                                ],
                            }
                        ],
                    }
                ]
            }
        ]
    }
    claims = _extract_claim_strings(toc_payload, "toc")
    paths = [p for p, _ in claims]
    assert "toc.development_objectives[0].intermediate_results[0].description" in paths
    assert "toc.development_objectives[0].intermediate_results[0].outputs[0].description" in paths
    assert not any(".indicators[" in p for p in paths)


def test_extract_claim_strings_skips_assumptions_and_risks():
    toc_payload = {
        "project_goal": "Improve service delivery",
        "critical_assumptions": ["Leadership support remains stable"],
        "assumptions": ["Local adoption remains strong"],
        "risks": ["Policy turnover delays adoption"],
    }
    claims = _extract_claim_strings(toc_payload, "toc")
    paths = [p for p, _ in claims]
    assert "toc.project_goal" in paths
    assert "toc.critical_assumptions[0]" not in paths
    assert "toc.assumptions[0]" not in paths
    assert "toc.risks[0]" not in paths


def test_extract_architect_claim_records_prioritizes_key_objective_paths():
    toc_payload = {
        "project_goal": "Improve access",
        "objectives": [
            {"objective_id": "OBJ-1", "title": "Title 1", "description": "Description 1"},
            {"objective_id": "OBJ-2", "title": "Title 2", "description": "Description 2"},
        ],
        "assumptions": ["Assumption 1"],
    }
    records = extract_architect_claim_records(toc_payload, max_claims=5, max_high_priority_claims=5)
    assert records
    paths = [str(r.get("statement_path") or "") for r in records]
    assert "toc.project_goal" in paths
    assert any(".objectives[0].description" in p or ".objectives[1].description" in p for p in paths)
    assert all(".assumptions[" not in p for p in paths)
    assert int(records[0].get("priority") or 0) >= int(records[-1].get("priority") or 0)


def test_summarize_architect_claim_citations_reports_coverage_ratios():
    claim_records = [
        {"statement_path": "toc.project_goal", "statement": "Goal", "priority": 5},
        {"statement_path": "toc.objectives[0].description", "statement": "Obj", "priority": 4},
    ]
    citations = [
        {
            "stage": "architect",
            "used_for": "toc_claim",
            "statement_path": "toc.project_goal",
            "citation_type": "rag_claim_support",
        },
        {
            "stage": "architect",
            "used_for": "toc_claim",
            "statement_path": "toc.objectives[0].description",
            "citation_type": "fallback_namespace",
        },
    ]
    stats = summarize_architect_claim_citations(claim_records=claim_records, citations=citations)
    assert stats["claims_total"] == 2
    assert stats["key_claims_total"] == 2
    assert stats["claim_paths_covered"] == 2
    assert stats["confident_claim_paths_covered"] == 1
    assert stats["fallback_claim_count"] == 1
    assert stats["key_claim_coverage_ratio"] == 1.0
    assert stats["fallback_claim_ratio"] == 0.5


def test_architect_claim_threshold_is_tuned_by_donor_and_section():
    default_threshold = architect_claim_confidence_threshold(donor_id="unknown", statement_path="toc.project_goal")
    usaid_goal_threshold = architect_claim_confidence_threshold(donor_id="usaid", statement_path="toc.project_goal")
    usaid_assumption_threshold = architect_claim_confidence_threshold(
        donor_id="usaid", statement_path="toc.critical_assumptions[0]"
    )

    assert default_threshold >= 0.35
    assert usaid_goal_threshold == 0.32
    assert usaid_assumption_threshold > usaid_goal_threshold


def test_summarize_architect_claim_citations_includes_traceability_and_threshold_stats():
    claim_records = [
        {"statement_path": "toc.project_goal", "statement": "Goal", "priority": 5},
        {"statement_path": "toc.objectives[0].description", "statement": "Objective", "priority": 4},
    ]
    citations = [
        {
            "stage": "architect",
            "used_for": "toc_claim",
            "statement_path": "toc.project_goal",
            "citation_type": "rag_claim_support",
            "traceability_status": "complete",
            "citation_confidence": 0.82,
            "confidence_threshold": 0.6,
        },
        {
            "stage": "architect",
            "used_for": "toc_claim",
            "statement_path": "toc.objectives[0].description",
            "citation_type": "rag_low_confidence",
            "traceability_status": "partial",
            "citation_confidence": 0.31,
            "confidence_threshold": 0.5,
        },
    ]
    stats = summarize_architect_claim_citations(claim_records=claim_records, citations=citations)
    assert stats["traceability_complete_citation_count"] == 1
    assert stats["traceability_partial_citation_count"] == 1
    assert stats["traceability_missing_citation_count"] == 0
    assert stats["traceability_gap_citation_count"] == 1
    assert stats["traceability_gap_rate"] == 0.5
    assert stats["threshold_considered_count"] == 2
    assert stats["threshold_hit_count"] == 1
    assert stats["threshold_hit_rate"] == 0.5
