package com.easyagents.flow.core.node;

import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository;
import org.junit.Assert;
import org.junit.Test;

import java.util.Arrays;
import java.util.Collections;
import java.util.Map;
import java.util.UUID;

/**
 * 内容模板节点变量渲染回归测试。
 */
public class TemplateNodeTest {

    /**
     * 验证模板可以使用中文参数名称。
     */
    @Test
    public void shouldRenderChineseParameterNames() {
        TemplateNode node = templateNode(
                "#(申请人)\n\n#(被申请人)",
                "申请人",
                "被申请人");
        Chain chain = chain(Map.of(
                "申请人", "申请内容",
                "被申请人", "答辩内容"));

        Map<String, Object> result = node.execute(chain);

        Assert.assertEquals(
                "申请内容\n\n答辩内容",
                result.get("finalContent"));
    }

    /**
     * 验证开启中文表达式后继续兼容英文参数名称。
     */
    @Test
    public void shouldKeepRenderingEnglishParameterNames() {
        TemplateNode node = templateNode(
                "#(applicant)\n\n#(respondent)",
                "applicant",
                "respondent");
        Chain chain = chain(Map.of(
                "applicant", "申请内容",
                "respondent", "答辩内容"));

        Map<String, Object> result = node.execute(chain);

        Assert.assertEquals(
                "申请内容\n\n答辩内容",
                result.get("finalContent"));
    }

    /**
     * 验证自动模板变量可以通过引用读取上游节点输出。
     */
    @Test
    public void shouldRenderManagedUpstreamReference() {
        String parameterName = "ref_node__llm_2e_output";
        TemplateNode node = templateNode(
                "模型输出：#(" + parameterName + ")");
        Parameter parameter = new Parameter(parameterName);
        parameter.setRefType(RefType.REF);
        parameter.setRef("node_llm.output");
        node.setParameters(Collections.singletonList(parameter));
        Chain chain = chain(Map.of(
                "node_llm.output", "回答内容"));

        Map<String, Object> result = node.execute(chain);

        Assert.assertEquals(
                "模型输出：回答内容",
                result.get("finalContent"));
    }

    /**
     * 创建指定输入参数的内容模板节点。
     *
     * @param template 模板内容
     * @param parameterNames 输入参数名称
     * @return 内容模板节点
     */
    private TemplateNode templateNode(
            String template,
            String... parameterNames) {
        TemplateNode node = new TemplateNode();
        node.setId("template-node");
        node.setName("内容模板");
        node.setTemplate(template);
        node.setParameters(Arrays.stream(parameterNames)
                .map(Parameter::new)
                .toList());
        node.setOutputDefs(Collections.singletonList(
                new Parameter("finalContent")));
        return node;
    }

    /**
     * 创建带初始化状态和输入变量的工作流。
     *
     * @param inputs 工作流输入
     * @return 工作流
     */
    private Chain chain(Map<String, Object> inputs) {
        Chain chain = new Chain(
                new ChainDefinition(),
                "template-node-" + UUID.randomUUID());
        chain.setChainStateRepository(
                new InMemoryChainStateRepository());
        chain.setNodeStateRepository(
                new InMemoryNodeStateRepository());
        ChainState state = chain.initializeState();
        state.getMemory().putAll(inputs);
        return chain;
    }
}
