package tech.easyflow.ai.easyagentsflow.knowledge;

import com.easyagents.core.document.Document;
import com.alibaba.fastjson2.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.knowledge.Knowledge;
import com.easyagents.flow.core.knowledge.KnowledgeProvider;
import com.easyagents.flow.core.node.KnowledgeNode;
import org.springframework.stereotype.Component;
import tech.easyflow.ai.rag.KnowledgeRetrievalRequest;
import tech.easyflow.ai.rag.KnowledgeRetrievalModes;
import tech.easyflow.ai.service.DocumentCollectionService;

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

/**
 * 为工作流知识库节点提供检索能力。
 */
@Component
public class KnowledgeProviderImpl implements KnowledgeProvider {

    @Resource
    private DocumentCollectionService documentCollectionService;

    /**
     * 获取知识库检索器。
     *
     * @param id 知识库 ID
     * @return 知识库检索器
     */
    @Override
    public Knowledge getKnowledge(Object id) {
        return new Knowledge() {
            /**
             * {@inheritDoc}
             */
            @Override
            public List<Map<String, Object>> search(
                    String keyword,
                    int limit,
                    KnowledgeNode knowledgeNode,
                    Chain chain) {
                KnowledgeRetrievalRequest request = new KnowledgeRetrievalRequest();
                request.setKnowledgeId(new BigInteger(id.toString()));
                request.setQuery(keyword);
                request.setLimit(limit);
                request.setRetrievalMode(KnowledgeRetrievalModes.parse(knowledgeNode.getRetrievalMode()));
                request.setCallerType("WORKFLOW");
                request.setCallerId(knowledgeNode.getId());
                List<Document> documents = documentCollectionService.search(request);
                if (limit > 0 && documents.size() > limit) {
                    documents = new ArrayList<>(documents.subList(0, limit));
                }
                List<Map<String, Object>> res = new ArrayList<>();
                for (Document document : documents) {
                    res.add(toWorkflowDocument(document, id));
                }
                return res;
            }
        };
    }

    /**
     * 将检索文档转换为工作流稳定对象，并保留旧序列化字段。
     *
     * @param document 检索文档
     * @param knowledgeId 知识库 ID
     * @return 工作流文档对象
     */
    private Map<String, Object> toWorkflowDocument(
            Document document, Object knowledgeId) {
        JSONObject result = JSONObject.from(document);
        result.put("title", document.getTitle());
        result.put("content", document.getContent());
        result.put(
                "documentId",
                document.getMetadata("documentId", document.getId()));
        result.put("knowledgeId", knowledgeId);
        return result;
    }
}
