package com.purui.prims.headless.chat.corrector;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.purui.prims.common.pojo.Constants;
import com.purui.prims.headless.api.pojo.DataSetSchema;
import com.purui.prims.headless.api.pojo.QueryConfig;
import com.purui.prims.headless.api.pojo.SchemaElement;
import com.purui.prims.headless.api.pojo.SemanticParseInfo;
import com.purui.prims.headless.api.pojo.SemanticSchema;
import com.purui.prims.headless.api.pojo.SqlInfo;
import com.purui.prims.headless.chat.parser.llm.ParseResult;
import com.purui.prims.headless.chat.ChatQueryContext;
import com.purui.prims.headless.chat.query.llm.s2sql.LLMReq;
import org.junit.Assert;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

class SchemaCorrectorTest {

    private String json = "{\n"
            + "          \"dataSetId\":  1,\n"
            + "          \"llmReq\":  {\n"
            + "                    \"queryText\":  \"xxx2024年播放量最高的十首歌\",\n"
            + "                    \"filterCondition\":  {\n"
            + "                              \"tableName\":  null\n"
            + "                    },\n"
            + "                    \"schema\":  {\n"
            + "                              \"domainName\":  \"歌曲\",\n"
            + "                              \"dataSetName\":  \"歌曲\",\n"
            + "                              \"fieldNameList\":  [\n"
            + "                                        \"商务组\",\n"
            + "                                        \"歌曲名\",\n"
            + "                                        \"播放量\",\n"
            + "                                        \"播放份额\",\n"
            + "                                        \"数据日期\"\n"
            + "                              ]\n"
            + "                    },\n"
            + "                    \"linking\":  [\n"
            + "\n"
            + "                    ],\n"
            + "                    \"currentDate\":  \"2024-02-24\",\n"
            + "                    \"priorExts\":  \"播放份额是小数; \",\n"
            + "                    \"sqlGenType\":  \"1_pass_self_consistency\"\n"
            + "          },\n"
            + "          \"request\":  null,\n"
            + "          \"linkingValues\":  [\n"
            + "\n"
            + "          ]\n"
            + "}";

    @Test
    void doCorrect() throws JsonProcessingException {
        Long dataSetId = 1L;
        ChatQueryContext queryContext = buildQueryContext(dataSetId);
        ObjectMapper objectMapper = new ObjectMapper();
        ParseResult parseResult = objectMapper.readValue(json, ParseResult.class);


        String sql = "select  歌曲名 from 歌曲 where 发行日期 >= '2024-01-01' "
                + "and 商务组 = 'xxx' order by 播放量 desc  limit 10";
        SemanticParseInfo semanticParseInfo = new SemanticParseInfo();
        SqlInfo sqlInfo = new SqlInfo();
        sqlInfo.setS2SQL(sql);
        sqlInfo.setCorrectS2SQL(sql);
        semanticParseInfo.setSqlInfo(sqlInfo);

        SchemaElement schemaElement = new SchemaElement();
        schemaElement.setDataSet(dataSetId);
        semanticParseInfo.setDataSet(schemaElement);


        semanticParseInfo.getProperties().put(Constants.CONTEXT, parseResult);

        SchemaCorrector schemaCorrector = new SchemaCorrector();
        schemaCorrector.removeFilterIfNotInLinkingValue(queryContext, semanticParseInfo);

        Assert.assertEquals("SELECT 歌曲名 FROM 歌曲 WHERE 发行日期 >= '2024-01-01' "
                + "ORDER BY 播放量 DESC LIMIT 10", semanticParseInfo.getSqlInfo().getCorrectS2SQL());

        parseResult = objectMapper.readValue(json, ParseResult.class);

        List<LLMReq.ElementValue> linkingValues = new ArrayList<>();
        LLMReq.ElementValue elementValue = new LLMReq.ElementValue();
        elementValue.setFieldName("商务组");
        elementValue.setFieldValue("xxx");
        linkingValues.add(elementValue);
        parseResult.setLinkingValues(linkingValues);
        semanticParseInfo.getProperties().put(Constants.CONTEXT, parseResult);

        semanticParseInfo.getSqlInfo().setCorrectS2SQL(sql);
        semanticParseInfo.getSqlInfo().setS2SQL(sql);
        schemaCorrector.removeFilterIfNotInLinkingValue(queryContext, semanticParseInfo);
        Assert.assertEquals("SELECT 歌曲名 FROM 歌曲 WHERE 发行日期 >= '2024-01-01' "
                + "AND 商务组 = 'xxx' ORDER BY 播放量 DESC LIMIT 10", semanticParseInfo.getSqlInfo().getCorrectS2SQL());

    }

    private ChatQueryContext buildQueryContext(Long dataSetId) {
        ChatQueryContext queryContext = new ChatQueryContext();
        List<DataSetSchema> dataSetSchemaList = new ArrayList<>();
        DataSetSchema dataSetSchema = new DataSetSchema();
        QueryConfig queryConfig = new QueryConfig();
        dataSetSchema.setQueryConfig(queryConfig);
        SchemaElement schemaElement = new SchemaElement();
        schemaElement.setDataSet(dataSetId);
        dataSetSchema.setDataSet(schemaElement);
        Set<SchemaElement> dimensions = new HashSet<>();
        SchemaElement element1 = new SchemaElement();
        element1.setDataSet(1L);
        element1.setName("歌曲名");
        dimensions.add(element1);

        SchemaElement element2 = new SchemaElement();
        element2.setDataSet(1L);
        element2.setName("商务组");
        dimensions.add(element2);

        SchemaElement element3 = new SchemaElement();
        element3.setDataSet(1L);
        element3.setName("发行日期");
        dimensions.add(element3);

        dataSetSchema.setDimensions(dimensions);
        dataSetSchemaList.add(dataSetSchema);

        SemanticSchema semanticSchema = new SemanticSchema(dataSetSchemaList);
        queryContext.setSemanticSchema(semanticSchema);
        return queryContext;
    }
}