from __future__ import annotations

from grantflow.swarm.nodes import mel_specialist as mel_module
from grantflow.tests.mel_helpers import _base_state


def test_mel_grounded_confidence_bonus_rewards_complete_rank1_hits():
    bonus = mel_module._mel_grounded_confidence_bonus(
        {
            "rank": 1,
            "retrieval_rank": 1,
            "retrieval_confidence": 0.46,
            "doc_id": "doc-1",
            "chunk_id": "chunk-1",
        },
        traceability_status="complete",
    )
    assert 0.0 < bonus <= 0.1


def test_mel_deterministic_mode_generates_logframe_and_citations(monkeypatch):
    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        return {
            "documents": [["USAID indicator guidance for water service quality"]],
            "metadatas": [
                [
                    {
                        "doc_id": "usaid_ads201_p12_c0",
                        "chunk_id": "usaid_ads201_p12_c0",
                        "page": 12,
                        "indicator_code": "EG.3.2-27",
                        "frequency": "quarterly",
                        "formula": "(Numerator / Denominator) * 100",
                        "data_source": "PMP indicator tracking dataset",
                    }
                ]
            ],
            "ids": [["usaid_ads201_p12_c0"]],
            "distances": [[0.05]],
        }

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    state = _base_state(llm_mode=False)

    out = mel_module.mel_assign_indicators(state)
    indicators = out["logframe_draft"]["indicators"]
    citations = [c for c in (out.get("citations") or []) if isinstance(c, dict) and c.get("stage") == "mel"]
    generation = out.get("mel_generation_meta") or {}

    assert indicators
    assert citations
    assert generation.get("llm_requested") is False
    assert generation.get("llm_used") is False
    assert str(generation.get("engine") or "").startswith("deterministic:")
    assert out["logframe_draft"]["rag_trace"]["used_results"] == 1
    assert out["logframe_draft"]["rag_trace"]["namespace_normalized"] == "usaid_ads201"
    assert out["logframe_draft"]["rag_trace"]["hits"][0]["traceability_status"] in {"complete", "partial", "missing"}
    assert "traceability_counts" in out["logframe_draft"]["rag_trace"]
    assert citations[0]["citation_type"] in {"rag_result", "rag_low_confidence"}
    assert citations[0]["namespace_normalized"] == "usaid_ads201"
    assert citations[0]["used_for"]
    assert str(citations[0]["toc_statement_path"]).startswith("toc.")
    assert str(citations[0]["statement_path"]).startswith("toc.")
    assert str(citations[0]["result_level"]) in {"impact", "outcome", "output"}
    assert indicators[0]["indicator_code"] == "EG.3.2-27"
    assert indicators[0]["frequency"] == "quarterly"
    assert indicators[0]["formula"] == "(Numerator / Denominator) * 100"
    assert indicators[0]["data_source"] == "PMP indicator tracking dataset"
    assert out.get("draft_versions")


def test_mel_llm_mode_without_api_key_uses_emergency_fallback(monkeypatch):
    monkeypatch.setattr(mel_module, "openai_compatible_llm_available", lambda: False)

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

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    state = _base_state(llm_mode=True)

    out = mel_module.mel_assign_indicators(state)
    indicators = out["logframe_draft"]["indicators"]
    generation = out.get("mel_generation_meta") or {}
    citations = [c for c in (out.get("citations") or []) if isinstance(c, dict) and c.get("stage") == "mel"]

    assert indicators
    assert indicators[0]["indicator_id"] == "IND_001"
    assert generation.get("llm_requested") is True
    assert generation.get("llm_used") is False
    assert generation.get("fallback_used") is True
    assert generation.get("fallback_class") == "emergency"
    assert "OPENAI_API_KEY / OPENROUTER_API_KEY missing" in str(generation.get("llm_fallback_reason") or "")
    assert citations
    assert citations[0]["citation_type"] == "fallback_namespace"


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

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    state = _base_state(llm_mode=False)

    out = mel_module.mel_assign_indicators(state)
    citations = [c for c in (out.get("citations") or []) if isinstance(c, dict) and c.get("stage") == "mel"]
    assert citations
    assert all(c.get("citation_type") == "strategy_reference" 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(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)
    assert all(str(c.get("toc_statement_path") or "").startswith("toc.") for c in citations)
    assert all(str(c.get("statement_path") or "").startswith("toc.") for c in citations)


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

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    state = _base_state(llm_mode=False)
    state["toc_draft"] = {
        "toc": {
            "project_goal": "Improve digital service quality in municipal agencies.",
            "development_objectives": [
                {
                    "description": "Civil servants apply AI-assisted service workflows in daily operations.",
                },
                {
                    "description": "Municipal service turnaround time is reduced for priority services.",
                },
            ],
        }
    }

    out = mel_module.mel_assign_indicators(state)
    indicators = out["logframe_draft"]["indicators"]
    citations = [c for c in (out.get("citations") or []) if isinstance(c, dict) and c.get("stage") == "mel"]
    generation = out.get("mel_generation_meta") or {}

    assert len(indicators) >= 2
    assert all(str(item.get("toc_statement_path") or "").startswith("toc.") for item in indicators)
    assert all(str(item.get("result_level") or "") in {"impact", "outcome", "output"} for item in indicators)
    assert all(str(item.get("frequency") or "") for item in indicators)
    assert all(str(item.get("formula") or "") for item in indicators)
    assert all(str(item.get("definition") or "") for item in indicators)
    assert all(str(item.get("data_source") or "") for item in indicators)
    assert all(isinstance(item.get("disaggregation"), list) for item in indicators)
    assert all(c.get("citation_type") == "strategy_reference" for c in citations)
    assert all(str(c.get("toc_statement_path") or "").startswith("toc.") for c in citations)
    assert all(str(c.get("statement_path") or "").startswith("toc.") for c in citations)
    assert all(str(c.get("result_level") or "") in {"impact", "outcome", "output"} for c in citations)
    assert generation.get("engine") == "deterministic:toc_results_template"
    assert generation.get("deterministic_source") == "toc_results_template"


def test_mel_query_variants_include_eu_and_worldbank_specific_terms():
    eu_state = {
        "donor_id": "eu",
        "input_context": {"project": "Digital Governance", "country": "Moldova"},
        "toc_draft": {
            "toc": {
                "overall_objective": {"title": "Improve digital governance performance"},
                "specific_objectives": [{"title": "Strengthen institutional capacity"}],
                "expected_outcomes": [{"expected_change": "Institutions demonstrate measurable improvements"}],
            }
        },
    }
    wb_state = {
        "donor_id": "worldbank",
        "input_context": {"project": "Public Sector Performance", "country": "Uzbekistan"},
        "toc_draft": {
            "toc": {
                "project_development_objective": "Improve service delivery performance",
                "objectives": [{"title": "Strengthen institutional performance"}],
                "results_chain": [{"description": "Agencies implement workflow improvements"}],
            }
        },
    }

    eu_query = mel_module._build_query_text(eu_state)
    wb_query = mel_module._build_query_text(wb_state)
    eu_variants = mel_module._query_variants(eu_state, eu_query, max_variants=6)
    wb_variants = mel_module._query_variants(wb_state, wb_query, max_variants=6)

    eu_combined = " || ".join(eu_variants).lower()
    wb_combined = " || ".join(wb_variants).lower()
    assert "means of verification" in eu_combined
    assert "specific objectives" in eu_combined
    assert "expected outcomes" in eu_combined
    assert "pdo indicators" in wb_combined
    assert "intermediate results indicators" in wb_combined
    assert "implementation status results report" in wb_combined
