package com.purui.prims.headless.chat;

import com.fasterxml.jackson.annotation.JsonIgnore;
import com.purui.prims.auth.api.authentication.pojo.User;
import com.purui.prims.common.config.ChatModelConfig;
import com.purui.prims.common.config.LLMConfig;
import com.purui.prims.common.config.PromptConfig;
import com.purui.prims.common.pojo.enums.Text2SQLType;
import com.purui.prims.common.util.ContextUtils;
import com.purui.prims.headless.api.pojo.QueryDataType;
import com.purui.prims.headless.api.pojo.SchemaMapInfo;
import com.purui.prims.headless.api.pojo.SemanticSchema;
import com.purui.prims.headless.api.pojo.enums.ChatWorkflowState;
import com.purui.prims.headless.api.pojo.enums.MapModeEnum;
import com.purui.prims.headless.api.pojo.request.QueryFilters;
import com.purui.prims.headless.chat.parser.ParserConfig;
import com.purui.prims.headless.chat.query.SemanticQuery;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class ChatQueryContext {

    private String queryText;
    private Integer chatId;
    private Set<Long> dataSetIds;
    private Map<Long, List<Long>> modelIdToDataSetIds;
    private User user;
    private boolean saveAnswer;
    private Text2SQLType text2SQLType = Text2SQLType.RULE_AND_LLM;
    private QueryFilters queryFilters;
    private List<SemanticQuery> candidateQueries = new ArrayList<>();
    private SchemaMapInfo mapInfo = new SchemaMapInfo();
    private MapModeEnum mapModeEnum = MapModeEnum.STRICT;
    @JsonIgnore
    private SemanticSchema semanticSchema;
    @JsonIgnore
    private ChatWorkflowState workflowState;
    private QueryDataType queryDataType = QueryDataType.ALL;
    private LLMConfig llmConfig;
    private ChatModelConfig modelConfig;
    private PromptConfig promptConfig;

    public List<SemanticQuery> getCandidateQueries() {
        ParserConfig parserConfig = ContextUtils.getBean(ParserConfig.class);
        int parseShowCount = Integer.valueOf(parserConfig.getParameterValue(ParserConfig.PARSER_SHOW_COUNT));
        candidateQueries = candidateQueries.stream()
                .sorted(Comparator.comparing(semanticQuery -> semanticQuery.getParseInfo().getScore(),
                        Comparator.reverseOrder()))
                .limit(parseShowCount)
                .collect(Collectors.toList());
        return candidateQueries;
    }
}
