package tech.easyflow.ai.service.impl;

import com.easyagents.core.model.embedding.EmbeddingModel;
import com.easyagents.core.store.DocumentStore;
import com.easyagents.core.store.StoreOptions;
import com.easyagents.core.store.StoreResult;
import com.mybatisflex.core.query.QueryWrapper;
import org.springframework.stereotype.Service;
import tech.easyflow.ai.entity.DocumentChunk;
import tech.easyflow.ai.entity.DocumentCollection;
import tech.easyflow.ai.entity.FaqItem;
import tech.easyflow.ai.entity.Model;
import tech.easyflow.ai.service.DocumentChunkService;
import tech.easyflow.ai.service.DocumentCollectionService;
import tech.easyflow.ai.service.FaqItemService;
import tech.easyflow.ai.service.KnowledgeEmbeddingService;
import tech.easyflow.ai.service.ModelService;
import tech.easyflow.ai.support.DocumentStoreLifecycleSupport;
import tech.easyflow.common.web.exceptions.BusinessException;

import javax.annotation.Resource;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 知识库向量重建服务实现。
 */
@Service
public class KnowledgeEmbeddingServiceImpl implements KnowledgeEmbeddingService {

    @Resource
    private DocumentCollectionService documentCollectionService;
    @Resource
    private DocumentChunkService documentChunkService;
    @Resource
    private FaqItemService faqItemService;
    @Resource
    private ModelService modelService;

    @Override
    public void rebuildKnowledgeVectors(BigInteger knowledgeId) {
        DocumentCollection knowledge = documentCollectionService.getById(knowledgeId);
        if (knowledge == null) {
            throw new BusinessException("知识库不存在");
        }
        DocumentStore documentStore = knowledge.toDocumentStore();
        if (documentStore == null) {
            throw new BusinessException("知识库没有配置向量库");
        }
        try {
            Model model = modelService.getModelInstance(knowledge.getVectorEmbedModelId());
            if (model == null) {
                throw new BusinessException("知识库没有配置向量模型");
            }
            EmbeddingModel embeddingModel = model.toEmbeddingModel();
            documentStore.setEmbeddingModel(embeddingModel);
            StoreOptions storeOptions = StoreOptions.ofCollectionName(knowledge.getVectorStoreCollection());
            storeOptions.setIndexName(knowledge.getVectorStoreCollection());

            if (knowledge.isFaqCollection()) {
                rebuildFaqVectors(knowledge, documentStore, storeOptions, embeddingModel);
                return;
            }
            rebuildDocumentVectors(knowledge, documentStore, storeOptions, embeddingModel);
        } finally {
            DocumentStoreLifecycleSupport.closeQuietly(documentStore);
        }
    }

    private void rebuildDocumentVectors(
        DocumentCollection knowledge,
        DocumentStore documentStore,
        StoreOptions storeOptions,
        EmbeddingModel embeddingModel
    ) {
        QueryWrapper wrapper = QueryWrapper.create()
            .eq(DocumentChunk::getDocumentCollectionId, knowledge.getId())
            .orderBy("sorting asc");
        List<DocumentChunk> chunks = documentChunkService.list(wrapper);
        List<BigInteger> ids = new ArrayList<>();
        List<com.easyagents.core.document.Document> documents = new ArrayList<>();
        for (DocumentChunk chunk : chunks) {
            ids.add(chunk.getId());
            com.easyagents.core.document.Document document = com.easyagents.core.document.Document.of(chunk.getContent());
            document.setId(chunk.getId());
            documents.add(document);
        }
        rewriteStore(documentStore, storeOptions, ids, documents);
        updateKnowledgeEmbeddingState(knowledge, embeddingModel);
    }

    private void rebuildFaqVectors(
        DocumentCollection knowledge,
        DocumentStore documentStore,
        StoreOptions storeOptions,
        EmbeddingModel embeddingModel
    ) {
        QueryWrapper wrapper = QueryWrapper.create()
            .eq(FaqItem::getCollectionId, knowledge.getId())
            .orderBy("order_no asc");
        List<FaqItem> faqItems = faqItemService.list(wrapper);
        List<BigInteger> ids = new ArrayList<>();
        List<com.easyagents.core.document.Document> documents = new ArrayList<>();
        for (FaqItem faqItem : faqItems) {
            ids.add(faqItem.getId());
            StringBuilder content = new StringBuilder();
            content.append("问题：").append(faqItem.getQuestion());
            if (faqItem.getAnswerText() != null && !faqItem.getAnswerText().isBlank()) {
                content.append("\n答案：").append(faqItem.getAnswerText());
            }
            com.easyagents.core.document.Document document =
                com.easyagents.core.document.Document.of(content.toString());
            document.setId(faqItem.getId());
            Map<String, Object> metadata = new HashMap<>();
            metadata.put("question", faqItem.getQuestion());
            metadata.put("answerText", faqItem.getAnswerText());
            metadata.put("categoryId", faqItem.getCategoryId());
            document.setMetadataMap(metadata);
            documents.add(document);
        }
        rewriteStore(documentStore, storeOptions, ids, documents);
        updateKnowledgeEmbeddingState(knowledge, embeddingModel);
    }

    private void rewriteStore(
        DocumentStore documentStore,
        StoreOptions storeOptions,
        List<BigInteger> ids,
        List<com.easyagents.core.document.Document> documents
    ) {
        if (!ids.isEmpty()) {
            documentStore.delete(ids, storeOptions);
        }
        if (documents.isEmpty()) {
            return;
        }
        StoreResult result = documentStore.store(documents, storeOptions);
        if (result == null || !result.isSuccess()) {
            throw new BusinessException("知识库向量重建失败");
        }
    }

    private void updateKnowledgeEmbeddingState(DocumentCollection knowledge, EmbeddingModel embeddingModel) {
        DocumentCollection update = new DocumentCollection();
        update.setId(knowledge.getId());
        Map<String, Object> options = knowledge.getOptions() == null
            ? new HashMap<>()
            : new HashMap<>(knowledge.getOptions());
        options.put(DocumentCollection.KEY_CAN_UPDATE_EMBEDDING_MODEL, false);
        update.setOptions(options);
        if (knowledge.getDimensionOfVectorModel() == null) {
            update.setDimensionOfVectorModel(Model.getEmbeddingDimension(embeddingModel));
        }
        documentCollectionService.updateById(update);
    }
}
