package com.purui.prims.headless.chat.parser.llm;


import com.purui.prims.common.util.ContextUtils;
import com.purui.prims.headless.chat.parser.SemanticParser;
import com.purui.prims.headless.chat.ChatContext;
import com.purui.prims.headless.chat.ChatQueryContext;
import com.purui.prims.headless.chat.query.llm.s2sql.LLMReq;
import com.purui.prims.headless.chat.query.llm.s2sql.LLMResp;
import com.purui.prims.headless.chat.query.llm.s2sql.LLMSqlResp;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.MapUtils;

/**
 * LLMSqlParser uses large language model to understand query semantics and
 * generate S2SQL statements to be executed by the semantic query engine.
 */
@Slf4j
public class LLMSqlParser implements SemanticParser {

    @Override
    public void parse(ChatQueryContext queryCtx, ChatContext chatCtx) {
        try {
            LLMRequestService requestService = ContextUtils.getBean(LLMRequestService.class);
            //1.determine whether to skip this parser.
            if (requestService.isSkip(queryCtx)) {
                return;
            }
            //2.get dataSetId from queryCtx and chatCtx.
            Long dataSetId = requestService.getDataSetId(queryCtx);
            if (dataSetId == null) {
                return;
            }
            log.info("Generate query statement for dataSetId:{}", dataSetId);

            //3.invoke LLM service to do parsing.
            tryParse(queryCtx, dataSetId);
        } catch (Exception e) {
            log.error("Failed to parse query:", e);
        }
    }

    private void tryParse(ChatQueryContext queryCtx, Long dataSetId) {
        LLMRequestService requestService = ContextUtils.getBean(LLMRequestService.class);
        LLMResponseService responseService = ContextUtils.getBean(LLMResponseService.class);
        int maxRetries = ContextUtils.getBean(LLMParserConfig.class).getRecallMaxRetries();

        LLMReq llmReq = requestService.getLlmReq(queryCtx, dataSetId);

        int currentRetry = 0;
        Map<String, LLMSqlResp> sqlRespMap = new HashMap<>();
        ParseResult parseResult = null;
        while (currentRetry < maxRetries) {
            log.info("currentRetry:{},start runText2SQL", currentRetry);
            try {
                LLMResp llmResp = requestService.runText2SQL(llmReq);
                if (Objects.nonNull(llmResp)) {
                    //deduplicate the S2SQL result list and build parserInfo
                    sqlRespMap = responseService.getDeduplicationSqlResp(currentRetry, llmResp);
                    if (MapUtils.isNotEmpty(sqlRespMap)) {
                        parseResult = ParseResult.builder().dataSetId(dataSetId).llmReq(llmReq)
                                .llmResp(llmResp).linkingValues(llmReq.getLinking()).build();
                        break;
                    }
                }
            } catch (Exception e) {
                log.error("currentRetry:{},runText2SQL error", currentRetry, e);
            }
            currentRetry++;
        }
        if (MapUtils.isEmpty(sqlRespMap)) {
            return;
        }
        for (Entry<String, LLMSqlResp> entry : sqlRespMap.entrySet()) {
            String sql = entry.getKey();
            double sqlWeight = entry.getValue().getSqlWeight();
            responseService.addParseInfo(queryCtx, parseResult, sql, sqlWeight);
        }
    }

}
