package tech.easyflow.ai.entity;

import com.easyagents.core.model.chat.tool.Tool;
import com.easyagents.core.store.DocumentStore;
import com.easyagents.rag.retrieval.RetrievalMode;
import com.mybatisflex.annotation.Column;
import com.easyagents.store.milvus.MilvusVectorStore;
import com.easyagents.store.milvus.MilvusVectorStoreConfig;
import com.mybatisflex.annotation.Table;
import tech.easyflow.ai.config.AiMilvusConfig;
import tech.easyflow.ai.chattime.availability.ChatTimeToolAvailabilityContext;
import tech.easyflow.ai.easyagents.tool.DocumentCollectionTool;
import tech.easyflow.ai.entity.base.DocumentCollectionBase;
import tech.easyflow.ai.rag.KnowledgeRetrievalModes;
import tech.easyflow.common.util.SpringContextUtil;
import tech.easyflow.common.util.StringUtil;
import tech.easyflow.system.permission.resource.VisibilityResource;

import java.math.BigDecimal;
import java.util.Map;

/**
 * 实体类。
 *
 * @author michael
 * @since 2024-08-23
 */

@Table("tb_document_collection")
public class DocumentCollection extends DocumentCollectionBase implements VisibilityResource {

    @Column(ignore = true)
    private Boolean approvalPending;

    @Column(ignore = true)
    private String currentApprovalActionType;

    @Column(ignore = true)
    private String displayPublishStatus;

    @Column(ignore = true)
    private String createdByName;

    public static final String TYPE_DOCUMENT = "DOCUMENT";
    public static final String TYPE_FAQ = "FAQ";

    /**
     * 文档块问题集合配置key
     */
    public static final String KEY_CHUNK_QUESTION_VECTOR_STORE_COLLECTION = "chunkQuestionVectorStoreCollection";

    /**
     * 文档块摘要集合配置key
     */
    public static final String KEY_CHUNK_SUMMARY_VECTOR_STORE_COLLECTION = "chunkSummaryVectorStoreCollection";

    /**
     * 知识召回最大条数
     */
    public static final String KEY_DOC_RECALL_MAX_NUM = "docRecallMaxNum";

    /**
     * 相似度最小值
     */
    public static final String KEY_SIMILARITY_THRESHOLD = "simThreshold";

    /**
     * 是否允许更新向量模型
     */
    public static final String KEY_CAN_UPDATE_EMBEDDING_MODEL = "canUpdateEmbeddingModel";

    /**
     * 是否启用重排模型
     */
    public static final String KEY_RERANK_ENABLE = "rerankEnable";
    public static final String KEY_SPLITTER_DEFAULT_STRATEGY = "splitter.defaultStrategyCode";
    public static final String KEY_SPLITTER_AUTO_RECOMMEND_ENABLED = "splitter.autoRecommendEnabled";
    public static final String KEY_SPLITTER_FALLBACK_STRATEGY = "splitter.fallbackStrategyCode";
    public static final String KEY_SPLITTER_STRATEGY_PROFILES = "splitter.strategyProfiles";

    public DocumentStore toDocumentStore() {
        if (StringUtil.noText(this.getVectorStoreCollection())) {
            return null;
        }
        return milvusStore();
    }

    public boolean isVectorStoreEnabled() {
        return this.getVectorStoreEnable() != null && this.getVectorStoreEnable();
    }

    public boolean isFaqCollection() {
        return TYPE_FAQ.equalsIgnoreCase(this.getCollectionType());
    }

    public boolean isDocumentCollection() {
        return !isFaqCollection();
    }

    public boolean isSearchEngineEnabled() {
        return true;
    }

    private DocumentStore milvusStore() {
        AiMilvusConfig aiMilvusConfig = SpringContextUtil.getBean(AiMilvusConfig.class);
        MilvusVectorStoreConfig milvusVectorStoreConfig = aiMilvusConfig.copyForCollection(this.getVectorStoreCollection());
        return new MilvusVectorStore(milvusVectorStoreConfig);
    }

    public Tool toFunction(boolean needEnglishName) {
        return toFunction(needEnglishName, RetrievalMode.HYBRID.name());
    }

    public Tool toFunction(boolean needEnglishName, String retrievalMode) {
        return toFunction(needEnglishName, retrievalMode, null);
    }

    /**
     * 构造知识库聊天工具。
     *
     * @param needEnglishName 是否使用英文名称
     * @param retrievalMode 检索模式
     * @param chatTimeContext 聊天时权限上下文
     * @return 聊天工具
     */
    public Tool toFunction(boolean needEnglishName, String retrievalMode, ChatTimeToolAvailabilityContext chatTimeContext) {
        return new DocumentCollectionTool(this, needEnglishName, KnowledgeRetrievalModes.parse(retrievalMode), chatTimeContext);
    }

    public Object getOptionsByKey(String key) {
        Map<String, Object> options = this.getOptions();
        if (KEY_DOC_RECALL_MAX_NUM.equals(key) && (options == null || !options.containsKey(KEY_DOC_RECALL_MAX_NUM))) {
            return 5;
        }
        if (KEY_SIMILARITY_THRESHOLD.equals(key) && (options == null || !options.containsKey(KEY_SIMILARITY_THRESHOLD))) {
            return 0.6f;
        }
        if (KEY_RERANK_ENABLE.equals(key)) {
            if (options == null || !options.containsKey(KEY_RERANK_ENABLE)) {
                return this.getRerankModelId() != null;
            }
            Object value = options.get(key);
            if (value instanceof Boolean) {
                return value;
            }
            if (value instanceof Number) {
                return ((Number) value).intValue() != 0;
            }
            if (value instanceof String) {
                return Boolean.parseBoolean((String) value);
            }
            return false;
        }

        if (options == null) {
            return null;
        }
        if (KEY_SIMILARITY_THRESHOLD.equals(key)) {
            BigDecimal score = (BigDecimal) options.get(key);
            return (float) score.doubleValue();
        }
        return options.get(key);
    }

    public Boolean getApprovalPending() {
        return approvalPending;
    }

    public void setApprovalPending(Boolean approvalPending) {
        this.approvalPending = approvalPending;
    }

    public String getCurrentApprovalActionType() {
        return currentApprovalActionType;
    }

    public void setCurrentApprovalActionType(String currentApprovalActionType) {
        this.currentApprovalActionType = currentApprovalActionType;
    }

    public String getDisplayPublishStatus() {
        return displayPublishStatus;
    }

    public void setDisplayPublishStatus(String displayPublishStatus) {
        this.displayPublishStatus = displayPublishStatus;
    }

    /**
     * 获取创建人展示名称。
     *
     * @return 创建人展示名称
     */
    public String getCreatedByName() {
        return createdByName;
    }

    /**
     * 设置创建人展示名称。
     *
     * @param createdByName 创建人展示名称
     */
    public void setCreatedByName(String createdByName) {
        this.createdByName = createdByName;
    }
}
