package org.dromara.knowledge.service.impl;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.common.core.exception.ServiceException;
import org.dromara.common.core.utils.MapstructUtils;
import org.dromara.common.core.utils.StringUtils;
import org.dromara.common.mybatis.core.page.PageQuery;
import org.dromara.common.mybatis.core.page.TableDataInfo;
import org.dromara.knowledge.domain.KnowledgeVector;
import org.dromara.knowledge.domain.bo.KnowledgeVectorBo;
import org.dromara.knowledge.domain.vo.KnowledgeSearchResultVo;
import org.dromara.knowledge.domain.vo.KnowledgeVectorVo;
import org.dromara.knowledge.mapper.KnowledgeVectorMapper;
import org.dromara.knowledge.service.KnowledgeFragmentService;
import org.dromara.knowledge.service.KnowledgeVectorService;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 知识片段向量Service业务层处理
 *
 * @author pengles
 */
@Slf4j
@RequiredArgsConstructor
@Service
public class KnowledgeVectorServiceImpl extends ServiceImpl<KnowledgeVectorMapper, KnowledgeVector> implements KnowledgeVectorService {

    // 支持的向量模型列表
    private static final List<String> AVAILABLE_MODELS = Arrays.asList(
            "text-embedding-ada-002", "text-embedding-3-small", "text-embedding-3-large",
            "bge-large-zh", "bge-base-zh", "m3e-base", "m3e-large"
    );
    private final KnowledgeVectorMapper baseMapper;
    private final KnowledgeFragmentService knowledgeFragmentService;
    private final org.dromara.knowledge.embedding.EmbeddingService embeddingService;
    private final org.dromara.knowledge.vectorstore.VectorStoreManager vectorStoreManager;

    /**
     * 查询知识片段向量
     */
    @Override
    public KnowledgeVectorVo queryById(Long id) {
        return baseMapper.selectVoById(id);
    }

    /**
     * 查询知识片段向量列表
     */
    @Override
    public TableDataInfo<KnowledgeVectorVo> queryPageList(KnowledgeVectorBo bo, PageQuery pageQuery) {
        LambdaQueryWrapper<KnowledgeVector> lqw = buildQueryWrapper(bo);
        Page<KnowledgeVectorVo> result = baseMapper.selectVoPage(pageQuery.build(), lqw);
        return TableDataInfo.build(result);
    }

    /**
     * 查询知识片段向量列表
     */
    @Override
    public List<KnowledgeVectorVo> queryList(KnowledgeVectorBo bo) {
        LambdaQueryWrapper<KnowledgeVector> lqw = buildQueryWrapper(bo);
        return baseMapper.selectVoList(lqw);
    }

    private LambdaQueryWrapper<KnowledgeVector> buildQueryWrapper(KnowledgeVectorBo bo) {
        Map<String, Object> params = bo.getParams();
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(StringUtils.isNotBlank(bo.getKnowledgeId()), KnowledgeVector::getKnowledgeId, bo.getKnowledgeId());
        lqw.eq(StringUtils.isNotBlank(bo.getDocId()), KnowledgeVector::getDocId, bo.getDocId());
        lqw.eq(StringUtils.isNotBlank(bo.getFragmentId()), KnowledgeVector::getFragmentId, bo.getFragmentId());
        lqw.eq(StringUtils.isNotBlank(bo.getVectorId()), KnowledgeVector::getVectorId, bo.getVectorId());
        lqw.eq(bo.getDimension() != null, KnowledgeVector::getDimension, bo.getDimension());
        lqw.like(StringUtils.isNotBlank(bo.getModelName()), KnowledgeVector::getModelName, bo.getModelName());
        lqw.eq(bo.getStatus() != null, KnowledgeVector::getStatus, bo.getStatus());
        return lqw;
    }

    /**
     * 新增知识片段向量
     */
    @Override
    public Boolean insertByBo(KnowledgeVectorBo bo) {
        KnowledgeVector add = MapstructUtils.convert(bo, KnowledgeVector.class);
        validEntityBeforeSave(add);
        boolean flag = baseMapper.insert(add) > 0;
        if (flag) {
            bo.setId(add.getId());
        }
        return flag;
    }

    /**
     * 修改知识片段向量
     */
    @Override
    public Boolean updateByBo(KnowledgeVectorBo bo) {
        KnowledgeVector update = MapstructUtils.convert(bo, KnowledgeVector.class);
        validEntityBeforeSave(update);
        return baseMapper.updateById(update) > 0;
    }

    /**
     * 保存前的数据校验
     */
    private void validEntityBeforeSave(KnowledgeVector entity) {
        //TODO 做一些数据校验,如唯一约束
    }

    /**
     * 批量删除知识片段向量
     */
    @Override
    public Boolean deleteWithValidByIds(Collection<Long> ids, Boolean isValid) {
        if (isValid) {
            //TODO 做一些业务上的校验,判断是否需要校验
        }
        return baseMapper.deleteBatchIds(ids) > 0;
    }

    /**
     * 根据向量ID查询知识片段向量
     */
    @Override
    public KnowledgeVectorVo queryByVectorId(String vectorId) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getVectorId, vectorId);
        return baseMapper.selectVoOne(lqw);
    }

    /**
     * 根据知识库ID查询知识片段向量列表
     */
    @Override
    public List<KnowledgeVectorVo> queryByKnowledgeId(String knowledgeId) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getKnowledgeId, knowledgeId);
        return baseMapper.selectVoList(lqw);
    }

    /**
     * 根据文档ID查询知识片段向量列表
     */
    @Override
    public List<KnowledgeVectorVo> queryByDocId(String docId) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getDocId, docId);
        return baseMapper.selectVoList(lqw);
    }

    /**
     * 根据知识片段ID查询知识片段向量列表
     */
    @Override
    public List<KnowledgeVectorVo> queryByFragmentId(String fragmentId) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getFragmentId, fragmentId);
        return baseMapper.selectVoList(lqw);
    }

    /**
     * 根据模型名称查询知识片段向量列表
     */
    @Override
    public List<KnowledgeVectorVo> queryByModelName(String modelName) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getModelName, modelName);
        return baseMapper.selectVoList(lqw);
    }

    /**
     * 根据状态查询知识片段向量列表
     */
    @Override
    public List<KnowledgeVectorVo> queryByStatus(Integer status) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getStatus, status);
        return baseMapper.selectVoList(lqw);
    }

    /**
     * 批量更新知识片段向量状态
     */
    @Override
    public Boolean updateStatusByIds(Collection<Long> ids, Integer status) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.in(KnowledgeVector::getId, ids);
        KnowledgeVector update = new KnowledgeVector();
        update.setStatus(status);
        return baseMapper.update(update, lqw) > 0;
    }

    /**
     * 根据知识库ID和状态查询知识片段向量列表
     */
    @Override
    public List<KnowledgeVectorVo> queryByKnowledgeIdAndStatus(String knowledgeId, Integer status) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getKnowledgeId, knowledgeId)
                .eq(KnowledgeVector::getStatus, status);
        return baseMapper.selectVoList(lqw);
    }

    @Override
    public String generateVector(String knowledgeId, String docId, String fragmentId, String content, String modelName) {
        try {
            // 验证模型是否可用
            if (!isModelAvailable(modelName)) {
                throw new ServiceException("不支持的向量模型: " + modelName);
            }

            double[] vectorData = embeddingService.embed(content, modelName);
            String vectorId = IdUtil.fastSimpleUUID();
            KnowledgeVector vector = new KnowledgeVector();
            vector.setKnowledgeId(knowledgeId);
            vector.setDocId(docId);
            vector.setFragmentId(fragmentId);
            vector.setVectorId(vectorId);
            vector.setVectorData(Arrays.toString(vectorData));
            vector.setDimension(vectorData.length);
            vector.setModelName(modelName);
            vector.setStatus(20);
            LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
            lqw.eq(KnowledgeVector::getFragmentId, fragmentId).eq(KnowledgeVector::getModelName, modelName);
            KnowledgeVector exist = baseMapper.selectOne(lqw);
            if (exist == null) {
                baseMapper.insert(vector);
            } else {
                exist.setVectorData(vector.getVectorData());
                exist.setDimension(vector.getDimension());
                exist.setStatus(20);
                baseMapper.updateById(exist);
                vectorId = exist.getVectorId();
            }
            vectorStoreManager.current().upsert(knowledgeId, docId, fragmentId, vectorId, vectorData, content, modelName);
            log.info("为片段 {} 生成向量成功，向量ID: {}", fragmentId, vectorId);
            return vectorId;
        } catch (Exception e) {
            log.error("生成向量失败: {}", e.getMessage(), e);
            throw new ServiceException("生成向量失败: " + e.getMessage());
        }
    }

    @Override
    public List<String> batchGenerateVectors(String knowledgeId, String docId,
                                             Map<String, String> fragmentContents, String modelName) {
        List<String> successIds = new ArrayList<>();

        for (Map.Entry<String, String> entry : fragmentContents.entrySet()) {
            try {
                String vectorId = generateVector(knowledgeId, docId, entry.getKey(), entry.getValue(), modelName);
                successIds.add(entry.getKey());
            } catch (Exception e) {
                log.error("为片段 {} 生成向量失败: {}", entry.getKey(), e.getMessage());
            }
        }

        log.info("批量生成向量完成，成功: {}/{}", successIds.size(), fragmentContents.size());
        return successIds;
    }

    @Override
    public List<KnowledgeSearchResultVo> vectorSimilaritySearch(String knowledgeId, String queryText,
                                                                String modelName, Integer topK, Double threshold) {
        try {
            // 将查询文本转换为向量
            double[] queryVector = embeddingService.embed(queryText, modelName);

            // 获取知识库中的所有向量
            return vectorStoreManager.current().search(knowledgeId, queryVector, modelName, topK, threshold);

            // 计算相似度并排序


        } catch (Exception e) {
            log.error("向量搜索失败: {}", e.getMessage(), e);
            throw new ServiceException("向量搜索失败: " + e.getMessage());
        }
    }

    @Override
    public double[] textToVector(String text, String modelName) {
        return embeddingService.embed(text, modelName);
    }

    @Override
    public double calculateSimilarity(double[] vector1, double[] vector2) {
        if (vector1.length != vector2.length) {
            throw new ServiceException("向量维度不匹配");
        }

        // 计算余弦相似度
        double dotProduct = 0;
        double norm1 = 0;
        double norm2 = 0;

        for (int i = 0; i < vector1.length; i++) {
            dotProduct += vector1[i] * vector2[i];
            norm1 += vector1[i] * vector1[i];
            norm2 += vector2[i] * vector2[i];
        }

        if (norm1 == 0 || norm2 == 0) {
            return 0;
        }

        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }

    @Override
    public int deleteVectorsByKnowledgeId(String knowledgeId) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getKnowledgeId, knowledgeId);
        return baseMapper.delete(lqw);
    }

    @Override
    public int deleteVectorsByDocId(String docId) {
        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getDocId, docId);
        return baseMapper.delete(lqw);
    }

    @Override
    public Map<String, Object> getVectorStatistics(String knowledgeId) {
        Map<String, Object> stats = new HashMap<>();

        LambdaQueryWrapper<KnowledgeVector> lqw = Wrappers.lambdaQuery();
        lqw.eq(KnowledgeVector::getKnowledgeId, knowledgeId);

        List<KnowledgeVector> vectors = baseMapper.selectList(lqw);

        stats.put("totalVectors", vectors.size());
        stats.put("completedVectors", vectors.stream().filter(v -> v.getStatus() == 20).count());
        stats.put("processingVectors", vectors.stream().filter(v -> v.getStatus() == 15).count());
        stats.put("failedVectors", vectors.stream().filter(v -> v.getStatus() == 30).count());

        // 按模型统计
        Map<String, Long> modelStats = vectors.stream()
                .collect(Collectors.groupingBy(KnowledgeVector::getModelName, Collectors.counting()));
        stats.put("modelStatistics", modelStats);

        return stats;
    }

    @Override
    public boolean isModelAvailable(String modelName) {
        return AVAILABLE_MODELS.contains(modelName);
    }

    @Override
    public List<String> getAvailableModels() {
        return new ArrayList<>(AVAILABLE_MODELS);
    }

    /**
     * 根据模型名称获取向量维度
     */
    private int getDimensionByModel(String modelName) {
        switch (modelName) {
            case "text-embedding-ada-002":
                return 1536;
            case "text-embedding-3-small":
                return 1536;
            case "text-embedding-3-large":
                return 3072;
            case "bge-large-zh":
                return 1024;
            case "bge-base-zh":
                return 768;
            case "m3e-base":
                return 768;
            case "m3e-large":
                return 1024;
            default:
                return 768; // 默认维度
        }
    }

    /**
     * 解析向量数据字符串
     */
    private double[] parseVectorData(String vectorDataStr) {
        try {
            if (StrUtil.isBlank(vectorDataStr)) {
                return null;
            }
            String cleanStr = vectorDataStr.replace("[", "").replace("]", "");
            String[] parts = cleanStr.split(",");
            double[] vector = new double[parts.length];
            for (int i = 0; i < parts.length; i++) {
                vector[i] = Double.parseDouble(parts[i].trim());
            }
            return vector;
        } catch (Exception e) {
            log.error("解析向量数据失败: {}", e.getMessage());
            return null;
        }
    }

}
