package tech.easyflow.ai.easyagentsflow.service;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mockito;
import tech.easyflow.datacenter.entity.DatacenterTable;
import tech.easyflow.datacenter.entity.DatacenterTableField;
import tech.easyflow.datacenter.meta.entity.DatacenterSource;
import tech.easyflow.datacenter.meta.service.DatacenterDatasetRegistryService;

import java.lang.reflect.Field;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;

/**
 * 工作流数据中枢内容准备服务测试。
 */
public class WorkflowDatacenterContentServiceTest {

    private static final BigInteger SOURCE_ID = BigInteger.valueOf(1001L);
    private static final BigInteger TABLE_ID = BigInteger.valueOf(2001L);

    private WorkflowDatacenterContentService service;
    private DatacenterDatasetRegistryService registryService;

    /**
     * 初始化服务及数据源元数据桩。
     *
     * @throws Exception 注入测试依赖失败时抛出
     */
    @Before
    public void setUp() throws Exception {
        service = new WorkflowDatacenterContentService();
        registryService = Mockito.mock(DatacenterDatasetRegistryService.class);
        injectField(service, "registryService", registryService);

        DatacenterSource source = Mockito.mock(DatacenterSource.class);
        Mockito.when(source.getSourceName()).thenReturn("ama 实验基线模型预算");
        Mockito.when(source.getSourceType()).thenReturn("EXCEL");

        DatacenterTableField modelId = mockField("col_id", "模型ID", "VARCHAR");
        DatacenterTableField inputPrice = mockField("token", "输入价格", "DECIMAL");
        DatacenterTable table = Mockito.mock(DatacenterTable.class);
        Mockito.when(table.getId()).thenReturn(TABLE_ID);
        Mockito.when(table.getTableName()).thenReturn("Sheet1");
        Mockito.when(table.getFields()).thenReturn(List.of(modelId, inputPrice));

        Mockito.when(registryService.getSourceRequired(SOURCE_ID)).thenReturn(source);
        Mockito.when(registryService.listManagedTables(SOURCE_ID, null))
                .thenReturn(new ArrayList<>(List.of(table)));
        Mockito.when(registryService.getTableWithFields(TABLE_ID)).thenReturn(table);
    }

    /**
     * 验证查询上下文会自动且幂等地绑定到系统提示词。
     */
    @Test
    public void testPrepareRootShouldBindQueryContextToSystemPromptIdempotently() {
        JSONObject root = buildWorkflowRoot();
        JSONObject llmData = root.getJSONArray("nodes").getJSONObject(1).getJSONObject("data");

        service.prepareRoot(root);
        service.prepareRoot(root);

        String systemPrompt = llmData.getString("systemPrompt");
        Assert.assertEquals("请根据问题生成查询语句\n\n{{queryDataContext}}", systemPrompt);
        Assert.assertEquals(1, countOccurrences(systemPrompt, "{{queryDataContext}}"));
        JSONObject contextParameter = findParameter(llmData, "queryDataContext");
        Assert.assertNotNull(contextParameter);
        String contextValue = contextParameter.getString("value");
        Assert.assertTrue(contextValue.contains("只输出 SQL"));
        Assert.assertTrue(contextValue.contains("Sheet1"));
        Assert.assertTrue(contextValue.contains("col_id"));
        Assert.assertTrue(contextValue.contains("模型ID"));
        Assert.assertTrue(contextValue.contains("token AS input_price"));
    }

    /**
     * 验证关闭查询上下文后会同步清理参数和系统提示词占位符。
     */
    @Test
    public void testPrepareRootShouldRemoveQueryContextBindingWhenDisabled() {
        JSONObject root = buildWorkflowRoot();
        JSONObject llmData = root.getJSONArray("nodes").getJSONObject(1).getJSONObject("data");
        service.prepareRoot(root);
        llmData.put("queryContextNodeIds", new JSONArray());

        service.prepareRoot(root);

        Assert.assertEquals("请根据问题生成查询语句", llmData.getString("systemPrompt"));
        Assert.assertNull(findParameter(llmData, "queryDataContext"));
    }

    /**
     * 构造包含查询节点和大模型节点的最小工作流。
     *
     * @return 工作流根对象
     */
    private JSONObject buildWorkflowRoot() {
        JSONObject datasetRef = new JSONObject();
        datasetRef.put("sourceId", SOURCE_ID);

        JSONObject queryData = new JSONObject();
        queryData.put("datasetRef", datasetRef);
        queryData.put("querySql", "{{query}}");
        JSONObject queryNode = buildNode(
                "query-node",
                WorkflowDatacenterContentService.SEARCH_NODE_TYPE,
                queryData);

        JSONArray queryContextNodeIds = new JSONArray();
        queryContextNodeIds.add("query-node");
        JSONObject llmData = new JSONObject();
        llmData.put("systemPrompt", "请根据问题生成查询语句");
        llmData.put("queryContextNodeIds", queryContextNodeIds);
        llmData.put("parameters", new JSONArray());
        JSONObject llmNode = buildNode(
                "llm-node",
                WorkflowDatacenterContentService.LLM_NODE_TYPE,
                llmData);

        JSONArray nodes = new JSONArray();
        nodes.add(queryNode);
        nodes.add(llmNode);
        JSONObject root = new JSONObject();
        root.put("nodes", nodes);
        return root;
    }

    /**
     * 构造工作流节点。
     *
     * @param id 节点标识
     * @param type 节点类型
     * @param data 节点数据
     * @return 工作流节点
     */
    private JSONObject buildNode(String id, String type, JSONObject data) {
        JSONObject node = new JSONObject();
        node.put("id", id);
        node.put("type", type);
        node.put("data", data);
        return node;
    }

    /**
     * 构造字段元数据桩。
     *
     * @param fieldName 字段名
     * @param fieldDesc 字段描述
     * @param jdbcType JDBC 类型
     * @return 字段元数据
     */
    private DatacenterTableField mockField(String fieldName, String fieldDesc, String jdbcType) {
        DatacenterTableField field = Mockito.mock(DatacenterTableField.class);
        Mockito.when(field.getFieldName()).thenReturn(fieldName);
        Mockito.when(field.getFieldDesc()).thenReturn(fieldDesc);
        Mockito.when(field.getJdbcType()).thenReturn(jdbcType);
        return field;
    }

    /**
     * 查找指定名称的节点参数。
     *
     * @param data 节点数据
     * @param name 参数名
     * @return 参数对象，不存在时返回 {@code null}
     */
    private JSONObject findParameter(JSONObject data, String name) {
        JSONArray parameters = data.getJSONArray("parameters");
        if (parameters == null) {
            return null;
        }
        for (int i = 0; i < parameters.size(); i++) {
            JSONObject parameter = parameters.getJSONObject(i);
            if (parameter != null && name.equals(parameter.getString("name"))) {
                return parameter;
            }
        }
        return null;
    }

    /**
     * 统计子串出现次数。
     *
     * @param value 原始文本
     * @param target 目标子串
     * @return 出现次数
     */
    private int countOccurrences(String value, String target) {
        int count = 0;
        int index = 0;
        while ((index = value.indexOf(target, index)) >= 0) {
            count++;
            index += target.length();
        }
        return count;
    }

    /**
     * 注入服务私有依赖。
     *
     * @param target 目标对象
     * @param fieldName 字段名
     * @param value 字段值
     * @throws Exception 反射注入失败时抛出
     */
    private void injectField(Object target, String fieldName, Object value) throws Exception {
        Field field = target.getClass().getDeclaredField(fieldName);
        field.setAccessible(true);
        field.set(target, value);
    }
}
