package tech.easyflow.ai.easyagents.tool;

import com.easyagents.core.document.Document;
import com.easyagents.core.model.chat.tool.BaseTool;
import com.easyagents.core.model.chat.tool.Parameter;
import com.easyagents.rag.retrieval.RetrievalMode;
import tech.easyflow.ai.chattime.availability.ChatTimeToolAvailabilityContext;
import tech.easyflow.ai.chattime.availability.ChatTimeToolAvailabilityService;
import tech.easyflow.ai.entity.DocumentCollection;
import tech.easyflow.ai.rag.KnowledgeRetrievalRequest;
import tech.easyflow.ai.service.DocumentCollectionService;
import tech.easyflow.common.util.SpringContextUtil;
import tech.easyflow.common.web.exceptions.BusinessException;

import java.math.BigInteger;
import java.util.List;
import java.util.Map;

/**
 * 知识库聊天工具。
 */
public class DocumentCollectionTool extends BaseTool {

    private BigInteger knowledgeId;
    private RetrievalMode retrievalMode = RetrievalMode.HYBRID;
    private ChatTimeToolAvailabilityContext chatTimeContext;

    /**
     * 默认构造器。
     */
    public DocumentCollectionTool() {
    }

    /**
     * 基于知识库实体构造聊天工具。
     *
     * @param documentCollection 知识库
     * @param needEnglishName 是否使用英文名
     */
    public DocumentCollectionTool(DocumentCollection documentCollection, boolean needEnglishName) {
        this(documentCollection, needEnglishName, RetrievalMode.HYBRID);
    }

    /**
     * 基于知识库实体构造聊天工具。
     *
     * @param documentCollection 知识库
     * @param needEnglishName 是否使用英文名
     * @param retrievalMode 检索模式
     */
    public DocumentCollectionTool(DocumentCollection documentCollection, boolean needEnglishName, RetrievalMode retrievalMode) {
        this(documentCollection, needEnglishName, retrievalMode, null);
    }

    /**
     * 基于知识库实体和聊天时权限上下文构造聊天工具。
     *
     * @param documentCollection 知识库
     * @param needEnglishName 是否使用英文名
     * @param retrievalMode 检索模式
     * @param chatTimeContext 聊天时权限上下文
     */
    public DocumentCollectionTool(DocumentCollection documentCollection,
                                  boolean needEnglishName,
                                  RetrievalMode retrievalMode,
                                  ChatTimeToolAvailabilityContext chatTimeContext) {
        this.knowledgeId = documentCollection.getId();
        this.retrievalMode = retrievalMode == null ? RetrievalMode.HYBRID : retrievalMode;
        this.chatTimeContext = chatTimeContext;
        this.name = ChatToolNameHelper.resolveToolName(
                needEnglishName,
                documentCollection.getEnglishName(),
                documentCollection.getTitle(),
                "knowledge",
                documentCollection.getId()
        );
        this.description = documentCollection.getDescription();
        this.parameters = getDefaultParameters();
    }


    public Parameter[] getDefaultParameters() {
        Parameter parameter = new Parameter();
        parameter.setName("input");
        parameter.setDescription("要查询的相关知识");
        parameter.setType("string");
        parameter.setRequired(true);
        return new Parameter[]{parameter};
    }

    public BigInteger getKnowledgeId() {
        return knowledgeId;
    }

    public void setKnowledgeId(BigInteger knowledgeId) {
        this.knowledgeId = knowledgeId;
    }

    public RetrievalMode getRetrievalMode() {
        return retrievalMode;
    }

    public void setRetrievalMode(RetrievalMode retrievalMode) {
        this.retrievalMode = retrievalMode == null ? RetrievalMode.HYBRID : retrievalMode;
    }

    /**
     * 获取聊天时权限上下文。
     *
     * @return 聊天时权限上下文
     */
    public ChatTimeToolAvailabilityContext getChatTimeContext() {
        return chatTimeContext;
    }

    /**
     * 设置聊天时权限上下文。
     *
     * @param chatTimeContext 聊天时权限上下文
     */
    public void setChatTimeContext(ChatTimeToolAvailabilityContext chatTimeContext) {
        this.chatTimeContext = chatTimeContext;
    }

    /**
     * 执行知识库检索。
     *
     * @param argsMap 工具入参
     * @return 检索结果拼接文本
     */
    @Override
    public Object invoke(Map<String, Object> argsMap) {
        DocumentCollectionService knowledgeService = SpringContextUtil.getBean(DocumentCollectionService.class);
        DocumentCollection knowledge = null;
        if (this.knowledgeId != null) {
            knowledge = knowledgeService.getById(this.knowledgeId);
        }
        if (knowledge == null) {
            throw new BusinessException("知识库不存在");
        }
        assertChatTimeAvailability(knowledge);

        KnowledgeRetrievalRequest request = new KnowledgeRetrievalRequest();
        request.setKnowledgeId(this.knowledgeId);
        request.setQuery((String) argsMap.get("input"));
        request.setRetrievalMode(this.retrievalMode);
        request.setCallerType("BOT_TOOL");
        request.setCallerId(this.knowledgeId == null ? null : this.knowledgeId.toString());
        List<Document> documents = knowledgeService.search(request);

        StringBuilder sb = new StringBuilder();
        if (documents != null) {
            for (int i = 0; i < documents.size(); i++) {
                Document document = documents.get(i);
                if (document == null || document.getContent() == null) {
                    continue;
                }
                if (sb.length() > 0) {
                    sb.append("\n\n---\n\n");
                }
                sb.append(document.getContent());
            }
        }
        return sb.toString();
    }

    /**
     * 当工具由聊天运行时装配时，执行聊天态权限兜底。
     *
     * @param knowledge 当前知识库实体
     */
    protected void assertChatTimeAvailability(DocumentCollection knowledge) {
        if (chatTimeContext == null) {
            return;
        }
        ChatTimeToolAvailabilityService availabilityService = SpringContextUtil.getBean(ChatTimeToolAvailabilityService.class);
        availabilityService.assertAvailable(chatTimeContext, knowledge, "当前用户无权在聊天中访问该知识库");
    }

}
