from __future__ import annotations

from pydantic import BaseModel, Field

from grantflow.core.strategies.factory import DonorFactory
from grantflow.swarm.nodes import mel_specialist as mel_module
from grantflow.tests.mel_helpers import _base_state


def test_mel_llm_mode_uses_structured_output_when_available(monkeypatch):
    monkeypatch.setattr(mel_module, "openai_compatible_llm_available", lambda: True)

    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        return {
            "documents": [["Official donor indicator guidance excerpt"]],
            "metadatas": [[{"doc_id": "usaid_ads201_p10_c1", "chunk_id": "usaid_ads201_p10_c1", "page": 10}]],
            "ids": [["usaid_ads201_p10_c1"]],
            "distances": [[0.1]],
        }

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    monkeypatch.setattr(
        mel_module,
        "_llm_structured_mel",
        lambda **kwargs: (
            {
                "indicators": [
                    {
                        "indicator_id": "IND_101",
                        "name": "Officials trained on AI governance",
                        "justification": "Directly tracks workforce upskilling objective.",
                        "citation": "USAID ADS 201 p.10",
                        "baseline": "0",
                        "target": "250",
                        "evidence_excerpt": "Official donor indicator guidance excerpt",
                    },
                    {
                        "indicator_id": "IND_102",
                        "name": "Policies updated with AI safeguards",
                        "justification": "Tracks policy institutionalization results.",
                        "citation": "USAID ADS 201 p.10",
                        "baseline": "0",
                        "target": "6",
                        "evidence_excerpt": "Official donor indicator guidance excerpt",
                    },
                ]
            },
            "llm:stub-mel-model",
            None,
        ),
    )

    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 len(indicators) == 2
    assert indicators[0]["indicator_id"] == "IND_101"
    assert indicators[1]["indicator_id"] == "IND_102"
    assert generation.get("llm_requested") is True
    assert generation.get("llm_used") is True
    assert generation.get("fallback_used") is False
    assert generation.get("engine") == "llm:stub-mel-model"
    assert generation.get("schema_contract_hint_present") is True
    assert generation.get("input_context_key_count") == 2
    assert generation.get("retrieval_prompt_hit_count") == 1
    assert generation.get("retrieval_trace_hint_present") is True
    assert len(citations) == 2
    assert all(
        c.get("citation_type") in {"rag_result", "rag_low_confidence", "fallback_namespace", "strategy_reference"}
        for c in citations
    )


def test_mel_llm_prompt_receives_full_input_context_and_schema_contract(monkeypatch):
    monkeypatch.setattr(mel_module, "openai_compatible_llm_available", lambda: True)

    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        return {
            "documents": [["Donor indicator guidance excerpt with baseline and target hints"]],
            "metadatas": [[{"doc_id": "usaid_ads201_p5_c1", "chunk_id": "usaid_ads201_p5_c1", "page": 5}]],
            "ids": [["usaid_ads201_p5_c1"]],
            "distances": [[0.08]],
        }

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    captured: dict[str, object] = {}

    def fake_llm_structured_mel(**kwargs):
        captured["input_context"] = kwargs.get("input_context")
        captured["schema_contract_hint"] = kwargs.get("schema_contract_hint")
        captured["retrieval_trace_hint"] = kwargs.get("retrieval_trace_hint")
        return (
            {
                "indicators": [
                    {
                        "indicator_id": "IND_501",
                        "name": "Civil servants certified on AI policy",
                        "justification": "Tracks capability uptake.",
                        "citation": "usaid_ads201_p5_c1",
                        "baseline": "0",
                        "target": "400",
                        "evidence_excerpt": "Donor indicator guidance excerpt with baseline and target hints",
                    }
                ]
            },
            "llm:stub-mel-model",
            None,
        )

    monkeypatch.setattr(mel_module, "_llm_structured_mel", fake_llm_structured_mel)

    state = _base_state(llm_mode=True)
    state["input_context"] = {
        "project": "AI Civil Service Training",
        "country": "Kazakhstan",
        "sector": "Public Administration",
        "duration_months": 24,
    }
    state["input"] = dict(state["input_context"])
    out = mel_module.mel_assign_indicators(state)
    generation = out.get("mel_generation_meta") or {}

    assert generation.get("llm_used") is True
    assert generation.get("schema_contract_hint_present") is True
    assert generation.get("input_context_key_count") == 4
    assert generation.get("retrieval_prompt_hit_count") == 1
    assert generation.get("retrieval_trace_hint_present") is True
    assert captured.get("input_context") == state["input_context"]
    schema_hint = str(captured.get("schema_contract_hint") or "")
    assert schema_hint
    assert "indicators" in schema_hint
    retrieval_hint = str(captured.get("retrieval_trace_hint") or "")
    assert retrieval_hint
    assert "used_results" in retrieval_hint


def test_mel_llm_mode_uses_strategy_mel_schema_contract(monkeypatch):
    monkeypatch.setattr(mel_module, "openai_compatible_llm_available", lambda: True)

    class DonorSpecificIndicator(BaseModel):
        indicator_id: str
        name: str
        justification: str
        citation: str
        baseline: str = "0"
        target: str = "1"
        evidence_excerpt: str | None = None

    class DonorSpecificMELDraft(BaseModel):
        mel_indicators: list[DonorSpecificIndicator] = Field(default_factory=list)

    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        return {
            "documents": [["Donor indicator guidance excerpt"]],
            "metadatas": [[{"doc_id": "usaid_ads201_p10_c1", "chunk_id": "usaid_ads201_p10_c1", "page": 10}]],
            "ids": [["usaid_ads201_p10_c1"]],
            "distances": [[0.1]],
        }

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    strategy = DonorFactory.get_strategy("usaid")
    monkeypatch.setattr(strategy, "get_mel_schema", lambda: DonorSpecificMELDraft)
    captured: dict[str, object] = {}

    def fake_llm_structured_mel(**kwargs):
        captured["schema_cls"] = kwargs.get("schema_cls")
        return (
            {
                "mel_indicators": [
                    {
                        "indicator_id": "IND_777",
                        "name": "Civil servants trained on AI governance",
                        "justification": "Tracks capability development outcome.",
                        "citation": "usaid_ads201_p10_c1",
                        "baseline": "0",
                        "target": "300",
                        "evidence_excerpt": "Donor indicator guidance excerpt",
                    }
                ]
            },
            "llm:stub-mel-model",
            None,
        )

    monkeypatch.setattr(mel_module, "_llm_structured_mel", fake_llm_structured_mel)

    state = _base_state(llm_mode=True)
    state["strategy"] = strategy
    state["donor_strategy"] = strategy
    out = mel_module.mel_assign_indicators(state)
    generation = out.get("mel_generation_meta") or {}
    indicators = out["logframe_draft"]["indicators"]

    assert captured.get("schema_cls") is DonorSpecificMELDraft
    assert generation.get("schema_name") == "DonorSpecificMELDraft"
    assert generation.get("llm_used") is True
    assert len(indicators) == 1
    assert indicators[0]["indicator_id"] == "IND_777"


def test_mel_llm_mode_preserves_strategy_specific_indicator_fields(monkeypatch):
    monkeypatch.setattr(mel_module, "openai_compatible_llm_available", lambda: True)

    class DonorSpecificIndicator(BaseModel):
        indicator_id: str
        name: str
        justification: str
        citation: str
        baseline: str = "0"
        target: str = "1"
        indicator_code: str | None = None
        frequency: str | None = None
        evidence_excerpt: str | None = None

    class DonorSpecificMELDraft(BaseModel):
        mel_indicators: list[DonorSpecificIndicator] = Field(default_factory=list)

    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        return {
            "documents": [["Donor indicator guidance excerpt"]],
            "metadatas": [[{"doc_id": "usaid_ads201_p10_c1", "chunk_id": "usaid_ads201_p10_c1", "page": 10}]],
            "ids": [["usaid_ads201_p10_c1"]],
            "distances": [[0.1]],
        }

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    strategy = DonorFactory.get_strategy("usaid")
    monkeypatch.setattr(strategy, "get_mel_schema", lambda: DonorSpecificMELDraft)

    def fake_llm_structured_mel(**kwargs):  # noqa: ARG001
        return (
            {
                "mel_indicators": [
                    {
                        "indicator_id": "IND_778",
                        "name": "Civil servants trained on AI governance",
                        "justification": "Tracks capability development outcome.",
                        "citation": "usaid_ads201_p10_c1",
                        "baseline": "0",
                        "target": "300",
                        "indicator_code": "EG.3.2-27",
                        "frequency": "quarterly",
                        "evidence_excerpt": "Donor indicator guidance excerpt",
                    }
                ]
            },
            "llm:stub-mel-model",
            None,
        )

    monkeypatch.setattr(mel_module, "_llm_structured_mel", fake_llm_structured_mel)

    state = _base_state(llm_mode=True)
    state["strategy"] = strategy
    state["donor_strategy"] = strategy
    out = mel_module.mel_assign_indicators(state)
    indicators = out["logframe_draft"]["indicators"]

    assert len(indicators) == 1
    assert indicators[0]["indicator_id"] == "IND_778"
    assert indicators[0]["indicator_code"] == "EG.3.2-27"
    assert indicators[0]["frequency"] == "quarterly"


def test_mel_llm_tries_fallback_model_chain_and_selects_first_success(monkeypatch):
    monkeypatch.setattr(mel_module, "openai_compatible_llm_available", lambda: True)
    monkeypatch.setattr(mel_module.config.llm, "mel_model", "model-primary")
    monkeypatch.setattr(mel_module.config.llm, "reasoning_model", "model-secondary")
    monkeypatch.setattr(mel_module.config.llm, "cheap_model", "model-cheap")

    def fake_query(*, namespace, query_texts, n_results):  # noqa: ARG001
        return {
            "documents": [["Official donor indicator guidance excerpt"]],
            "metadatas": [[{"doc_id": "usaid_ads201_p10_c1", "chunk_id": "usaid_ads201_p10_c1", "page": 10}]],
            "ids": [["usaid_ads201_p10_c1"]],
            "distances": [[0.1]],
        }

    monkeypatch.setattr(mel_module.vector_store, "query", fake_query)
    calls: list[str] = []

    def fake_llm_structured_mel(**kwargs):
        model_name = str(kwargs.get("model_name") or "")
        calls.append(model_name)
        if model_name == "model-primary":
            return None, None, "primary model unavailable"
        if model_name == "model-secondary":
            return (
                {
                    "indicators": [
                        {
                            "indicator_id": "IND_201",
                            "name": "Officials trained on AI governance",
                            "justification": "Tracks workforce capability improvements.",
                            "citation": "usaid_ads201_p10_c1",
                            "baseline": "0",
                            "target": "300",
                            "evidence_excerpt": "Official donor indicator guidance excerpt",
                        }
                    ]
                },
                "llm:model-secondary",
                None,
            )
        return None, None, "unexpected model"

    monkeypatch.setattr(mel_module, "_llm_structured_mel", fake_llm_structured_mel)

    state = _base_state(llm_mode=True)
    out = mel_module.mel_assign_indicators(state)
    generation = out.get("mel_generation_meta") or {}

    assert generation.get("llm_used") is True
    assert generation.get("engine") == "llm:model-secondary"
    assert generation.get("llm_selected_model") == "model-secondary"
    assert generation.get("llm_attempt_count") == 2
    assert generation.get("llm_models_tried") == ["model-primary", "model-secondary"]
    assert calls == ["model-primary", "model-secondary"]
