package cn.com.poc.thirdparty.resource.demand.ai.aggregate.impl;

import cn.com.poc.thirdparty.resource.demand.ai.constants.KnowledgeSearchTypeEnum;
import cn.com.poc.thirdparty.resource.demand.ai.entity.qaknowledge.QAKnowledgeConfig;
import cn.com.poc.thirdparty.resource.demand.dgTools.DgtoolsAbstractHttpClient;
import cn.com.poc.thirdparty.resource.demand.ai.route.DgtoolsApiRoute;
import cn.com.poc.thirdparty.resource.demand.dgTools.result.AbstractResult;
import cn.com.poc.thirdparty.resource.demand.ai.aggregate.DemandKnowledgeService;
import cn.com.poc.thirdparty.resource.demand.ai.entity.knowledge.*;
import cn.com.poc.thirdparty.resource.demand.member.service.DemandAuthService;
import cn.com.yict.framemax.core.config.Config;
import cn.com.yict.framemax.core.i18n.I18nMessageException;
import cn.com.yict.framemax.data.model.PagingInfo;
import cn.hutool.core.lang.Assert;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.http.Header;
import org.apache.http.message.BasicHeader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Service
public class DemandKnowledgeServiceImpl implements DemandKnowledgeService {


    private Logger logger = LoggerFactory.getLogger(DemandKnowledgeService.class);
    @Resource
    private DgtoolsAbstractHttpClient dgToolsAbstractHttpClient;

    @Resource
    private DemandAuthService demandAuthService;

    @Override
    public String trainKnowledge(String fileURL, SegmentationConfigRequest segmentationConfig) {
        Assert.notBlank(fileURL);
        TrainKnowledgeRequest request = new TrainKnowledgeRequest();
        request.setDocumentUrl(fileURL);
        request.setSegmentationConfig(segmentationConfig);
        TrainKnowledgeResult trainKnowledgeResult = dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.TRAIN_KNOWLEDGE, request, getHeaders());
        if (null == trainKnowledgeResult) {
            throw new I18nMessageException("exception/abnormal.knowledge.base.training");
        }
        return trainKnowledgeResult.getKnowledgeId();
    }

    @Override
    public String trainKnowledgeEvent(String fileURL, String knowledgeType, SegmentationConfigRequest segmentationConfig, List<QAKnowledgeConfig> qaKnowledgeConfigs) {
        Assert.notBlank(fileURL);
        TrainKnowledgeRequest request = new TrainKnowledgeRequest();
        request.setDocumentUrl(fileURL);
        request.setSegmentationConfig(segmentationConfig);
        request.setKnowledgeType(knowledgeType);
        request.setQaKnowledgeConfigs(qaKnowledgeConfigs);
        TrainKnowledgeResult trainKnowledgeResult = dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.TRAIN_KNOWLEDGE_EVENT, request, getHeaders());
        if (null == trainKnowledgeResult) {
            throw new I18nMessageException("exception/abnormal.knowledge.base.training");
        }
        return trainKnowledgeResult.getKnowledgeId();
    }

    @Override
    public String trainKnowledgeStatus(String knowledgeId) {
        Assert.notBlank(knowledgeId);
        TrainKnowledgeStatusRequest request = new TrainKnowledgeStatusRequest();
        request.setKnowledgeId(knowledgeId);
        TrainKnowledgeStatusResult trainKnowledgeStatusResult = dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.TRAIN_KNOWLEDGE_STATUS, request, getHeaders());
        if (null == trainKnowledgeStatusResult) {
            throw new I18nMessageException("exception/abnormal.training.status.of.knowledge.base.acquisition");
        }
        return trainKnowledgeStatusResult.getTrainStatus();
    }

    @Override
    public void delKnowledge(String knowledgeId) {
        Assert.notBlank(knowledgeId);
        DelKnowledgeRequest request = new DelKnowledgeRequest();
        request.setKnowledgeId(knowledgeId);
        AbstractResult abstractResult = dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.DEL_KNOWLEDGE, request, getHeaders());
        if (null == abstractResult) {
            throw new I18nMessageException("exception/delete.knowledge.base.exception");
        }
    }

    @Override
    public SearchKnowledgeResult searchKnowledge(String query, List<String> knowledgeIds, Integer topK, Double score, KnowledgeSearchTypeEnum searchTypeEnum) {
        Assert.notBlank(query);
        SearchKnowledgeResult result = new SearchKnowledgeResult();
        result.setDocuments(new ArrayList<>());
        result.setScore(new ArrayList<>());
        result.setKnowledgeIds(new ArrayList<>());
        if (CollectionUtils.isEmpty(knowledgeIds)) {
            return new SearchKnowledgeResult();
        }
        if (topK == null) {
            topK = 3;
        }
        SearchKnowledgeRequest searchKnowledgeRequest = new SearchKnowledgeRequest();
        searchKnowledgeRequest.setQuery(query);
        searchKnowledgeRequest.setKnowLedgeIds(knowledgeIds);
        searchKnowledgeRequest.setTopK(topK);
        searchKnowledgeRequest.setScore(score);
        VectorSearchConfig vectorSearchConfig = new VectorSearchConfig();
        vectorSearchConfig.setSearchType(searchTypeEnum);
        vectorSearchConfig.setApiKey(Config.get("large-model.apikey"));
        searchKnowledgeRequest.setVectorSearchConfig(vectorSearchConfig);
        result = dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.SEARCH_KNOWLEDGE, searchKnowledgeRequest, getHeaders());
        if (null == result) {
            logger.warn("query knowledge base exception");
        }
        return result;
    }

    @Override
    public GetKnowledgeChunkInfoResult getKnowledgeChunkInfos(List<String> knowledgeIds, String query, PagingInfo pagingInfo) {
        GetKnowledgeChunkInfoRequest request = new GetKnowledgeChunkInfoRequest();
        request.setKnowledgeIds(knowledgeIds);
        request.setQuery(query);
        return dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.GET_KNOWLEDGE_CHUNK_INFOS, request, getHeaders(), pagingInfo);
    }

    @Override
    public void openKnowledgeChunk(String knowledgeId, String chunkRelationId, String isOpen) {
        UpsertChunkInfoRequest request = new UpsertChunkInfoRequest();
        request.setKnowledgeId(knowledgeId);
        request.setChunkRelationId(chunkRelationId);
        request.setIsOpen(isOpen);
        dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.OPEN_KNOWLEDGE_CHUNK, request, getHeaders());
    }

    @Override
    public void deleteKnowledgeChunk(String knowledgeId, String chunkRelationId) {
        UpsertChunkInfoRequest request = new UpsertChunkInfoRequest();
        request.setKnowledgeId(knowledgeId);
        request.setChunkRelationId(chunkRelationId);
        dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.DELETE_KNOWLEDGE_CHUNK, request, getHeaders());
    }

    @Override
    public void updateKnowledgeChunkDoc(String knowledgeId, String chunkRelationId, String content) {
        UpsertChunkInfoRequest request = new UpsertChunkInfoRequest();
        request.setKnowledgeId(knowledgeId);
        request.setChunkRelationId(chunkRelationId);
        request.setChunkContent(content);
        dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.UPDATE_KNOWLEDGE_CHUNK_DOC, request, getHeaders());
    }

    @Override
    public void addKnowledgeChunk(String knowledgeId, String content, Integer chunkSort) {
        UpsertChunkInfoRequest request = new UpsertChunkInfoRequest();
        request.setKnowledgeId(knowledgeId);
        request.setChunkSort(chunkSort);
        request.setChunkContent(content);
        dgToolsAbstractHttpClient.doRequest(DgtoolsApiRoute.DgtoolsAI.ADD_KNOWLEDGE_CHUNK, request, getHeaders());
    }

    private List<Header> getHeaders() {
        List<Header> headers = new ArrayList<>();
        headers.add(DgtoolsApiRoute.JSON_HEADER);
        headers.add(DgtoolsApiRoute.AI_HEADER);
        headers.add(new BasicHeader(DgtoolsApiRoute.HEADER_X_PLATFORM_AUTHORIZATION, demandAuthService.getToken()));
        return headers;
    }
}
