from __future__ import annotations

from grantflow.swarm.nodes import architect_retrieval as architect_retrieval_module
from grantflow.swarm.nodes.architect_retrieval import (
    pick_best_architect_evidence_hit,
    retrieve_architect_evidence,
    score_architect_evidence_hit,
)


def test_architect_evidence_ranking_prefers_more_relevant_hit():
    statement = "Improve water sanitation outcomes in Kenya through community systems"
    hits = [
        {
            "rank": 1,
            "label": "Generic donor guide",
            "excerpt": "Procurement and reporting templates",
            "source": "a.pdf",
        },
        {
            "rank": 2,
            "label": "WASH technical guidance",
            "excerpt": "Kenya water sanitation outcomes and community systems strengthening guidance",
            "source": "b.pdf",
            "page": 12,
        },
    ]
    best_hit, confidence = pick_best_architect_evidence_hit(statement, hits)
    assert best_hit["source"] == "b.pdf"
    assert confidence > 0.3


def test_architect_evidence_scoring_penalizes_generic_excerpt_and_rewards_worldbank_tokens():
    statement = "Improve public service delivery performance through results framework indicators"
    generic_hit = {
        "rank": 1,
        "label": "Implementation templates",
        "excerpt": "Budget templates procurement reporting compliance annex and finance forms",
        "source": "generic.pdf",
    }
    wb_hit = {
        "rank": 2,
        "label": "World Bank results framework guide",
        "excerpt": "Results framework indicators for public service delivery performance monitoring",
        "source": "wb.pdf",
        "page": 5,
    }
    generic_score = score_architect_evidence_hit(
        statement,
        generic_hit,
        donor_id="worldbank",
        statement_path="toc.results[0].description",
    )
    wb_score = score_architect_evidence_hit(
        statement,
        wb_hit,
        donor_id="worldbank",
        statement_path="toc.results[0].description",
    )
    assert wb_score > generic_score
    best_hit, confidence = pick_best_architect_evidence_hit(
        statement,
        [generic_hit, wb_hit],
        donor_id="worldbank",
        statement_path="toc.results[0].description",
    )
    assert best_hit["source"] == "wb.pdf"
    assert confidence == wb_score


def test_retrieve_architect_evidence_normalizes_traceability_and_deduplicates_hits(monkeypatch):
    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        return {
            "documents": [["Relevant donor paragraph", "Duplicate entry"]],
            "metadatas": [
                [
                    {"source": "guide.pdf", "chunk_id": "ch_1", "page": 2},
                    {"source": "guide.pdf", "chunk_id": "ch_1", "page": 2},
                ]
            ],
            "ids": [["id_1", "id_1"]],
            "distances": [[0.1, 0.2]],
        }

    monkeypatch.setattr(architect_retrieval_module.vector_store, "query", fake_query)

    summary, hits = retrieve_architect_evidence(
        {"donor_id": "usaid", "input_context": {"project": "Water", "country": "Kenya"}},
        "usaid_ads201",
    )
    assert len(hits) == 1
    hit = hits[0]
    assert hit["doc_id"]
    assert hit["chunk_id"] == "ch_1"
    assert hit["traceability_status"] == "complete"
    assert hit["traceability_complete"] is True
    assert hit["namespace_normalized"] == "usaid_ads201"
    assert summary["traceability_counts"]["complete"] == 1
    assert summary["namespace_normalized"] == "usaid_ads201"
    assert summary.get("collection", "").startswith("grantflow_")
    assert summary["hits"][0]["traceability_complete"] is True


def test_architect_query_variants_prioritize_compact_donor_shaped_queries(monkeypatch):
    captured: dict[str, object] = {}

    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        captured["query_texts"] = query_texts
        return {"documents": [[]], "metadatas": [[]], "ids": [[]], "distances": [[]]}

    monkeypatch.setattr(architect_retrieval_module.vector_store, "query", fake_query)

    retrieve_architect_evidence(
        {
            "donor_id": "worldbank",
            "input_context": {
                "project": "Public Sector Performance and Service Delivery Capacity Strengthening",
                "country": "Uzbekistan",
                "sector": "governance",
                "theme": "service_delivery",
            },
            "revision_hint": (
                "Address the following issues before finalizing the draft: "
                "Replace repeated boilerplate. Grounding gate warning: architect_retrieval_no_hits."
            ),
            "toc_draft": {
                "toc": {
                    "project_development_objective": "Improve service delivery performance in Uzbekistan",
                    "results_chain": [{"description": "Agencies implement workflow improvements"}],
                }
            },
        },
        "worldbank_ads301",
    )
    variants = list(captured.get("query_texts") or [])
    assert variants
    assert "architect_retrieval_no_hits" not in variants[0].lower()
    assert "replace repeated boilerplate" not in variants[0].lower()
    assert len(variants[0].split()) < len(variants[-1].split())
    assert "worldbank" in variants[0].lower()
    assert "results framework" in variants[0].lower() or "project development objective" in variants[0].lower()
