from __future__ import annotations

from grantflow.swarm.critic_rules import evaluate_rule_based_critic


def test_rule_based_critic_applies_usaid_donor_specific_checks():
    state = {
        "donor_id": "usaid",
        "draft_versions": [
            {"version_id": "toc_v1", "sequence": 1, "section": "toc", "content": {}},
            {"version_id": "logframe_v1", "sequence": 2, "section": "logframe", "content": {}},
        ],
        "toc_validation": {"valid": True, "errors": [], "schema_name": "USAIDTOC"},
        "toc_draft": {"toc": {"project_goal": "Improve services"}},
        "logframe_draft": {"indicators": [{"indicator_id": "IND_001"}]},
        "citations": [
            {"stage": "architect", "statement_path": "toc.project_goal", "label": "doc", "used_for": "toc_claim"},
            {"stage": "mel", "label": "doc", "used_for": "IND_001"},
        ],
    }

    report = evaluate_rule_based_critic(state)
    checks = [c.model_dump() if hasattr(c, "model_dump") else c.dict() for c in report.checks]
    flaws = [f.model_dump() if hasattr(f, "model_dump") else f.dict() for f in report.fatal_flaws]

    assert any(c["code"] == "USAID_DO_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "USAID_CRITICAL_ASSUMPTIONS_PRESENT" and c["status"] == "warn" for c in checks)
    assert any(f["code"] == "USAID_DO_MISSING" and f["section"] == "toc" for f in flaws)
    assert any(f["code"] == "USAID_ASSUMPTIONS_MISSING" and f["section"] == "toc" for f in flaws)


def test_rule_based_critic_applies_katch_evaluation_rfq_checks():
    state = {
        "donor_id": "un_agencies",
        "input_context": {"proposal_mode": "evaluation_rfq", "rfq_profile": "katch_final_assessment"},
        "draft_versions": [
            {"version_id": "toc_v1", "sequence": 1, "section": "toc", "content": {}},
            {"version_id": "logframe_v1", "sequence": 2, "section": "logframe", "content": {}},
        ],
        "toc_validation": {"valid": True, "errors": [], "schema_name": "EvaluationRFQTOC"},
        "toc_draft": {
            "proposal_mode": "evaluation_rfq",
            "rfq_profile": "katch_final_assessment",
            "toc": {
                "proposal_mode": "evaluation_rfq",
                "rfq_profile": "katch_final_assessment",
                "brief": "Technical response for KATCH.",
                "evaluation_purpose": "Assess project performance.",
                "methodology_overview": "Mixed methods design.",
                "deliverables": [{"deliverable": "Draft Evaluation Report", "timing": "Week 4", "purpose": "Draft"}],
                "team_composition": [{"role": "Team Lead", "responsibility": "Lead the assignment"}],
            },
        },
        "logframe_draft": {"indicators": [{"indicator_id": "IND_001"}]},
        "citations": [
            {"stage": "architect", "statement_path": "toc.brief", "label": "doc", "used_for": "toc_claim"},
            {"stage": "mel", "label": "doc", "used_for": "IND_001"},
        ],
    }

    report = evaluate_rule_based_critic(state)
    checks = [c.model_dump() if hasattr(c, "model_dump") else c.dict() for c in report.checks]
    flaws = [f.model_dump() if hasattr(f, "model_dump") else f.dict() for f in report.fatal_flaws]

    assert any(c["code"] == "KATCH_ORGANIZATION_INFORMATION_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_METHOD_COMPONENTS_COMPLETE" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_DELIVERABLE_PLAN_COMPLETE" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_COMPLIANCE_MATRIX_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_KEY_PERSONNEL_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_FINANCIAL_SUMMARY_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_COST_STRUCTURE_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_SUBMISSION_PACKAGE_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_ATTACHMENT_MANIFEST_PRESENT" and c["status"] == "warn" for c in checks)
    assert any(c["code"] == "KATCH_EVALUATION_QUESTION_MATRIX_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_METHODS_COVERAGE_MATRIX_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(c["code"] == "KATCH_DELIVERABLES_SCHEDULE_PRESENT" and c["status"] == "fail" for c in checks)
    assert any(f["code"] == "KATCH_TEAM_COMPOSITION_MISSING" for f in flaws)
    assert any(f["code"] == "KATCH_COMPLIANCE_MATRIX_MISSING" for f in flaws)
    assert any(f["code"] == "KATCH_KEY_PERSONNEL_MISSING" for f in flaws)
    assert any(f["code"] == "KATCH_COST_STRUCTURE_MISSING" for f in flaws)
    assert any(f["code"] == "KATCH_SUBMISSION_PACKAGE_MISSING" for f in flaws)
    assert any(f["code"] == "KATCH_ATTACHMENT_MANIFEST_WEAK" for f in flaws)
    assert any(f["code"] == "KATCH_EVALUATION_QUESTION_MATRIX_MISSING" for f in flaws)
    assert any(f["code"] == "KATCH_METHODS_COVERAGE_MATRIX_MISSING" for f in flaws)
    assert any(f["code"] == "KATCH_DELIVERABLES_SCHEDULE_MISSING" for f in flaws)


def test_rule_based_critic_uses_donor_specific_language_for_grounding_and_measurement():
    state = {
        "donor_id": "worldbank",
        "architect_rag_enabled": True,
        "input_context": {"project": "Public sector performance", "country": "Uzbekistan"},
        "draft_versions": [
            {"version_id": "toc_v1", "sequence": 1, "section": "toc", "content": {}},
            {"version_id": "logframe_v1", "sequence": 2, "section": "logframe", "content": {}},
        ],
        "toc_validation": {"valid": True, "errors": [], "schema_name": "WorldBankTOC"},
        "toc_draft": {
            "toc": {
                "project_development_objective": "Improve service delivery outcomes.",
                "objectives": [
                    {"title": "Objective 1", "description": "Strengthen local service performance."},
                    {"title": "Objective 2", "description": "Improve planning quality."},
                ],
            }
        },
        "logframe_draft": {
            "indicators": [
                {"indicator_id": "IND_001", "baseline": "TBD", "target": "TBD"},
            ]
        },
        "citations": [
            {
                "stage": "architect",
                "used_for": "toc_claim",
                "statement_path": "toc.project_development_objective",
                "citation_type": "fallback_namespace",
                "citation_confidence": 0.1,
                "confidence_threshold": 0.4,
            },
            {
                "stage": "architect",
                "used_for": "toc_claim",
                "statement_path": "toc.objectives[0].description",
                "citation_type": "fallback_namespace",
                "citation_confidence": 0.1,
                "confidence_threshold": 0.4,
            },
            {
                "stage": "architect",
                "used_for": "toc_claim",
                "statement_path": "toc.objectives[1].description",
                "citation_type": "fallback_namespace",
                "citation_confidence": 0.1,
                "confidence_threshold": 0.4,
            },
        ],
        "toc_generation_meta": {
            "claim_coverage": {
                "key_claim_coverage_ratio": 1.0,
                "fallback_claim_ratio": 1.0,
            }
        },
    }

    report = evaluate_rule_based_critic(state)
    flaws = [f.model_dump() if hasattr(f, "model_dump") else f.dict() for f in report.fatal_flaws]
    by_code = {f["code"]: f for f in flaws}

    assert (
        "world bank results framework and pdo package"
        in by_code["TOC_CLAIM_GROUNDING_FALLBACK_DOMINANT"]["message"].lower()
    )
    assert (
        "results-framework and isr-style review"
        in by_code["LOGFRAME_BASELINE_TARGET_PLACEHOLDERS_CRITICAL"]["message"].lower()
    )


def test_rule_based_critic_uses_usaid_specific_traceability_language():
    state = {
        "donor_id": "usaid",
        "architect_rag_enabled": True,
        "input_context": {"project": "Civil service AI skills", "country": "Kazakhstan"},
        "draft_versions": [
            {"version_id": "toc_v1", "sequence": 1, "section": "toc", "content": {}},
            {"version_id": "logframe_v1", "sequence": 2, "section": "logframe", "content": {}},
        ],
        "toc_validation": {"valid": True, "errors": [], "schema_name": "USAIDTOC"},
        "toc_draft": {
            "toc": {
                "project_goal": "Improve responsible AI skills in civil service.",
                "objectives": [
                    {"title": "Objective 1", "description": "Strengthen institutional guidance."},
                    {"title": "Objective 2", "description": "Improve training quality."},
                ],
            }
        },
        "logframe_draft": {"indicators": [{"indicator_id": "IND_001", "baseline": "0", "target": "40"}]},
        "citations": [
            {
                "stage": "architect",
                "used_for": "toc_claim",
                "statement_path": "toc.project_goal",
                "citation_type": "rag_claim_support",
                "citation_confidence": 0.8,
                "confidence_threshold": 0.4,
                "doc_id": "doc-1",
            },
            {
                "stage": "architect",
                "used_for": "toc_claim",
                "statement_path": "toc.objectives[0].description",
                "citation_type": "rag_claim_support",
                "citation_confidence": 0.8,
                "confidence_threshold": 0.4,
                "source": "policy.pdf",
            },
            {
                "stage": "architect",
                "used_for": "toc_claim",
                "statement_path": "toc.objectives[1].description",
                "citation_type": "rag_claim_support",
                "citation_confidence": 0.8,
                "confidence_threshold": 0.4,
            },
        ],
    }

    report = evaluate_rule_based_critic(state)
    flaws = [f.model_dump() if hasattr(f, "model_dump") else f.dict() for f in report.fatal_flaws]
    by_code = {f["code"]: f for f in flaws}

    assert (
        "usaid results hierarchy and monitoring package"
        in by_code["TOC_CLAIM_TRACEABILITY_GAP_CRITICAL"]["message"].lower()
    )


def test_rule_based_critic_uses_donor_specific_schema_language():
    state = {
        "donor_id": "eu",
        "input_context": {"project": "Digital governance", "country": "Moldova"},
        "draft_versions": [
            {"version_id": "toc_v1", "sequence": 1, "section": "toc", "content": {}},
            {"version_id": "logframe_v1", "sequence": 2, "section": "logframe", "content": {}},
        ],
        "toc_validation": {"valid": False, "errors": ["missing field"], "schema_name": "EUTOC"},
        "toc_draft": {"toc": {"overall_objective": {"title": "Objective"}}},
        "logframe_draft": {"indicators": []},
        "citations": [],
    }

    report = evaluate_rule_based_critic(state)
    flaws = [f.model_dump() if hasattr(f, "model_dump") else f.dict() for f in report.fatal_flaws]
    by_code = {f["code"]: f for f in flaws}

    assert "eu intervention-logic structure" in by_code["TOC_SCHEMA_INVALID"]["fix_hint"].lower()
    assert "eu intervention logic and verification package" in by_code["TOC_SCHEMA_INVALID"]["message"].lower()
