package com.purui.prims.headless.chat.query.llm.s2sql;

import com.purui.prims.auth.api.authentication.pojo.User;
import com.purui.prims.headless.api.pojo.DataSetSchema;
import com.purui.prims.headless.api.pojo.SqlInfo;
import com.purui.prims.headless.api.pojo.request.SemanticQueryReq;
import com.purui.prims.headless.chat.query.QueryManager;
import com.purui.prims.headless.chat.query.llm.LLMSemanticQuery;
import com.purui.prims.headless.chat.utils.QueryReqBuilder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

@Slf4j
@Component
public class LLMSqlQuery extends LLMSemanticQuery {

    public static final String QUERY_MODE = "LLM_S2SQL";

    public LLMSqlQuery() {
        QueryManager.register(this);
    }

    @Override
    public String getQueryMode() {
        return QUERY_MODE;
    }

    @Override
    public SemanticQueryReq buildSemanticQueryReq() {
        return QueryReqBuilder.buildS2SQLReq(parseInfo.getSqlInfo(), parseInfo.getDataSetId());
    }

    @Override
    public void initS2Sql(DataSetSchema semanticSchema, User user) {
        SqlInfo sqlInfo = parseInfo.getSqlInfo();
        sqlInfo.setCorrectS2SQL(sqlInfo.getS2SQL());
    }
}
