package com.emoon.openplatform.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.emoon.openplatform.domain.AiChatModelDo;
import com.emoon.openplatform.domain.AiKnowledgeInfoDo;
import com.emoon.openplatform.domain.bo.QueryVectorBo;
import com.emoon.openplatform.domain.dto.SchemaMessageDTO;
import com.emoon.openplatform.domain.request.SchemaRequest;
import com.emoon.openplatform.mapper.AiChatModelMapper;
import com.emoon.openplatform.mapper.AiKnowledgeInfoMapper;
import com.emoon.openplatform.service.IChatService;
import com.emoon.openplatform.service.IVectorStoreService;
import com.emoon.openplatform.util.RerankUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.common.core.utils.StringUtils;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.List;
import java.util.Map;

/**
 * @author destiny
 * @description:
 * @date 2025/10/13 16:26
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class ChatServiceImpl implements IChatService {

    public static final String RERANK_MODEL_NAME = "bge-reranker-v2-m3";
    private final AiKnowledgeInfoMapper aiKnowledgeInfoMapper;

    private final AiChatModelMapper aiChatModelMapper;

    private final IVectorStoreService vectorStoreService;

    /**
     * 根据请求参数获取与指定知识库相关的模式信息。
     * <p>
     * 该方法首先校验 kid 是否为空，然后根据 kid 查询对应的知识库信息，
     * 接着查询所使用的向量化模型，并基于这些信息从向量存储中检索最相似的内容。
     * 最后使用重排序模型对结果进行重新排序并返回封装后的数据。
     *
     * @param schemaRequest 包含 kid、messages 和 maxDistance 等字段的请求对象，用于查询相关 schema 信息
     * @return 返回包含原始 prompt 及其最近邻向量匹配结果的 SchemaMessageDTO 对象
     * @throws IllegalArgumentException 当 kid 为空或知识库/模型不存在时抛出
     * @throws RuntimeException         当重排序过程中发生 IO 异常时抛出
     */
    @Override
    public SchemaMessageDTO getSchema(SchemaRequest schemaRequest) {
        if (StringUtils.isEmpty(schemaRequest.getKid())) {
            throw new IllegalArgumentException("kid 不能为空");
        }
        // 通过 kid 查询知识库信息
        AiKnowledgeInfoDo aiKnowledgeInfoDo = aiKnowledgeInfoMapper.selectOne(
            new LambdaQueryWrapper<AiKnowledgeInfoDo>().eq(AiKnowledgeInfoDo::getKnowledgeCode, schemaRequest.getKid()));
        if (aiKnowledgeInfoDo == null) {
            throw new IllegalArgumentException("知识库不存在");
        }
        AiChatModelDo aiChatModelDo = aiChatModelMapper.selectOne(
            new LambdaQueryWrapper<AiChatModelDo>().eq(AiChatModelDo::getModelName, aiKnowledgeInfoDo.getEmbeddingModelName()));
        if (aiChatModelDo == null) {
            throw new IllegalArgumentException("向量化模型不存在");
        }

        // 构建向量查询参数对象
        QueryVectorBo queryVectorBo = QueryVectorBo.builder()
            .query(schemaRequest.getMessages())
            .kid(schemaRequest.getKid())
            .apiKey(aiChatModelDo.getApiKey())
            .baseUrl(aiChatModelDo.getApiHost())
            .vectorModelName(aiKnowledgeInfoDo.getVectorModelName())
            .embeddingModelName(aiKnowledgeInfoDo.getEmbeddingModelName())
            .maxResults(aiKnowledgeInfoDo.getRetrieveLimit())
            .build();

        // 执行向量检索操作，获取最相近的结果列表
        List<Map<String, Object>> nearestList = vectorStoreService.getQueryVector(queryVectorBo, schemaRequest.getMaxDistance());

        // 查询重排序模型并对检索结果进行重排序处理（当前固定使用 bge-reranker-v2-m3 模型）
        List<Map<String, Object>> maps = null;
        AiChatModelDo rerankModel = aiChatModelMapper.selectOne(
            new LambdaQueryWrapper<AiChatModelDo>().eq(AiChatModelDo::getModelName, RERANK_MODEL_NAME));
        try {
            maps = RerankUtil.rerankNearestList(schemaRequest.getMessages(), nearestList, rerankModel.getApiHost(), RERANK_MODEL_NAME);
        } catch (IOException e) {
            log.error("重排序失败: {}", e.getMessage(), e);
            throw new RuntimeException(e);
        }

        // 将最终结果构造成 DTO 并返回
        return SchemaMessageDTO.builder()
            .prompt(schemaRequest.getMessages())
            .nearest(maps)
            .build();
    }

}
