package com.easyagents.agent.runtime.prompt;

import com.easyagents.agent.runtime.AgentDefinition;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import org.junit.Assert;
import org.junit.Test;

import java.util.List;

/**
 * 测试运行时系统提示词组合器。
 */
public class SystemPromptComposerTest {

    @Test
    public void shouldKeepUserPromptWhenNoRuntimeProtocolIsRequired() {
        AgentDefinition definition = definition("  user prompt  ");

        Assert.assertEquals("  user prompt  ", SystemPromptComposer.compose(definition));
    }

    @Test
    public void shouldReturnKnowledgeProtocolWhenUserPromptIsBlank() {
        AgentDefinition definition = definition("   ");
        definition.setKnowledgeSpecs(List.of(knowledge("faq")));

        String prompt = SystemPromptComposer.compose(definition);

        Assert.assertTrue(prompt.startsWith("Knowledge tool protocol:"));
        Assert.assertFalse(prompt.startsWith("\n"));
    }

    @Test
    public void shouldAppendKnowledgeProtocolAfterUserPrompt() {
        AgentDefinition definition = definition("user prompt");
        definition.setKnowledgeSpecs(List.of(knowledge("faq")));

        String prompt = SystemPromptComposer.compose(definition);

        Assert.assertTrue(prompt.startsWith("user prompt\n\nKnowledge tool protocol:"));
        Assert.assertTrue(prompt.contains("prefer making one retrieval call"));
        Assert.assertTrue(prompt.contains("reformulate the query and retry once"));
    }

    @Test
    public void shouldAppendKnowledgeProtocolOnlyOnceForMultipleKnowledgeTools() {
        AgentDefinition definition = definition("user prompt");
        definition.setKnowledgeSpecs(List.of(knowledge("faq"), knowledge("policy")));

        String prompt = SystemPromptComposer.compose(definition);

        Assert.assertEquals(1, occurrences(prompt, "Knowledge tool protocol:"));
    }

    @Test
    public void shouldSkipKnowledgeProtocolWhenToolCallingIsDisabled() {
        AgentDefinition definition = definition("user prompt");
        definition.setKnowledgeSpecs(List.of(knowledge("faq")));
        definition.getExecutionOptions().setToolCallingEnabled(false);

        Assert.assertEquals("user prompt", SystemPromptComposer.compose(definition));
    }

    @Test
    public void shouldPreserveAsyncToolProtocolBehavior() {
        AgentDefinition definition = definition(null);
        definition.setToolSpecs(List.of(asyncTool()));

        String prompt = SystemPromptComposer.compose(definition);

        Assert.assertTrue(prompt.startsWith("Async tool protocol:"));
        Assert.assertTrue(prompt.contains("These are internal execution phases."));
        Assert.assertTrue(prompt.contains("call its submit sub-tool first with the user-provided arguments by default"));
        Assert.assertTrue(prompt.contains("immediately call observe"));
    }

    @Test
    public void shouldComposeUserKnowledgeAndAsyncProtocolsInStableOrder() {
        AgentDefinition definition = definition("user prompt");
        definition.setKnowledgeSpecs(List.of(knowledge("faq")));
        definition.setToolSpecs(List.of(asyncTool()));

        String prompt = SystemPromptComposer.compose(definition);
        int userIndex = prompt.indexOf("user prompt");
        int knowledgeIndex = prompt.indexOf("Knowledge tool protocol:");
        int asyncIndex = prompt.indexOf("Async tool protocol:");

        Assert.assertTrue(userIndex >= 0);
        Assert.assertTrue(knowledgeIndex > userIndex);
        Assert.assertTrue(asyncIndex > knowledgeIndex);
        Assert.assertEquals(1, occurrences(prompt, "Knowledge tool protocol:"));
        Assert.assertEquals(1, occurrences(prompt, "Async tool protocol:"));
    }

    /**
     * 创建测试智能体定义。
     *
     * @param systemPrompt 用户系统提示词
     * @return 智能体定义
     */
    private AgentDefinition definition(String systemPrompt) {
        AgentDefinition definition = new AgentDefinition();
        definition.setSystemPrompt(systemPrompt);
        return definition;
    }

    /**
     * 创建测试知识库定义。
     *
     * @param runtimeName 知识库运行名
     * @return 知识库定义
     */
    private AgentKnowledgeSpec knowledge(String runtimeName) {
        AgentKnowledgeSpec spec = new AgentKnowledgeSpec();
        spec.setKnowledgeId("knowledge-" + runtimeName);
        spec.setRuntimeName(runtimeName);
        return spec;
    }

    /**
     * 创建测试异步工具定义。
     *
     * @return 异步工具定义
     */
    private AgentToolSpec asyncTool() {
        AgentToolSpec toolSpec = new AgentToolSpec();
        toolSpec.setName("demo_submit");
        toolSpec.getMetadata().put("asyncTool", true);
        return toolSpec;
    }

    /**
     * 统计文本片段出现次数。
     *
     * @param value 待检查文本
     * @param fragment 目标片段
     * @return 出现次数
     */
    private int occurrences(String value, String fragment) {
        int count = 0;
        int offset = 0;
        while ((offset = value.indexOf(fragment, offset)) >= 0) {
            count++;
            offset += fragment.length();
        }
        return count;
    }
}
