"""Step-budget recovery must preserve history and allow a real next turn."""

from unittest.mock import AsyncMock

import pytest
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import END, START, MessagesState, StateGraph
from pydantic import ValidationError

from deerflow.config.system_settings import ConversationContextSettings
from deerflow.runtime.runs.manager import RunManager
from deerflow.runtime.runs.schemas import RunStatus
from deerflow.runtime.runs.worker import RunContext, run_agent


@pytest.mark.parametrize("value", [19, 10001, 25.5, True, "100"])
def test_invalid_budget(value):
    from app.gateway.routers.system_settings import ConversationContextUpdateRequest

    with pytest.raises(ValidationError):
        ConversationContextSettings(max_run_steps=value)
    with pytest.raises(ValidationError):
        ConversationContextUpdateRequest(max_run_steps=value)


@pytest.mark.asyncio
@pytest.mark.parametrize("budget", [3, 4])
async def test_limit_is_saved_without_error_and_continue_works(budget):
    calls = []

    def model(state):
        messages = state["messages"]
        if any(isinstance(m, HumanMessage) and m.content == "继续" for m in messages):
            # A provider requires every earlier call to have a result.
            tool_ids = {m.tool_call_id for m in messages if isinstance(m, ToolMessage)}
            assert all(c["id"] in tool_ids for m in messages if isinstance(m, AIMessage) for c in m.tool_calls)
            return {"messages": [AIMessage(content="已继续完成")]}
        return {
            "messages": [
                AIMessage(
                    content="正在处理",
                    tool_calls=[
                        {
                            "id": f"call-{len(messages)}",
                            "name": "search",
                            "args": {},
                        }
                    ],
                )
            ]
        }

    def tool(state):
        call = state["messages"][-1].tool_calls[0]
        calls.append(call["id"])
        return {"messages": [ToolMessage(content="已有检索结果", tool_call_id=call["id"])]}

    graph = StateGraph(MessagesState)
    graph.add_node("model", model)
    graph.add_node("tools", tool)
    graph.add_edge(START, "model")
    graph.add_conditional_edges("model", lambda s: "tools" if s["messages"][-1].tool_calls else END)
    graph.add_edge("tools", "model")
    saver = InMemorySaver()
    agent = graph.compile(checkpointer=saver)
    config = {"configurable": {"thread_id": "step-test"}, "recursion_limit": budget}
    manager = RunManager()
    record = await manager.create("step-test", "lead_agent")
    bridge = AsyncMock()
    await run_agent(bridge, manager, record, ctx=RunContext(checkpointer=saver), agent_factory=lambda config: agent, graph_input={"messages": [HumanMessage(content="研究")]}, config=config, stream_modes=["messages", "values"])
    assert record.status == RunStatus.interrupted
    assert record.error is None
    assert not any(c.args[1] == "error" for c in bridge.publish.call_args_list)
    events = {c.args[1]: c.args[2] for c in bridge.publish.call_args_list}
    assert "步骤上限" in events["messages"][0]["content"]
    assert "步骤上限" in events["values"]["messages"][-1]["content"]
    bridge.publish_end.assert_awaited_once()
    state = await agent.aget_state(config)
    assert not state.next
    assert "已有检索结果" in [m.content for m in state.values["messages"]]
    assert "步骤上限" in state.values["messages"][-1].content
    count = len(calls)
    result = await agent.ainvoke({"messages": [HumanMessage(content="继续")]}, config)
    assert result["messages"][-1].content == "已继续完成"
    assert len(calls) == count


def test_server_budget_overrides_client_and_preserves_other_graphs(monkeypatch):
    from app.gateway.services import build_run_config

    monkeypatch.setattr("deerflow.config.system_settings.get_conversation_context_settings", lambda: ConversationContextSettings(max_run_steps=800))
    for assistant in (None, "lead_agent", "custom-agent"):
        assert build_run_config("t", {"recursion_limit": 99999}, None, assistant_id=assistant)["recursion_limit"] == 800
    assert build_run_config("t", {"recursion_limit": 250}, None, assistant_id="canvas_agent")["recursion_limit"] == 250


@pytest.mark.parametrize("client_limit", [None, 1, 300])
@pytest.mark.parametrize("server_limit", [20, 800])
def test_open_chat_uses_server_budget(monkeypatch, client_limit, server_limit):
    from app.gateway.routers.open_chat import OpenChatRequest, _build_run_context

    monkeypatch.setattr("deerflow.config.system_settings.get_conversation_context_settings", lambda: ConversationContextSettings(max_run_steps=server_limit))
    body = OpenChatRequest(message="hello", **({"recursion_limit": client_limit} if client_limit is not None else {}))
    config, _, _ = _build_run_context(body)
    assert config["recursion_limit"] == server_limit
    assert config["configurable"]["ephemeral"] is True


@pytest.mark.asyncio
async def test_setting_requires_admin_and_roundtrips(monkeypatch):
    from fastapi import HTTPException

    from app.gateway.routers import system_settings as routes
    from deerflow.config.system_settings import SystemSettings

    current = SystemSettings()
    monkeypatch.setattr(routes, "load_system_settings", lambda: current)

    def save(value):
        nonlocal current
        current = value
        return value

    monkeypatch.setattr(routes, "save_system_settings", save)
    guard = AsyncMock(side_effect=HTTPException(status_code=403))
    monkeypatch.setattr(routes, "_require_admin", guard)
    body = routes.ConversationContextUpdateRequest(max_run_steps=600)
    with pytest.raises(HTTPException):
        await routes.update_conversation_context_setting(body, None)
    assert current.conversation_context.max_run_steps == 1000
    guard.side_effect = None
    result = await routes.update_conversation_context_setting(body, None)
    assert result.max_run_steps == 600
    assert (await routes.get_conversation_context_setting()).max_run_steps == 600
    assert current.conversation_context.recent_turn_window_enabled is False