from __future__ import annotations

import asyncio
import json
from types import SimpleNamespace

import pytest
from langchain.agents.middleware.types import ModelRequest, ModelResponse
from langchain_core.language_models.fake_chat_models import FakeListChatModel
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from deerflow.agents.middlewares.clarification_middleware import (
    ClarificationMiddleware,
)
from deerflow.agents.middlewares.knowledge_rag_middleware import (
    KnowledgeRagMiddleware,
)
from deerflow.agents.middlewares.markdown_file_deliverable_middleware import (
    MarkdownFileDeliverableMiddleware,
)
from deerflow.agents.middlewares.rag_citation_middleware import (
    RagCitationMiddleware,
)
from deerflow.agents.middlewares.smart_reasoning_retrieval_middleware import (
    SMART_REASONING_DUPLICATE_SYNTHESIS_FALLBACK,
    SMART_REASONING_FINAL_SYNTHESIS_PROMPT,
    SMART_REASONING_LIMIT_FALLBACK,
    SMART_REASONING_TOOL_NAMES,
    SMART_RETRIEVAL_HISTORY_MARKER,
    SmartReasoningRetrievalMiddleware,
    sanitize_smart_retrieval_history,
)


def _request(messages, *, tools=None, runtime=None) -> ModelRequest:
    return ModelRequest(
        model=FakeListChatModel(responses=["unused"]),
        messages=list(messages),
        tools=list(tools or []),
        runtime=runtime,
    )


def _tool_names(request: ModelRequest) -> list[str]:
    names = []
    for tool in request.tools:
        if isinstance(tool, dict):
            names.append(
                str(
                    tool.get("name")
                    or (tool.get("function") or {}).get("name")
                    or ""
                )
            )
        else:
            names.append(str(getattr(tool, "name", "") or ""))
    return names


def _assert_tool_protocol_is_valid(messages) -> None:
    pending: set[str] = set()
    for message in messages:
        if isinstance(message, AIMessage):
            assert not pending
            calls = list(message.tool_calls)
            if not calls:
                calls = list((message.additional_kwargs or {}).get("tool_calls") or [])
            pending = {
                str(call.get("id") or "")
                for call in calls
                if call.get("id")
            }
        elif isinstance(message, ToolMessage):
            assert message.tool_call_id in pending
            pending.remove(message.tool_call_id)
        else:
            assert not pending
    assert not pending


def _raw_baseline(
    user: HumanMessage,
    *,
    query: str | None = None,
    results: list[dict] | None = None,
):
    raw = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "raw-baseline",
                "name": "ordinary_knowledge_search",
                "args": {"queries": [query or str(user.content)]},
            }
        ],
    )
    raw_result = ToolMessage(
        content=json.dumps(
            {
                "ok": True,
                "operation": "knowledge_search",
                "results": results or [],
            },
            ensure_ascii=False,
        ),
        tool_call_id="raw-baseline",
        name="ordinary_knowledge_search",
    )
    return [user, raw, raw_result]


def test_new_question_drops_old_library_evidence_and_query_but_keeps_current_turn():
    old_ai = AIMessage(
        content="我先核对知识库。",
        tool_calls=[
            {
                "id": "old-smart",
                "name": "ordinary_knowledge_search",
                "args": {"queries": ["OLD-QUERY-SECRET"]},
            },
            {
                "id": "old-other",
                "name": "local_calculator",
                "args": {"expression": "1+1"},
            },
        ],
        additional_kwargs={
            "tool_calls": [
                {
                    "id": "old-smart",
                    "type": "function",
                    "function": {
                        "name": "ordinary_knowledge_search",
                        "arguments": '{"queries":["OLD-QUERY-SECRET"]}',
                    },
                },
                {
                    "id": "old-other",
                    "type": "function",
                    "function": {
                        "name": "local_calculator",
                        "arguments": '{"expression":"1+1"}',
                    },
                },
            ]
        },
        response_metadata={"finish_reason": "tool_calls"},
    )
    old_smart_result = ToolMessage(
        content="OLD-KB-A-EVIDENCE-SECRET",
        tool_call_id="old-smart",
        name="ordinary_knowledge_search",
        additional_kwargs={"mapping_id": "old-library-a"},
    )
    old_other_result = ToolMessage(
        content="calculator-result-2",
        tool_call_id="old-other",
        name="local_calculator",
    )
    current_ai = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "current-smart",
                "name": "ordinary_wiki_search",
                "args": {"patterns": ["CURRENT-LIBRARY-B"]},
            }
        ],
    )
    current_result = ToolMessage(
        content="CURRENT-KB-B-EVIDENCE",
        tool_call_id="current-smart",
        name="ordinary_wiki_search",
    )
    original = [
        HumanMessage(content="旧问题"),
        old_ai,
        old_smart_result,
        old_other_result,
        AIMessage(content="旧回答"),
        HumanMessage(content="切换到知识库 B 后的新问题"),
        current_ai,
        current_result,
    ]

    captured = {}

    def handler(request):
        captured["messages"] = request.messages
        return AIMessage(content="新回答")

    result = SmartReasoningRetrievalMiddleware().wrap_model_call(
        _request(original), handler
    )

    assert isinstance(result, AIMessage)
    sent = captured["messages"]
    rendered = "\n".join(str(message.model_dump()) for message in sent)
    assert "OLD-KB-A-EVIDENCE-SECRET" not in rendered
    assert "OLD-QUERY-SECRET" not in rendered
    assert "old-library-a" not in rendered
    assert SMART_RETRIEVAL_HISTORY_MARKER in rendered
    assert "CURRENT-KB-B-EVIDENCE" in rendered
    assert "CURRENT-LIBRARY-B" in rendered

    sanitized_old_ai = sent[1]
    assert isinstance(sanitized_old_ai, AIMessage)
    assert [call["name"] for call in sanitized_old_ai.tool_calls] == [
        "local_calculator"
    ]
    assert sent[2] is old_other_result
    assert sent[-2:] == [current_ai, current_result]
    _assert_tool_protocol_is_valid(sent)

    # Request rewriting is transient: checkpoint/UI objects retain the full
    # old trace and current evidence.
    assert old_ai.tool_calls[0]["args"]["queries"] == ["OLD-QUERY-SECRET"]
    assert old_smart_result.content == "OLD-KB-A-EVIDENCE-SECRET"


def test_hidden_context_after_question_does_not_make_current_evidence_historical():
    messages = [
        HumanMessage(content="当前问题"),
        AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "current-read",
                    "name": "ordinary_document_chunks",
                    "args": {"chunk_id": "safe-current"},
                }
            ],
        ),
        ToolMessage(
            content="当前轮原文证据",
            tool_call_id="current-read",
            name="ordinary_document_chunks",
        ),
        HumanMessage(
            content="隐藏业务上下文",
            name="business_context",
            additional_kwargs={"hide_from_ui": True},
        ),
    ]

    assert sanitize_smart_retrieval_history(messages) is None


def test_skill_capability_question_is_answered_without_knowledge_retrieval():
    middleware = SmartReasoningRetrievalMiddleware()
    request = _request(
        [HumanMessage(content="你有哪些技能？")],
        tools=[
            {"name": "ordinary_knowledge_search"},
            {"name": "ordinary_wiki_search"},
            {"name": "skill_list"},
        ],
    )
    captured: list[ModelRequest] = []

    def handler(current_request):
        captured.append(current_request)
        return AIMessage(content="我可以列出并按需使用当前可用技能。")

    result = middleware.wrap_model_call(request, handler)

    assert len(captured) == 1
    assert captured[0].messages == request.messages
    assert result.content == "我可以列出并按需使用当前可用技能。"
    assert result.tool_calls == []


def test_non_knowledge_skill_tool_call_is_preserved_verbatim():
    middleware = SmartReasoningRetrievalMiddleware()
    response = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "skill-list-call",
                "name": "skill_list",
                "args": {"category": "available"},
            }
        ],
        response_metadata={"finish_reason": "tool_calls"},
    )

    result = middleware.wrap_model_call(
        _request(
            [HumanMessage(content="列出可用技能")],
            tools=[
                {"name": "ordinary_knowledge_search"},
                {"name": "skill_list"},
            ],
        ),
        lambda _request: response,
    )

    assert result is response
    assert result.tool_calls == [
        {
            "id": "skill-list-call",
            "name": "skill_list",
            "args": {"category": "available"},
            "type": "tool_call",
        }
    ]


def test_model_selected_knowledge_search_is_preserved_without_extra_calls():
    middleware = SmartReasoningRetrievalMiddleware()
    response = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "kb-search-call",
                "name": "ordinary_knowledge_search",
                "args": {"queries": ["全民普法院 机构"]},
            }
        ],
        response_metadata={"finish_reason": "tool_calls"},
    )

    result = middleware.wrap_model_call(
        _request(
            [HumanMessage(content="库里的全民普法院是什么机构？")],
            tools=[
                {"name": "ordinary_knowledge_search"},
                {"name": "ordinary_wiki_search"},
            ],
        ),
        lambda _request: response,
    )

    assert result is response
    assert result.tool_calls == [
        {
            "id": "kb-search-call",
            "name": "ordinary_knowledge_search",
            "args": {"queries": ["全民普法院 机构"]},
            "type": "tool_call",
        }
    ]


def test_direct_answer_after_raw_search_does_not_inject_wiki_search():
    middleware = SmartReasoningRetrievalMiddleware()
    user = HumanMessage(content="根据刚才命中的原文概括一下")
    raw_history = _raw_baseline(
        user,
        results=[
            {
                "title": "普法工作说明",
                "content": "全民普法是面向社会公众开展的法治宣传教育。",
            }
        ],
    )

    result = middleware.wrap_model_call(
        _request(
            raw_history,
            tools=[
                {"name": "ordinary_knowledge_search"},
                {"name": "ordinary_wiki_search"},
                {"name": "ordinary_wiki_read_page"},
            ],
        ),
        lambda _request: AIMessage(
            content="全民普法是面向社会公众开展的法治宣传教育。"
        ),
    )

    assert result.content == "全民普法是面向社会公众开展的法治宣传教育。"
    assert result.tool_calls == []


def test_current_turn_duplicate_queries_and_wiki_pages_are_removed():
    middleware = SmartReasoningRetrievalMiddleware()
    tools = [
        {"name": "ordinary_knowledge_search"},
        {"name": "ordinary_wiki_search"},
        {"name": "ordinary_wiki_read_page"},
    ]
    user = HumanMessage(content="梳理工程进展")
    raw_history = _raw_baseline(user, query="星河工程 进展")
    wiki_search = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-dedupe-search",
                "name": "ordinary_wiki_search",
                "args": {"pattern": "星河工程|进展"},
            }
        ],
    )
    wiki_result = ToolMessage(
        content=json.dumps(
            {
                "ok": True,
                "results": [
                    {"page_ref": "wkpage-one"},
                    {"page_ref": "wkpage-two"},
                ],
            },
            ensure_ascii=False,
        ),
        tool_call_id="wiki-dedupe-search",
        name="ordinary_wiki_search",
    )
    first_read = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-dedupe-read-one",
                "name": "ordinary_wiki_read_page",
                "args": {"page_ref": "wkpage-one"},
            }
        ],
    )
    first_result = ToolMessage(
        content='{"ok":true,"results":[{"complete":true}]}',
        tool_call_id="wiki-dedupe-read-one",
        name="ordinary_wiki_read_page",
    )
    history = [
        *raw_history,
        wiki_search,
        wiki_result,
        first_read,
        first_result,
    ]

    result = middleware.wrap_model_call(
        _request(history, tools=tools),
        lambda _request: AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "raw-dedupe-mixed",
                    "name": "ordinary_knowledge_search",
                    "args": {
                        "queries": [
                            "  星河工程   进展 ",
                            "星河工程 建设单位",
                            "星河工程 建设单位",
                        ]
                    },
                },
                {
                    "id": "wiki-dedupe-repeat",
                    "name": "ordinary_wiki_read_page",
                    "args": {"page_ref": "wkpage-one"},
                },
                {
                    "id": "wiki-dedupe-new",
                    "name": "ordinary_wiki_read_page",
                    "args": {"page_ref": "wkpage-two"},
                },
            ],
        ),
    )

    assert [call["name"] for call in result.tool_calls] == [
        "ordinary_knowledge_search",
        "ordinary_wiki_read_page",
    ]
    assert result.tool_calls[0]["args"]["queries"] == ["星河工程 建设单位"]
    assert result.tool_calls[1]["args"] == {"page_ref": "wkpage-two"}
    assert result.additional_kwargs[
        "knowledge_retrieval_duplicate_actions_skipped"
    ] == 3


def test_new_wiki_cursor_is_preserved_while_repeated_cursor_is_removed():
    middleware = SmartReasoningRetrievalMiddleware()
    user = HumanMessage(content="继续读取相关 Wiki 页面")
    prior_read = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-cursor-prior",
                "name": "ordinary_wiki_read_page",
                "args": {"page_ref": "wkpage-one", "cursor": "cursor-1"},
            }
        ],
    )
    prior_result = ToolMessage(
        content='{"ok":true,"results":[{"complete":false,"next_cursor":"cursor-2"}]}',
        tool_call_id="wiki-cursor-prior",
        name="ordinary_wiki_read_page",
    )

    result = middleware.wrap_model_call(
        _request(
            [user, prior_read, prior_result],
            tools=[{"name": "ordinary_wiki_read_page"}],
        ),
        lambda _request: AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "wiki-cursor-repeat",
                    "name": "ordinary_wiki_read_page",
                    "args": {"page_ref": "wkpage-one", "cursor": "cursor-1"},
                },
                {
                    "id": "wiki-cursor-next",
                    "name": "ordinary_wiki_read_page",
                    "args": {"page_ref": "wkpage-one", "cursor": "cursor-2"},
                },
            ],
        ),
    )

    assert [call["id"] for call in result.tool_calls] == ["wiki-cursor-next"]
    assert result.additional_kwargs[
        "knowledge_retrieval_duplicate_actions_skipped"
    ] == 1


def test_regex_wiki_pattern_and_document_reads_are_deduplicated():
    middleware = SmartReasoningRetrievalMiddleware()
    tools = [
        {"name": "ordinary_knowledge_search"},
        {"name": "ordinary_wiki_search"},
        {"name": "ordinary_wiki_read_page"},
        {"name": "ordinary_chunk_regex_search"},
        {"name": "ordinary_document_chunks"},
    ]
    user = HumanMessage(content="核对工程编号")
    raw_history = _raw_baseline(user)
    actions = [
        AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "wiki-pattern-prior",
                    "name": "ordinary_wiki_search",
                    "args": {"pattern": "工程|编号"},
                }
            ],
        ),
        ToolMessage(
            content='{"ok":true,"results":[]}',
            tool_call_id="wiki-pattern-prior",
            name="ordinary_wiki_search",
        ),
        AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "regex-prior",
                    "name": "ordinary_chunk_regex_search",
                    "args": {"pattern": "编号-[0-9]+"},
                }
            ],
        ),
        ToolMessage(
            content='{"ok":true,"results":[]}',
            tool_call_id="regex-prior",
            name="ordinary_chunk_regex_search",
        ),
        AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "chunk-prior",
                    "name": "ordinary_document_chunks",
                    "args": {
                        "identifier": "chunk-1",
                        "identifier_type": "chunk_id",
                        "page": 1,
                        "page_size": 20,
                    },
                }
            ],
        ),
        ToolMessage(
            content='{"ok":true,"results":[]}',
            tool_call_id="chunk-prior",
            name="ordinary_document_chunks",
        ),
        AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "document-page-prior",
                    "name": "ordinary_document_chunks",
                    "args": {
                        "identifier": "document-1",
                        "identifier_type": "knowledge_id",
                        "page": 1,
                        "page_size": 20,
                    },
                }
            ],
        ),
        ToolMessage(
            content='{"ok":true,"results":[]}',
            tool_call_id="document-page-prior",
            name="ordinary_document_chunks",
        ),
    ]

    result = middleware.wrap_model_call(
        _request([*raw_history, *actions], tools=tools),
        lambda _request: AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "wiki-pattern-repeat",
                    "name": "ordinary_wiki_search",
                    "args": {"pattern": "工程|编号"},
                },
                {
                    "id": "regex-repeat",
                    "name": "ordinary_chunk_regex_search",
                    "args": {"pattern": "编号-[0-9]+"},
                },
                {
                    "id": "chunk-repeat-with-inert-pagination",
                    "name": "ordinary_document_chunks",
                    "args": {
                        "identifier": "chunk-1",
                        "identifier_type": "chunk_id",
                        "page": 99,
                        "page_size": 1,
                    },
                },
                {
                    "id": "document-page-repeat",
                    "name": "ordinary_document_chunks",
                    "args": {
                        "identifier": "document-1",
                        "identifier_type": "knowledge_id",
                        "page": 1,
                        "page_size": 20,
                    },
                },
                {
                    "id": "document-page-next",
                    "name": "ordinary_document_chunks",
                    "args": {
                        "identifier": "document-1",
                        "identifier_type": "knowledge_id",
                        "page": 2,
                        "page_size": 20,
                    },
                },
            ],
        ),
    )

    assert [call["id"] for call in result.tool_calls] == [
        "document-page-next"
    ]
    assert result.additional_kwargs[
        "knowledge_retrieval_duplicate_actions_skipped"
    ] == 4


def test_duplicate_only_call_preserves_model_answer_and_ends_without_reread():
    middleware = SmartReasoningRetrievalMiddleware()
    tools = [
        {"name": "ordinary_knowledge_search"},
        {"name": "ordinary_wiki_search"},
        {"name": "ordinary_wiki_read_page"},
    ]
    user = HumanMessage(content="梳理新闻")
    raw_history = _raw_baseline(user, query="新闻")
    wiki_search = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-only-duplicate-search",
                "name": "ordinary_wiki_search",
                "args": {"pattern": "新闻"},
            }
        ],
    )
    wiki_result = ToolMessage(
        content='{"ok":true,"results":[]}',
        tool_call_id="wiki-only-duplicate-search",
        name="ordinary_wiki_search",
    )

    result = middleware.wrap_model_call(
        _request([*raw_history, wiki_search, wiki_result], tools=tools),
        lambda _request: AIMessage(
            content="依据现有资料回答。",
            tool_calls=[
                {
                    "id": "raw-only-duplicate",
                    "name": "ordinary_knowledge_search",
                    "args": {"queries": ["新闻"]},
                }
            ],
            response_metadata={"finish_reason": "tool_calls"},
        ),
    )

    assert result.content == "依据现有资料回答。"
    assert result.tool_calls == []
    assert result.response_metadata["finish_reason"] == "stop"


@pytest.mark.parametrize("asynchronous", [False, True])
def test_duplicate_only_empty_retries_once_for_tool_free_synthesis(asynchronous):
    middleware = SmartReasoningRetrievalMiddleware()
    tools = [
        {"name": "ordinary_knowledge_search"},
        {"name": "ordinary_wiki_search"},
        {"name": "ordinary_wiki_read_page"},
        {"name": "local_calculator"},
    ]
    user = HumanMessage(content="开幕式主题曲是什么？")
    raw_history = _raw_baseline(
        user,
        results=[
            {
                "title": "开幕式原文",
                "content": "开幕式主题曲《友爱之声》由合唱团演唱。",
            }
        ],
    )
    wiki_search = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-synthesis-search",
                "name": "ordinary_wiki_search",
                "args": {"pattern": "开幕式|主题曲"},
            }
        ],
    )
    wiki_result = ToolMessage(
        content='{"ok":true,"results":[]}',
        tool_call_id="wiki-synthesis-search",
        name="ordinary_wiki_search",
    )
    requests: list[ModelRequest] = []

    def response_for(request: ModelRequest) -> AIMessage:
        requests.append(request)
        if len(requests) == 1:
            return AIMessage(
                content="",
                tool_calls=[
                    {
                        "id": "raw-duplicate-empty",
                        "name": "ordinary_knowledge_search",
                        "args": {"queries": ["开幕式主题曲是什么？"]},
                    }
                ],
                response_metadata={"finish_reason": "tool_calls"},
            )
        return AIMessage(content="开幕式主题曲是《互助之声》。")

    request = _request(
        [*raw_history, wiki_search, wiki_result],
        tools=tools,
    )
    if asynchronous:

        async def handler(current_request):
            return response_for(current_request)

        result = asyncio.run(middleware.awrap_model_call(request, handler))
    else:
        result = middleware.wrap_model_call(request, response_for)

    assert len(requests) == 2
    assert _tool_names(requests[1]) == ["local_calculator"]
    assert requests[1].tool_choice == "none"
    assert requests[1].messages[-1].name == (
        "smart_reasoning_duplicate_synthesis"
    )
    assert requests[1].messages[-1].additional_kwargs["hide_from_ui"] is True
    assert result.content == "开幕式主题曲是《友爱之声》。"
    assert result.tool_calls == []
    assert result.additional_kwargs["knowledge_retrieval_stop_reason"] == (
        "duplicate_only_synthesis"
    )


def test_duplicate_only_empty_uses_safe_fallback_if_synthesis_is_still_empty():
    middleware = SmartReasoningRetrievalMiddleware()
    tools = [
        {"name": "ordinary_knowledge_search"},
        {"name": "ordinary_wiki_search"},
    ]
    user = HumanMessage(content="梳理新闻")
    raw_history = _raw_baseline(user, query="新闻")
    wiki_search = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-empty-synthesis-search",
                "name": "ordinary_wiki_search",
                "args": {"pattern": "新闻"},
            }
        ],
    )
    wiki_result = ToolMessage(
        content='{"ok":true,"results":[]}',
        tool_call_id="wiki-empty-synthesis-search",
        name="ordinary_wiki_search",
    )
    call_count = 0

    def handler(_request):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return AIMessage(
                content="",
                tool_calls=[
                    {
                        "id": "raw-empty-synthesis-repeat",
                        "name": "ordinary_knowledge_search",
                        "args": {"queries": ["新闻"]},
                    }
                ],
            )
        return AIMessage(content="")

    result = middleware.wrap_model_call(
        _request(
            [*raw_history, wiki_search, wiki_result],
            tools=tools,
        ),
        handler,
    )

    assert call_count == 2
    assert result.content == SMART_REASONING_DUPLICATE_SYNTHESIS_FALLBACK


def test_final_answer_uses_raw_document_name_over_wiki_and_model_rewrite():
    middleware = SmartReasoningRetrievalMiddleware()
    tools = [
        {"name": "ordinary_knowledge_search"},
        {"name": "ordinary_wiki_search"},
        {"name": "ordinary_wiki_read_page"},
    ]
    user = HumanMessage(content="开幕式有哪些内容？")
    wiki_search = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-search-name",
                "name": "ordinary_wiki_search",
                "args": {"pattern": "开幕式"},
            }
        ],
    )
    wiki_result = ToolMessage(
        content='{"ok":true,"results":[{"page_ref":"wkpage-opening"}]}',
        tool_call_id="wiki-search-name",
        name="ordinary_wiki_search",
    )
    wiki_read = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "wiki-read-name",
                "name": "ordinary_wiki_read_page",
                "args": {"page_ref": "wkpage-opening"},
            }
        ],
    )
    wiki_page = ToolMessage(
        content=json.dumps(
            {
                "ok": True,
                "results": [
                    {
                        "content_markdown": "Wiki 生成稿称主题曲《互助之声》。"
                    }
                ],
            },
            ensure_ascii=False,
        ),
        tool_call_id="wiki-read-name",
        name="ordinary_wiki_read_page",
    )
    raw_search = AIMessage(
        content="",
        tool_calls=[
            {
                "id": "raw-search-name",
                "name": "ordinary_knowledge_search",
                "args": {"queries": ["开幕式 主题曲"]},
            }
        ],
    )
    raw_result = ToolMessage(
        content=json.dumps(
            {
                "ok": True,
                "results": [
                    {
                        "title": "开幕新闻原文",
                        "content": "文体表演包括《锦绣之韵》和主题曲《友爱之声》。",
                    }
                ],
            },
            ensure_ascii=False,
        ),
        tool_call_id="raw-search-name",
        name="ordinary_knowledge_search",
    )

    result = middleware.wrap_model_call(
        _request(
            [
                user,
                wiki_search,
                wiki_result,
                wiki_read,
                wiki_page,
                raw_search,
                raw_result,
            ],
            tools=tools,
        ),
        lambda _request: AIMessage(
            content="开幕式文体表演包括《锦绣之韵》和主题曲《互助之声》。"
        ),
    )

    assert isinstance(result, AIMessage)
    assert result.content == "开幕式文体表演包括《锦绣之韵》和主题曲《友爱之声》。"
    assert result.additional_kwargs["knowledge_retrieval_fidelity_corrections"] == 1


def test_raw_name_guard_leaves_supported_and_ambiguous_names_unchanged():
    from deerflow.agents.middlewares import smart_reasoning_retrieval_middleware as smart

    supported, supported_count = smart._repair_raw_name_fidelity(
        "主题曲《友爱之声》",
        ["原文写明主题曲《友爱之声》。"],
    )
    ambiguous, ambiguous_count = smart._repair_raw_name_fidelity(
        "主题曲《第三首歌》",
        ["甲文写主题曲《第一首歌》，乙文写主题曲《第二首歌》。"],
    )
    swapped, swapped_count = smart._repair_raw_name_fidelity(
        "主题曲《锦绣之韵》",
        ["文体表演包括《锦绣之韵》和主题曲《友爱之声》。"],
    )
    unrelated, unrelated_count = smart._repair_raw_name_fidelity(
        "作为一般背景，可参照文件名《十四五规划》。",
        ["库内材料只提到文件名《内部规定》。"],
    )

    assert supported == "主题曲《友爱之声》"
    assert supported_count == 0
    assert ambiguous == "主题曲《第三首歌》"
    assert ambiguous_count == 0
    assert swapped == "主题曲《友爱之声》"
    assert swapped_count == 1
    assert unrelated == "作为一般背景，可参照文件名《十四五规划》。"
    assert unrelated_count == 0


@pytest.mark.parametrize("tool_name", sorted(SMART_REASONING_TOOL_NAMES))
def test_every_smart_retrieval_tool_is_removed_from_an_old_turn(tool_name):
    messages = [
        HumanMessage(content="旧问题"),
        AIMessage(
            content="",
            tool_calls=[
                {
                    "id": "old-call",
                    "name": tool_name,
                    "args": {"secret": "old-argument"},
                }
            ],
        ),
        ToolMessage(
            content="old-result-body",
            tool_call_id="old-call",
            name=tool_name,
        ),
        HumanMessage(content="新问题"),
    ]

    sanitized = sanitize_smart_retrieval_history(messages)

    assert sanitized is not None
    rendered = "\n".join(str(message.model_dump()) for message in sanitized)
    assert "old-argument" not in rendered
    assert "old-result-body" not in rendered
    assert SMART_RETRIEVAL_HISTORY_MARKER in rendered
    _assert_tool_protocol_is_valid(sanitized)


def test_raw_provider_non_retrieval_call_in_mixed_bundle_stays_paired():
    messages = [
        HumanMessage(content="旧问题"),
        AIMessage(
            content="",
            additional_kwargs={
                "tool_calls": [
                    {
                        "id": "raw-smart",
                        "type": "function",
                        "function": {
                            "name": "ordinary_wiki_search",
                            "arguments": '{"patterns":["secret"]}',
                        },
                    },
                    {
                        "id": "raw-other",
                        "type": "function",
                        "function": {
                            "name": "local_calculator",
                            "arguments": '{"expression":"2+2"}',
                        },
                    },
                ]
            },
            response_metadata={"finish_reason": "tool_calls"},
        ),
        ToolMessage(
            content="secret-result",
            tool_call_id="raw-smart",
            name="ordinary_wiki_search",
        ),
        ToolMessage(
            content="4",
            tool_call_id="raw-other",
            name="local_calculator",
        ),
        HumanMessage(content="新问题"),
    ]

    sanitized = sanitize_smart_retrieval_history(messages)

    assert sanitized is not None
    old_ai = sanitized[1]
    assert isinstance(old_ai, AIMessage)
    raw_calls = old_ai.additional_kwargs["tool_calls"]
    assert [call["id"] for call in raw_calls] == ["raw-other"]
    assert old_ai.response_metadata["finish_reason"] == "tool_calls"
    assert any(
        isinstance(message, ToolMessage) and message.tool_call_id == "raw-other"
        for message in sanitized
    )
    assert "secret-result" not in "\n".join(
        str(message.model_dump()) for message in sanitized
    )
    _assert_tool_protocol_is_valid(sanitized)


@pytest.mark.parametrize("asynchronous", [False, True])
def test_fortieth_call_hides_retrieval_tools_and_forces_plain_final(asynchronous):
    messages = [HumanMessage(content="需要多轮检索的问题")]
    messages.extend(AIMessage(content=f"intermediate-{index}") for index in range(39))
    middleware = SmartReasoningRetrievalMiddleware()
    request = _request(
        messages,
        tools=[
            {"name": "ordinary_wiki_search"},
            {
                "type": "function",
                "function": {"name": "ordinary_document_chunks"},
            },
            {"name": "local_calculator"},
        ],
    )
    captured = {}

    def response_for(patched):
        captured["request"] = patched
        return ModelResponse(
            result=[
                AIMessage(
                    content="",
                    tool_calls=[
                        {
                            "id": "ignored-call",
                            "name": "ordinary_wiki_search",
                            "args": {"patterns": ["repeat"]},
                        }
                    ],
                    additional_kwargs={
                        "tool_calls": [
                            {
                                "id": "ignored-call",
                                "type": "function",
                                "function": {
                                    "name": "ordinary_wiki_search",
                                    "arguments": '{"patterns":["repeat"]}',
                                },
                            }
                        ]
                    },
                    response_metadata={"finish_reason": "tool_calls"},
                )
            ]
        )

    if asynchronous:

        async def async_handler(patched):
            return response_for(patched)

        result = asyncio.run(middleware.awrap_model_call(request, async_handler))
    else:
        result = middleware.wrap_model_call(request, response_for)

    patched = captured["request"]
    assert _tool_names(patched) == ["local_calculator"]
    assert patched.tool_choice == "none"
    assert SMART_REASONING_FINAL_SYNTHESIS_PROMPT in str(
        patched.messages[-1].content
    )
    assert patched.messages[-1].additional_kwargs["hide_from_ui"] is True

    assert isinstance(result, ModelResponse)
    final = result.result[0]
    assert isinstance(final, AIMessage)
    assert final.tool_calls == []
    assert final.content == SMART_REASONING_LIMIT_FALLBACK
    assert final.additional_kwargs["knowledge_retrieval_stop_reason"] == (
        "iteration_limit"
    )
    assert final.response_metadata["finish_reason"] == "stop"


def test_no_forty_first_model_call_is_allowed():
    messages = [HumanMessage(content="当前问题")]
    messages.extend(AIMessage(content=f"call-{index}") for index in range(40))
    called = False

    def handler(_request):
        nonlocal called
        called = True
        return AIMessage(content="should not run")

    result = SmartReasoningRetrievalMiddleware().wrap_model_call(
        _request(messages, tools=[{"name": "ordinary_wiki_search"}]), handler
    )

    assert called is False
    assert isinstance(result, AIMessage)
    assert result.content == SMART_REASONING_LIMIT_FALLBACK
    assert result.additional_kwargs["knowledge_retrieval_stop_reason"] == (
        "iteration_limit"
    )


def test_middleware_mounts_only_for_smart_mode_and_keeps_clarification_last(
    monkeypatch,
):
    import deerflow.agents.lead_agent.agent as factory

    app_config = SimpleNamespace(
        token_usage=SimpleNamespace(enabled=False),
        skill_evolution=SimpleNamespace(
            enabled=False, review_worker=SimpleNamespace(enabled=False)
        ),
        knowledge=SimpleNamespace(enabled=True, rag_limit=5, rag_enabled=True),
        tool_search=SimpleNamespace(enabled=False),
        skill_stop=SimpleNamespace(enabled=False, skill_names=[]),
        skills=SimpleNamespace(container_path="/mnt/skills"),
        get_model_config=lambda _name: None,
    )
    monkeypatch.setattr(factory, "build_lead_runtime_middlewares", lambda **_kwargs: [])
    monkeypatch.setattr(
        factory, "_create_context_reduction_middleware", lambda **_kwargs: None
    )
    monkeypatch.setattr(
        factory,
        "get_memory_runtime_settings",
        lambda: SimpleNamespace(enabled=False),
    )
    monkeypatch.setattr("deerflow.models.vision.vision_enabled", lambda *_args, **_kwargs: False)

    classic = factory._build_middlewares(
        {
            "context": {
                "knowledge_retrieval_mode": "classic",
                "rag_mode_enabled": True,
            }
        },
        model_name=None,
        app_config=app_config,
    )
    smart = factory._build_middlewares(
        {
            "context": {
                "knowledge_retrieval_mode": "smart_reasoning",
                "rag_mode_enabled": True,
            }
        },
        model_name=None,
        app_config=app_config,
    )
    text_only_coordinator = factory._build_middlewares(
        {
            "context": {
                "allowed_tools": [
                    "agent_orchestration",
                    "ask_clarification",
                    "present_files",
                ],
                "excluded_tools": ["write_file"],
            }
        },
        model_name=None,
        app_config=app_config,
    )

    assert not any(
        isinstance(middleware, SmartReasoningRetrievalMiddleware)
        for middleware in classic
    )
    assert sum(
        isinstance(middleware, SmartReasoningRetrievalMiddleware)
        for middleware in smart
    ) == 1
    assert sum(
        isinstance(middleware, RagCitationMiddleware) for middleware in classic
    ) == 1
    assert not any(
        isinstance(middleware, RagCitationMiddleware) for middleware in smart
    )
    assert sum(
        isinstance(middleware, KnowledgeRagMiddleware) for middleware in classic
    ) == 1
    assert not any(
        isinstance(middleware, KnowledgeRagMiddleware) for middleware in smart
    )
    assert any(
        isinstance(middleware, MarkdownFileDeliverableMiddleware)
        for middleware in classic
    )
    assert not any(
        isinstance(middleware, MarkdownFileDeliverableMiddleware)
        for middleware in text_only_coordinator
    )
    assert isinstance(classic[-1], ClarificationMiddleware)
    assert isinstance(smart[-1], ClarificationMiddleware)
    assert isinstance(text_only_coordinator[-1], ClarificationMiddleware)


def test_smart_agent_keeps_normal_tools_skills_and_mcp_with_selected_space(
    monkeypatch,
) -> None:
    import deerflow.agents.analysis_project as analysis_project
    import deerflow.agents.lead_agent.agent as factory
    import deerflow.tools as tool_factory

    captured: dict[str, object] = {}

    def fake_get_available_tools(**kwargs):
        captured["tool_kwargs"] = kwargs
        return [
            {"name": "web_search"},
            {"name": "skill_list"},
            {"name": "ordinary_knowledge_search"},
        ]

    def fake_apply_prompt_template(**kwargs):
        captured["prompt_kwargs"] = kwargs
        return "<citations>keep web citations</citations>"

    monkeypatch.setattr(tool_factory, "get_available_tools", fake_get_available_tools)
    monkeypatch.setattr(factory, "validate_agent_id", lambda value: value)
    monkeypatch.setattr(
        factory,
        "load_agent_config",
        lambda _agent_id: SimpleNamespace(
            name="普通智能体",
            skills=["normal-skill"],
            model=None,
            tool_groups=None,
        ),
    )
    monkeypatch.setattr(factory, "load_agent_soul", lambda _agent_id: "")
    monkeypatch.setattr(
        factory,
        "_resolve_model_name",
        lambda _name, *, app_config: "test-model",
    )
    monkeypatch.setattr(factory, "create_chat_model", lambda **_kwargs: "model")
    monkeypatch.setattr(factory, "apply_prompt_template", fake_apply_prompt_template)
    monkeypatch.setattr(factory, "_build_middlewares", lambda *_args, **_kwargs: [])
    monkeypatch.setattr(factory, "_log_es_query_routing_registration", lambda **_kwargs: None)
    monkeypatch.setattr(factory, "create_agent", lambda **kwargs: kwargs)
    monkeypatch.setattr(analysis_project, "project_system_prompt", lambda _cfg: "")

    app_config = SimpleNamespace(
        get_model_config=lambda _name: SimpleNamespace(supports_thinking=True),
    )
    config = {
        "context": {
            "agent_id": "ordinary-agent",
            "knowledge_space_strict": True,
            "knowledge_retrieval_mode": "smart_reasoning",
            "subagent_enabled": True,
        }
    }

    built = factory._make_lead_agent(config, app_config=app_config)

    tool_kwargs = captured["tool_kwargs"]
    assert isinstance(tool_kwargs, dict)
    assert tool_kwargs["include_mcp"] is True
    excluded = set(tool_kwargs["excluded_tools"])
    assert "llmwiki_search" in excluded
    assert "web_search" not in excluded
    assert "skill_list" not in excluded
    assert "task" not in excluded
    assert not (excluded & SMART_REASONING_TOOL_NAMES)
    assert built["tools"] == [
        {"name": "web_search"},
        {"name": "skill_list"},
        {"name": "ordinary_knowledge_search"},
    ]
    assert captured["prompt_kwargs"]["available_skills"] == {"normal-skill"}
    assert captured["prompt_kwargs"]["subagent_enabled"] is True
    assert "<citations>keep web citations</citations>" in built["system_prompt"]
    assert "不表示每个问题都必须检索" in built["system_prompt"]
    assert config["metadata"]["knowledge_space_strict"] is False


def test_rag_citation_middleware_fails_closed_for_smart_mode() -> None:
    smart_request = _request(
        [HumanMessage(content="查制度")],
        runtime=SimpleNamespace(
            context={
                "knowledge_retrieval_mode": "smart_reasoning",
                "rag_mode_enabled": True,
            },
            config={
                "configurable": {"rag_mode_enabled": True},
            },
        ),
    )
    assert RagCitationMiddleware._enabled(smart_request) is False
    assert RagCitationMiddleware._patch(smart_request) is None

    classic_request = _request(
        [HumanMessage(content="查制度")],
        runtime=SimpleNamespace(
            context={
                "knowledge_retrieval_mode": "classic",
                "rag_mode_enabled": True,
            },
            config={},
        ),
    )
    assert RagCitationMiddleware._enabled(classic_request) is True


def test_smart_prompt_keeps_normal_tools_and_makes_knowledge_optional() -> None:
    import deerflow.agents.lead_agent.agent as factory

    base = "before<citations>must add Sources</citations>middle<rag_citation_mode>must cite [1]</rag_citation_mode>after"
    cleaned = factory._without_classic_citation_contracts(base)
    prompt = factory._SMART_REASONING_KNOWLEDGE_SCOPE_PROMPT
    assert cleaned == "before<citations>must add Sources</citations>middleafter"
    assert "不表示每个问题都必须检索" in prompt
    assert "你有哪些技能/会做什么/有什么工具" in prompt
    assert "直接回答，不得为了展示流程而调用知识库工具" in prompt
    assert "知识库工具与其它正常工具并列" in prompt
    assert "正常问答、技能、Web、MCP、子智能体及其它工具仍按原有权限" in prompt
    assert "不存在强制顺序" in prompt
    assert "不得重复相同 query" in prompt
    assert "专名、作品/文件/活动标题、日期、时间、数字、单位和数量关系" in prompt
    assert "证据未明确给出的细节不要补写" in prompt
    assert "不要求使用 [1]、[2] 等编号" in prompt
    assert "不要求生成文末参考文献列表" in prompt