package tech.easyflow.core.runtime;

import org.junit.Assert;
import org.junit.Test;

import java.util.List;
import java.util.Map;

/**
 * {@link ChatAssistantAccumulator} 测试。
 */
public class ChatAssistantAccumulatorTest {

    /**
     * 应同时保留展示链和可回放的结构化消息链。
     */
    @Test
    @SuppressWarnings("unchecked")
    public void shouldBuildStructuredPayloadWithDisplayChains() {
        ChatAssistantAccumulator accumulator = new ChatAssistantAccumulator();
        accumulator.appendReasoning("思考-1");
        accumulator.appendToolCall("call-1", "kb_search", "{\"query\":\"java\"}");
        accumulator.appendToolResult("call-1", "kb_search", "{\"hits\":1}");
        accumulator.appendReasoning("思考-2");
        accumulator.appendContent("最终回答");

        Map<String, Object> payload = accumulator.buildPayload("最终回答");
        List<Map<String, Object>> messageChain = (List<Map<String, Object>>) payload.get("messageChain");
        List<Map<String, Object>> chains = (List<Map<String, Object>>) payload.get("chains");

        Assert.assertEquals(3, messageChain.size());
        Assert.assertEquals("assistant", messageChain.get(0).get("role"));
        Assert.assertEquals("思考-1", messageChain.get(0).get("reasoningContent"));
        Assert.assertEquals("tool", messageChain.get(1).get("role"));
        Assert.assertEquals("assistant", messageChain.get(2).get("role"));
        Assert.assertEquals("最终回答", messageChain.get(2).get("content"));
        Assert.assertEquals("思考-2", messageChain.get(2).get("reasoningContent"));

        Assert.assertFalse(chains.isEmpty());
        Assert.assertEquals("思考-1思考-2", chains.get(0).get("reasoning_content"));
        Assert.assertEquals("{\"query\":\"java\"}", chains.get(1).get("arguments"));
        Assert.assertEquals("{\"hits\":1}", chains.get(1).get("result"));
    }

    /**
     * 连续多轮 tool-only assistant 不应被错误合并到同一条 assistant 历史。
     */
    @Test
    @SuppressWarnings("unchecked")
    public void shouldKeepToolOnlyAssistantSegmentsSeparatedAcrossRecursiveRounds() {
        ChatAssistantAccumulator accumulator = new ChatAssistantAccumulator();
        accumulator.appendToolCall("call-1", "tool_one", "{\"step\":1}");
        accumulator.appendToolResult("call-1", "tool_one", "{\"ok\":1}");
        accumulator.appendToolCall("call-2", "tool_two", "{\"step\":2}");
        accumulator.appendToolResult("call-2", "tool_two", "{\"ok\":2}");

        Map<String, Object> payload = accumulator.buildPayload(null);
        List<Map<String, Object>> messageChain = (List<Map<String, Object>>) payload.get("messageChain");

        Assert.assertEquals(4, messageChain.size());
        Assert.assertEquals("assistant", messageChain.get(0).get("role"));
        Assert.assertEquals("tool", messageChain.get(1).get("role"));
        Assert.assertEquals("assistant", messageChain.get(2).get("role"));
        Assert.assertEquals("tool", messageChain.get(3).get("role"));

        List<Map<String, Object>> firstToolCalls = (List<Map<String, Object>>) messageChain.get(0).get("toolCalls");
        List<Map<String, Object>> secondToolCalls = (List<Map<String, Object>>) messageChain.get(2).get("toolCalls");
        Assert.assertEquals(1, firstToolCalls.size());
        Assert.assertEquals("call-1", firstToolCalls.get(0).get("id"));
        Assert.assertEquals(1, secondToolCalls.size());
        Assert.assertEquals("call-2", secondToolCalls.get(0).get("id"));
    }

    /**
     * 工具展示名应进入展示链和 assistant toolCalls，但不覆盖真实工具名。
     */
    @Test
    @SuppressWarnings("unchecked")
    public void shouldKeepToolDisplayNameWithoutOverridingToolName() {
        ChatAssistantAccumulator accumulator = new ChatAssistantAccumulator();
        accumulator.appendToolCall("call-1", "mcp_123_search", "知识库 MCP - search", "{\"q\":\"java\"}");
        accumulator.appendToolResult("call-1", "mcp_123_search", "知识库 MCP - search", "{\"ok\":true}");

        Map<String, Object> payload = accumulator.buildPayload(null);
        List<Map<String, Object>> chains = (List<Map<String, Object>>) payload.get("chains");
        List<Map<String, Object>> messageChain = (List<Map<String, Object>>) payload.get("messageChain");
        List<Map<String, Object>> toolCalls = (List<Map<String, Object>>) messageChain.get(0).get("toolCalls");

        Assert.assertEquals("mcp_123_search", chains.get(0).get("name"));
        Assert.assertEquals("知识库 MCP - search", chains.get(0).get("toolDisplayName"));
        Assert.assertEquals("mcp_123_search", toolCalls.get(0).get("name"));
        Assert.assertEquals("知识库 MCP - search", toolCalls.get(0).get("toolDisplayName"));
    }
}
