package tech.easyflow.ai.easyagents.memory;

import com.easyagents.core.message.AiMessage;
import com.easyagents.core.message.Message;
import com.easyagents.core.message.ToolCall;
import com.easyagents.core.message.ToolMessage;
import com.easyagents.core.model.client.OpenAIChatMessageSerializer;
import com.easyagents.llm.deepseek.DeepseekConfig;
import org.junit.Assert;
import org.junit.Test;
import tech.easyflow.core.runtime.ChatAssistantAccumulator;
import tech.easyflow.core.runtime.ChatRuntimeMessage;

import java.util.List;

/**
 * {@link RuntimeChatMemory} 测试。
 */
public class RuntimeChatMemoryTest {

    /**
     * 应当从结构化 payload 中恢复完整的 assistant/tool 消息链。
     */
    @Test
    public void shouldRestoreStructuredAssistantAndToolHistory() {
        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("最终回答");

        ChatRuntimeMessage runtimeMessage = new ChatRuntimeMessage();
        runtimeMessage.setRole("assistant");
        runtimeMessage.setContentText("最终回答");
        runtimeMessage.setContentPayload(accumulator.buildPayload("最终回答"));

        RuntimeChatMemory memory = new RuntimeChatMemory("c1", List.of(runtimeMessage));
        List<Message> messages = memory.getMessages(10);

        Assert.assertEquals(3, messages.size());
        Assert.assertTrue(messages.get(0) instanceof AiMessage);
        Assert.assertTrue(messages.get(1) instanceof ToolMessage);
        Assert.assertTrue(messages.get(2) instanceof AiMessage);

        AiMessage toolAssistant = (AiMessage) messages.get(0);
        Assert.assertEquals("思考-1", toolAssistant.getReasoningContent());
        Assert.assertEquals(1, toolAssistant.getToolCalls().size());
        Assert.assertEquals("call-1", toolAssistant.getToolCalls().get(0).getId());
        Assert.assertEquals("{\"query\":\"java\"}", toolAssistant.getToolCalls().get(0).getArguments());

        ToolMessage toolMessage = (ToolMessage) messages.get(1);
        Assert.assertEquals("call-1", toolMessage.getToolCallId());
        Assert.assertEquals("{\"hits\":1}", toolMessage.getContent());

        AiMessage finalAssistant = (AiMessage) messages.get(2);
        Assert.assertEquals("最终回答", finalAssistant.getTextContent());
        Assert.assertEquals("思考-2", finalAssistant.getReasoningContent());
    }

    /**
     * 旧结构历史应保持纯文本兼容。
     */
    @Test
    public void shouldFallbackToPlainAssistantMessageForLegacyPayload() {
        ChatRuntimeMessage runtimeMessage = new ChatRuntimeMessage();
        runtimeMessage.setRole("assistant");
        runtimeMessage.setContentText("旧版回答");

        RuntimeChatMemory memory = new RuntimeChatMemory("c2", List.of(runtimeMessage));
        List<Message> messages = memory.getMessages(10);

        Assert.assertEquals(1, messages.size());
        Assert.assertTrue(messages.get(0) instanceof AiMessage);
        Assert.assertEquals("旧版回答", ((AiMessage) messages.get(0)).getTextContent());
    }

    /**
     * 内部聊天结构化历史回放后，DeepSeek 序列化应继续带上 reasoning_content 与 tool 链。
     */
    @Test
    public void shouldSerializeRestoredRuntimeHistoryForDeepseekFollowUp() {
        ChatAssistantAccumulator accumulator = new ChatAssistantAccumulator();
        accumulator.appendReasoning("先思考");
        accumulator.appendToolCall("call-1", "kb_search", "{\"query\":\"java\"}");
        accumulator.appendToolResult("call-1", "kb_search", "{\"hits\":1}");

        ChatRuntimeMessage runtimeMessage = new ChatRuntimeMessage();
        runtimeMessage.setRole("assistant");
        runtimeMessage.setContentPayload(accumulator.buildPayload(null));

        RuntimeChatMemory memory = new RuntimeChatMemory("c3", List.of(runtimeMessage));
        List<Message> restoredMessages = memory.getMessages(10);

        DeepseekConfig config = new DeepseekConfig();
        config.setSupportThinking(Boolean.TRUE);
        config.setThinkingProtocol("deepseek");
        config.setNeedReasoningContentForToolMessage(Boolean.TRUE);

        List<java.util.Map<String, Object>> serialized = new OpenAIChatMessageSerializer()
                .serializeMessages(restoredMessages, config);

        Assert.assertEquals(2, serialized.size());
        Assert.assertEquals("assistant", serialized.get(0).get("role"));
        Assert.assertEquals("先思考", serialized.get(0).get("reasoning_content"));
        Assert.assertTrue(serialized.get(0).containsKey("tool_calls"));
        Assert.assertEquals("tool", serialized.get(1).get("role"));
        Assert.assertEquals("call-1", serialized.get(1).get("tool_call_id"));
    }

    /**
     * 连续多轮 tool-only assistant 回放后仍应保持独立顺序。
     */
    @Test
    public void shouldRestoreSeparatedToolOnlyAssistantSegments() {
        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}");

        ChatRuntimeMessage runtimeMessage = new ChatRuntimeMessage();
        runtimeMessage.setRole("assistant");
        runtimeMessage.setContentPayload(accumulator.buildPayload(null));

        RuntimeChatMemory memory = new RuntimeChatMemory("c4", List.of(runtimeMessage));
        List<Message> messages = memory.getMessages(10);

        Assert.assertEquals(4, messages.size());
        Assert.assertTrue(messages.get(0) instanceof AiMessage);
        Assert.assertTrue(messages.get(1) instanceof ToolMessage);
        Assert.assertTrue(messages.get(2) instanceof AiMessage);
        Assert.assertTrue(messages.get(3) instanceof ToolMessage);

        AiMessage firstAssistant = (AiMessage) messages.get(0);
        AiMessage secondAssistant = (AiMessage) messages.get(2);
        Assert.assertEquals(1, firstAssistant.getToolCalls().size());
        Assert.assertEquals("call-1", firstAssistant.getToolCalls().get(0).getId());
        Assert.assertEquals(1, secondAssistant.getToolCalls().size());
        Assert.assertEquals("call-2", secondAssistant.getToolCalls().get(0).getId());
    }
}
