from __future__ import annotations

from grantflow.core.strategies.factory import DonorFactory
from grantflow.swarm.nodes import architect_generation as architect_generation_module
from grantflow.swarm.nodes.architect_generation import _fallback_structured_toc, generate_toc_under_contract


def test_architect_llm_validation_failure_retries_once_and_recovers(monkeypatch):
    strategy = DonorFactory.get_strategy("usaid")
    monkeypatch.setattr(architect_generation_module, "openai_compatible_llm_available", lambda: True)
    schema_cls = strategy.get_toc_schema()
    valid_payload, _ = _fallback_structured_toc(
        schema_cls,
        donor_id="usaid",
        project="Water Sanitation",
        country="Kenya",
        revision_hint="",
        evidence_hits=[],
    )

    calls = []

    def fake_llm_structured_toc(*args, **kwargs):
        calls.append(kwargs.get("validation_error_hint"))
        if len(calls) == 1:
            return {"invalid": "payload"}, "llm:mock", None
        return valid_payload, "llm:mock", None

    monkeypatch.setattr(architect_generation_module, "_llm_structured_toc", fake_llm_structured_toc)

    state = {
        "donor": "usaid",
        "donor_id": "usaid",
        "strategy": strategy,
        "donor_strategy": strategy,
        "input": {"project": "Water Sanitation", "country": "Kenya"},
        "input_context": {"project": "Water Sanitation", "country": "Kenya"},
        "llm_mode": True,
        "critic_notes": {},
    }

    toc, validation, generation_meta, claim_citations = generate_toc_under_contract(
        state=state, strategy=strategy, evidence_hits=[]
    )
    assert toc
    assert validation["valid"] is True
    assert generation_meta["llm_used"] is True
    assert generation_meta.get("llm_validation_repair_attempted") is True
    assert len(calls) == 2
    assert calls[0] is None
    assert isinstance(calls[1], str) and calls[1]
    assert isinstance(claim_citations, list)


def test_architect_llm_mode_without_api_key_uses_emergency_fallback(monkeypatch):
    strategy = DonorFactory.get_strategy("usaid")
    monkeypatch.setattr(architect_generation_module, "openai_compatible_llm_available", lambda: False)
    monkeypatch.setattr(architect_generation_module, "_llm_structured_toc", lambda *args, **kwargs: (None, None, None))

    state = {
        "donor": "usaid",
        "donor_id": "usaid",
        "strategy": strategy,
        "donor_strategy": strategy,
        "input": {"project": "Water Sanitation", "country": "Kenya"},
        "input_context": {"project": "Water Sanitation", "country": "Kenya"},
        "llm_mode": True,
        "critic_notes": {},
    }

    toc, validation, generation_meta, _claim_citations = generate_toc_under_contract(
        state=state, strategy=strategy, evidence_hits=[]
    )
    assert toc
    assert validation["valid"] is True
    assert generation_meta.get("llm_used") is False
    assert generation_meta.get("llm_requested") is True
    assert generation_meta.get("llm_available") is False
    assert generation_meta.get("llm_attempted") is False
    assert generation_meta.get("fallback_used") is True
    assert generation_meta.get("fallback_class") == "emergency"
    assert generation_meta.get("architect_mode") == "llm"
    assert str(generation_meta.get("engine") or "").startswith("fallback:")
    assert "missing" in str(generation_meta.get("llm_fallback_reason") or "").lower()


def test_architect_llm_prompt_receives_full_input_context_and_schema_contract(monkeypatch):
    strategy = DonorFactory.get_strategy("usaid")
    monkeypatch.setattr(architect_generation_module, "openai_compatible_llm_available", lambda: True)

    schema_cls = strategy.get_toc_schema()
    valid_payload, _ = _fallback_structured_toc(
        schema_cls,
        donor_id="usaid",
        project="Digital Governance",
        country="Kazakhstan",
        revision_hint="",
        evidence_hits=[],
    )
    captured: dict[str, object] = {}

    def fake_llm_structured_toc(*args, **kwargs):
        captured["input_context"] = kwargs.get("input_context")
        captured["schema_contract_hint"] = kwargs.get("schema_contract_hint")
        captured["schema_json_contract_hint"] = kwargs.get("schema_json_contract_hint")
        return valid_payload, "llm:mock", None

    monkeypatch.setattr(architect_generation_module, "_llm_structured_toc", fake_llm_structured_toc)

    state = {
        "donor_id": "usaid",
        "donor_strategy": strategy,
        "input_context": {
            "project": "Digital Governance",
            "country": "Kazakhstan",
            "sector": "Public Administration",
            "budget_usd": 2500000,
        },
        "llm_mode": True,
        "critic_notes": {},
    }
    _toc, validation, generation_meta, _claim_citations = generate_toc_under_contract(
        state=state,
        strategy=strategy,
        evidence_hits=[],
    )
    assert validation["valid"] is True
    assert generation_meta["llm_used"] is True
    assert generation_meta["schema_contract_hint_present"] is True
    assert generation_meta["schema_json_hint_present"] is True
    assert generation_meta["input_context_key_count"] == 4
    assert isinstance(captured.get("input_context"), dict)
    assert captured["input_context"] == state["input_context"]
    schema_hint = str(captured.get("schema_contract_hint") or "")
    assert schema_hint
    assert "project_goal" in schema_hint
    schema_json_hint = str(captured.get("schema_json_contract_hint") or "")
    assert schema_json_hint
    assert "properties" in schema_json_hint


def test_architect_llm_tries_fallback_model_chain_and_selects_first_success(monkeypatch):
    strategy = DonorFactory.get_strategy("usaid")
    monkeypatch.setattr(architect_generation_module, "openai_compatible_llm_available", lambda: True)
    monkeypatch.setattr(architect_generation_module.config.llm, "architect_model", "model-primary")
    monkeypatch.setattr(architect_generation_module.config.llm, "reasoning_model", "model-secondary")
    monkeypatch.setattr(architect_generation_module.config.llm, "cheap_model", "model-cheap")

    schema_cls = strategy.get_toc_schema()
    valid_payload, _ = _fallback_structured_toc(
        schema_cls,
        donor_id="usaid",
        project="Digital Governance",
        country="Kazakhstan",
        revision_hint="",
        evidence_hits=[],
    )
    calls: list[str] = []

    def fake_llm_structured_toc(*args, **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 valid_payload, "llm:model-secondary", None
        return None, None, "unexpected model"

    monkeypatch.setattr(architect_generation_module, "_llm_structured_toc", fake_llm_structured_toc)

    state = {
        "donor_id": "usaid",
        "donor_strategy": strategy,
        "input_context": {"project": "Digital Governance", "country": "Kazakhstan"},
        "llm_mode": True,
        "critic_notes": {},
    }
    _toc, validation, generation_meta, _claim_citations = generate_toc_under_contract(
        state=state,
        strategy=strategy,
        evidence_hits=[],
    )
    assert validation["valid"] is True
    assert generation_meta["llm_used"] is True
    assert generation_meta["engine"] == "llm:model-secondary"
    assert generation_meta["llm_selected_model"] == "model-secondary"
    assert generation_meta["llm_attempt_count"] == 2
    assert generation_meta["llm_models_tried"] == ["model-primary", "model-secondary"]
    assert calls[:2] == ["model-primary", "model-secondary"]
    assert len(calls) >= 2


def test_architect_llm_soft_quality_repair_retries_once(monkeypatch):
    strategy = DonorFactory.get_strategy("usaid")
    monkeypatch.setattr(architect_generation_module, "openai_compatible_llm_available", lambda: True)

    schema_cls = strategy.get_toc_schema()
    valid_payload, _ = _fallback_structured_toc(
        schema_cls,
        donor_id="usaid",
        project="Water Sanitation",
        country="Kenya",
        revision_hint="",
        evidence_hits=[],
    )
    placeholder_payload = dict(valid_payload)
    placeholder_payload["project_goal"] = "TBD"

    hints: list[str | None] = []

    def fake_llm_structured_toc(*args, **kwargs):
        hints.append(kwargs.get("validation_error_hint"))
        if len(hints) == 1:
            return placeholder_payload, "llm:mock", None
        return valid_payload, "llm:mock", None

    monkeypatch.setattr(architect_generation_module, "_llm_structured_toc", fake_llm_structured_toc)

    state = {
        "donor_id": "usaid",
        "donor_strategy": strategy,
        "input_context": {"project": "Water Sanitation", "country": "Kenya"},
        "llm_mode": True,
        "critic_notes": {},
    }
    toc, validation, generation_meta, _claim_citations = generate_toc_under_contract(
        state=state,
        strategy=strategy,
        evidence_hits=[],
    )
    assert validation["valid"] is True
    assert toc["project_goal"] != "TBD"
    assert generation_meta["llm_used"] is True
    assert generation_meta.get("llm_quality_repair_attempted") is True
    assert int(generation_meta.get("llm_quality_issue_count") or 0) >= 1
    assert len(hints) == 2
    assert hints[0] is None
    assert isinstance(hints[1], str)
    assert "Soft quality issues detected" in str(hints[1])
