package com.pavis.admin.aigc.service.impl;

import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.pavis.admin.aigc.common.dto.EmbeddingR;
import com.pavis.admin.aigc.common.model.req.ChatReq;
import com.pavis.admin.aigc.constant.EmbedConst;
import com.pavis.admin.aigc.core.llm.exception.ServiceException;
import com.pavis.admin.aigc.core.llm.provider.EmbeddingProvider;
import com.pavis.admin.aigc.core.llm.provider.EmbeddingStoreFactory;
import com.pavis.admin.aigc.core.llm.provider.KnowledgeStoreFactory;
import com.pavis.admin.aigc.core.llm.provider.ModelStoreFactory;
import com.pavis.admin.aigc.mapper.DocChunkMapper;
import com.pavis.admin.aigc.model.entity.DocChunkDO;
import com.pavis.admin.aigc.model.resp.KnowledgeResp;
import com.pavis.admin.aigc.service.KnowledgeService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.UrlDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.annotation.Resource;
import lombok.RequiredArgsConstructor;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import org.springframework.transaction.annotation.Transactional;
import top.continew.starter.extension.crud.service.BaseServiceImpl;
import com.pavis.admin.aigc.mapper.DocMapper;
import com.pavis.admin.aigc.model.entity.DocDO;
import com.pavis.admin.aigc.model.query.DocQuery;
import com.pavis.admin.aigc.model.req.DocReq;
import com.pavis.admin.aigc.model.resp.DocDetailResp;
import com.pavis.admin.aigc.model.resp.DocResp;
import com.pavis.admin.aigc.service.DocService;

import java.util.ArrayList;
import java.util.List;

/**
 * 智能文档业务实现
 *
 * @author semi
 * @since 2025/05/26 17:28
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class DocServiceImpl extends BaseServiceImpl<DocMapper, DocDO, DocResp, DocDetailResp, DocQuery, DocReq> implements DocService {
    private final KnowledgeStoreFactory knowledgeStoreFactory;

    private final ModelStoreFactory modelStoreFactory;

    private final EmbeddingStoreFactory embeddingStoreFactory;

    private final DocMapper docMapper;

    private final DocChunkMapper docChunkMapper;

    @Resource
    private KnowledgeService aigcKnowledgeService;

    @Override
    public void embedDocsSlice(DocDO data, String url) {
        List<EmbeddingR> list = embeddingDocs(new ChatReq().setDocsName(data.getName())
            .setKnowledgeId(data.getKnowledgeId().toString())
            .setUrl(url));
        list.forEach(i -> {
            DocChunkDO docChunkDO = new DocChunkDO();
            docChunkDO.setKnowledgeId(data.getKnowledgeId());
            docChunkDO.setDocId(data.getId());
            docChunkDO.setEmbedStoreId(i.getVectorId());
            docChunkDO.setContent(i.getText());
            docChunkDO.setCreateUser(data.getCreateUser());
            docChunkDO.setWordNum(i.getText().length());
            docChunkDO.setName(data.getName());
            // knowledgeService.addDocsSlice(docChunkDO);
            docChunkMapper.insert(docChunkDO);
        });

        DocDO docDO = new DocDO();
        docDO.setId(data.getId());
        docDO.setVectorStatus("completed");
        docDO.setSliceNum(list.size());
        docDO.setUpdateUser(data.getCreateUser());
        baseMapper.updateById(docDO);
    }

    public List<EmbeddingR> embeddingDocs(ChatReq req) {
        log.info(">>>>>>>>>>>>>> Docs文档向量解析开始，KnowledgeId={}, DocsName={}", req.getKnowledgeId(), req.getDocsName());
        Document document = UrlDocumentLoader.load(req.getUrl(), new ApacheTikaDocumentParser());
        document.metadata().put(EmbedConst.KNOWLEDGE, req.getKnowledgeId()).put(EmbedConst.FILENAME, req.getDocsName());
        
        List<EmbeddingR> list = new ArrayList<>();
        try {
            DocumentSplitter splitter = EmbeddingProvider.splitter();
            List<TextSegment> segments = splitter.split(document);

            EmbeddingModel embeddingModel = getEmbeddingModel(Long.valueOf(req.getKnowledgeId()));
            EmbeddingStore<TextSegment> embeddingStore = getEmbeddingStore(Long.valueOf(req.getKnowledgeId()));
            List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
            List<String> ids = embeddingStore.addAll(embeddings, segments);

            for (int i = 0; i < ids.size(); i++) {
                list.add(new EmbeddingR().setVectorId(ids.get(i)).setText(segments.get(i).text()));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        log.info(">>>>>>>>>>>>>> Docs文档向量解析结束，KnowledgeId={}, DocsName={}", req.getKnowledgeId(), req.getDocsName());
        return list;
    }

    public EmbeddingModel getEmbeddingModel(Long knowledgeId) {
        if (knowledgeStoreFactory.containsKnowledge(knowledgeId)) {
            KnowledgeResp data = knowledgeStoreFactory.getKnowledge(knowledgeId);
            if (modelStoreFactory.containsEmbeddingModel(data.getEmbedModelId().toString())) {
                return modelStoreFactory.getEmbeddingModel(data.getEmbedModelId().toString());
            }
        }
        throw new ServiceException("没有找到匹配的Embedding向量数据库");
    }

    public EmbeddingStore<TextSegment> getEmbeddingStore(Long knowledgeId) {
        if (knowledgeStoreFactory.containsKnowledge(knowledgeId)) {
            KnowledgeResp data = knowledgeStoreFactory.getKnowledge(knowledgeId);
            if (embeddingStoreFactory.containsEmbeddingStore(data.getEmbedStoreId().toString())) {
                return embeddingStoreFactory.getEmbeddingStore(data.getEmbedStoreId().toString());
            }
        }
        throw new ServiceException("没有找到匹配的Embedding向量数据库");
    }


    @Override
    @Transactional
    public void clearDocSlices(String docsId) {
        if (StrUtil.isBlank(docsId)) {
            return;
        }
        // remove from embedding store
        List<String> vectorIds = aigcKnowledgeService.listSliceVectorIdsOfDoc(docsId);
        if (vectorIds.isEmpty()) {
            return;
        }
        DocDO docs = docMapper.selectById(docsId);
        EmbeddingStore<TextSegment> embeddingStore = getEmbeddingStore(docs.getKnowledgeId());
        embeddingStore.removeAll(vectorIds);
        // remove from docSlice
        aigcKnowledgeService.removeSlicesOfDoc(docsId);
    }
}